88 lines
2.9 KiB
Python
88 lines
2.9 KiB
Python
"""懒加载管道共用件：tushare pro 单例、限频重试、取值清洗、stock_sync_state 读写。
|
||
|
||
finance / dividend 等按需单查管道共用；语义与 company.py 内的私有版本一致
|
||
（那处历史代码未迁移，新管道一律从这里取）。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import time
|
||
from datetime import datetime, timedelta, timezone
|
||
|
||
from sqlalchemy import select
|
||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from ..config import settings
|
||
from ..models import StockSyncState
|
||
|
||
# 频率超限特征（等待 62s 重试一次；与 company.py / data.etf_sync._call_retry 同款语义）
|
||
_RATE_MARKS = ("频率超限", "每分钟")
|
||
|
||
_pro = None # 惰性单例（get_pro 每次都 ts.set_token 写文件，没必要重复）
|
||
|
||
|
||
def get_pro_lazy():
|
||
if not settings.tushare_token:
|
||
raise RuntimeError("未配置 TUSHARE_TOKEN，无法拉取 tushare 数据（backend/.env）")
|
||
global _pro
|
||
if _pro is None:
|
||
from .tushare_provider import get_pro
|
||
|
||
_pro = get_pro()
|
||
return _pro
|
||
|
||
|
||
def call_retry(fn, *args, **kwargs):
|
||
"""同步调用 tushare 接口；「每分钟」级频率超限等 62s 重试一次。"""
|
||
try:
|
||
return fn(*args, **kwargs)
|
||
except Exception as e: # noqa: BLE001
|
||
msg = str(e)
|
||
if any(m in msg for m in _RATE_MARKS) and "小时" not in msg:
|
||
time.sleep(62)
|
||
return fn(*args, **kwargs)
|
||
raise
|
||
|
||
|
||
def utcnow() -> datetime:
|
||
return datetime.now(timezone.utc).replace(tzinfo=None)
|
||
|
||
|
||
def fresh(updated_at: datetime | None, days: int) -> bool:
|
||
return updated_at is not None and updated_at >= utcnow() - timedelta(days=days)
|
||
|
||
|
||
def s_clean(v) -> str | None:
|
||
"""pandas NaN / 空串 / None -> None，其余 strip。"""
|
||
if v is None or (isinstance(v, float) and v != v):
|
||
return None
|
||
s = str(v).strip()
|
||
return s or None
|
||
|
||
|
||
def f_clean(v) -> float | None:
|
||
try:
|
||
f = float(v)
|
||
except (TypeError, ValueError):
|
||
return None
|
||
return None if f != f else f # NaN -> None
|
||
|
||
|
||
async def read_sync_state(session: AsyncSession, ts_code: str, kind: str) -> StockSyncState | None:
|
||
return (await session.execute(
|
||
select(StockSyncState).where(
|
||
StockSyncState.ts_code == ts_code, StockSyncState.kind == kind
|
||
))).scalar_one_or_none()
|
||
|
||
|
||
async def upsert_sync_state(session: AsyncSession, ts_code: str, kind: str, *, has_data: bool) -> None:
|
||
"""写入同步状态（不 commit，由调用方统一提交）。"""
|
||
stmt = pg_insert(StockSyncState).values(
|
||
ts_code=ts_code, kind=kind, last_synced_at=utcnow(), has_data=has_data,
|
||
)
|
||
stmt = stmt.on_conflict_do_update(
|
||
index_elements=["ts_code", "kind"],
|
||
set_={"last_synced_at": stmt.excluded.last_synced_at, "has_data": stmt.excluded.has_data},
|
||
)
|
||
await session.execute(stmt)
|