"""全市场数据同步（选股专用，未复权；与回测 candles 表隔离）。 设计：trade_cal 取近 N 个交易日 -> 逐日 pro.daily(trade_date=...) / pro.daily_basic(trade_date=...) 一次返回全市场当日数据 -> 按 trade_date 删旧插新批量入库（幂等）。 同步为进程内后台任务（MVP 不引入任务队列），前端轮询 /api/screener/sync/status。 daily 与 daily_basic 分步独立落库：daily_basic 积分不足时日线仍可用，错误写入状态不中断任务。 """ from __future__ import annotations import asyncio import time from datetime import datetime, timedelta from sqlalchemy import delete, func, insert, select from sqlalchemy.ext.asyncio import AsyncSession from ..config import settings from ..models import DailySnapshot, MarketDaily, StockBasic, TradeCalendar from .llm import ScreenerError # 进程内单例任务状态（uvicorn --reload 单进程场景够用） _sync_state: dict = { "running": False, "step": None, "total_days": 0, "done_days": 0, "error": None, "started_at": None, "finished_at": None, } _sync_task: asyncio.Task | None = None _sync_lock = asyncio.Lock() _BATCH = 5000 # executemany 分批行数 # Tushare 积分/权限不足的特征文案（daily_basic 常见门槛） _PERM_MARKS = ("抱歉，您没有访问该项目权限", "积分", "权限") # 频率超限特征（等待 62s 重试一次） _RATE_MARKS = ("频率超限", "每分钟") 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 _get_pro(): """token 检查 + 返回 pro api 客户端（同步对象，调用需 to_thread 包裹）。""" if not settings.tushare_token: raise ScreenerError("未配置 TUSHARE_TOKEN，无法同步全市场数据（backend/.env）") import tushare as ts ts.set_token(settings.tushare_token) return ts.pro_api() def _parse_d(s: str) -> datetime: return datetime.strptime(str(s), "%Y%m%d") def _fetch_calendar_sync(pro) -> list[str]: """拉取宽范围交易日历（近 18 个月 + 未来 3 个月），返回 YYYYMMDD 列表。""" time.sleep(settings.screener_sync_interval) end = (datetime.now() + timedelta(days=90)).strftime("%Y%m%d") start = (datetime.now() - timedelta(days=550)).strftime("%Y%m%d") cal = _call_retry(pro.trade_cal, exchange="SSE", start_date=start, end_date=end, is_open="1") return sorted(cal["cal_date"].tolist()) async def _recent_trade_dates(session: AsyncSession, pro, days: int) -> list[str]: """近 N 个交易日（YYYYMMDD，倒序）。日历本地缓存，仅在覆盖不到当天时刷新一次。 trade_cal 低积分版限频 1 次/小时：刷新被限频时沿用缓存（日历略旧无害—— daily 对未生成日期返回空，同步会自然跳过）。 """ cached = (await session.execute(select(TradeCalendar.trade_date).order_by(TradeCalendar.trade_date.desc()))).scalars().all() today = datetime.now().strftime("%Y%m%d") have_today = bool(cached) and cached[0] >= today if not have_today: try: dates = await asyncio.to_thread(_fetch_calendar_sync, pro) await session.execute(delete(TradeCalendar)) await session.execute(insert(TradeCalendar), [{"trade_date": d} for d in dates]) await session.commit() cached = dates[::-1] except Exception as e: # noqa: BLE001 —— 限频且无缓存时才致命 if not cached: raise ScreenerError(f"获取交易日历失败（且本地无缓存）: {str(e)[:150]}") from e _sync_state["step"] = "交易日历刷新受限，沿用本地缓存" recent = [d for d in cached if d <= today][:days] if not recent: raise ScreenerError("交易日历为空") return recent def _fetch_daily(pro, d: str) -> list[dict]: """拉取某交易日全市场日线（未复权）。当日数据未生成（盘前/盘中）返回空。""" time.sleep(settings.screener_sync_interval) df = _call_retry(pro.daily, trade_date=d) if df is None or df.empty: return [] rows = [] for _, r in df.iterrows(): rows.append({ "trade_date": _parse_d(d), "ts_code": r["ts_code"], "open": float(r["open"]), "high": float(r["high"]), "low": float(r["low"]), "close": float(r["close"]), "pre_close": float(r["pre_close"]), "change": None if r.get("change") != r.get("change") else float(r["change"]), "pct_chg": None if r.get("pct_chg") != r.get("pct_chg") else float(r["pct_chg"]), "vol": float(r["vol"]), # 手 "amount": float(r["amount"]), # 千元 }) return rows def _fetch_basic(pro, d: str) -> list[dict]: """拉取某交易日每日指标快照（daily_basic，低积分版限频 1 次/分钟）。 失败（积分不足等）时记录错误返回空，不拖垮日线同步。 """ time.sleep(settings.screener_sync_interval) try: df = _call_retry(pro.daily_basic, trade_date=d) except Exception as e: # noqa: BLE001 msg = str(e) if any(m in msg for m in _PERM_MARKS): _sync_state["error"] = ( f"Tushare 无法获取每日指标(daily_basic)：{msg[:150]}。" "市值/市盈率等条件不可用；纯指标选股不受影响。" ) return [] raise if df is None or df.empty: return [] rows = [] for _, r in df.iterrows(): def _f(key: str) -> float | None: v = r.get(key) return None if v is None or v != v else float(v) rows.append({ "trade_date": _parse_d(d), "ts_code": r["ts_code"], "close": _f("close"), "turnover_rate": _f("turnover_rate"), "turnover_rate_f": _f("turnover_rate_f"), "volume_ratio": _f("volume_ratio"), "pe": _f("pe"), "pe_ttm": _f("pe_ttm"), "pb": _f("pb"), "total_mv": _f("total_mv"), "circ_mv": _f("circ_mv"), # 万元 }) return rows def _sync_stock_list_sync(pro) -> list[dict]: """拉取在市股票列表。""" time.sleep(settings.screener_sync_interval) df = _call_retry(pro.stock_basic, exchange="", list_status="L", fields="ts_code,symbol,name,area,industry,market,exchange,list_status,list_date,delist_date") rows = [] for _, r in df.iterrows(): rows.append({ "ts_code": r["ts_code"], "symbol": r["symbol"], "name": r["name"], "area": r.get("area") or None, "industry": r.get("industry") or None, "market": r.get("market") or None, "exchange": r["exchange"] or "", "list_status": r["list_status"], "list_date": r.get("list_date") or "", "delist_date": r.get("delist_date") or None, }) return rows def _norm_date(v) -> str: """把 DB 读出的 trade_date（可能是 datetime 或 str）归一为 YYYYMMDD。""" if hasattr(v, "strftime"): return v.strftime("%Y%m%d") return str(v)[:10].replace("-", "") async def _existing_dates(session: AsyncSession, model) -> set[str]: """某表已落库的交易日集合（YYYYMMDD 字符串，便于比对）。""" res = await session.execute(select(func.distinct(model.trade_date))) return {_norm_date(r[0]) for r in res} async def _replace_day(session: AsyncSession, model, rows: list[dict], d_str: str) -> None: """按交易日删旧插新（幂等），executemany 分批。""" d = _parse_d(d_str) await session.execute(delete(model).where(model.trade_date == d)) for i in range(0, len(rows), _BATCH): await session.execute(insert(model), rows[i : i + _BATCH]) await session.commit() async def _run_sync(days: int, force: bool) -> None: """后台任务主体：stock_basic -> 逐日日线 -> 最新交易日快照。异常写状态。 daily_basic 只拉最新交易日（快照条件仅作用于最新截面，且低积分 token 限频 1 次/分钟）。 """ from ..db import async_session # 延迟导入避免循环 try: pro = await asyncio.to_thread(_get_pro) # 1) 股票列表（已有数据则跳过——stock_basic 低积分版限频 1 次/小时） async with async_session() as session: stocks_now = int(await session.scalar(select(func.count()).select_from(StockBasic)) or 0) if stocks_now == 0 or force: _sync_state["step"] = "正在同步股票列表" try: rows = await asyncio.to_thread(_sync_stock_list_sync, pro) async with async_session() as session: await session.execute(delete(StockBasic)) for i in range(0, len(rows), _BATCH): await session.execute(insert(StockBasic), rows[i : i + _BATCH]) await session.commit() except Exception as e: # noqa: BLE001 —— 受限时沿用现有列表继续 if stocks_now > 0: _sync_state["step"] = f"股票列表同步受限（沿用现有 {stocks_now} 只）" else: raise # 2) 逐交易日日线（增量；当日未生成则跳过） async with async_session() as session: dates = await _recent_trade_dates(session, pro, days) have_daily = set() if force else await _existing_dates(session, MarketDaily) todo = [d for d in dates if d not in have_daily] _sync_state["total_days"] = len(todo) _sync_state["done_days"] = 0 for d in todo: _sync_state["step"] = f"正在同步 {d} 日线（{_sync_state['done_days'] + 1}/{len(todo)}）" daily_rows = await asyncio.to_thread(_fetch_daily, pro, d) if daily_rows: # 盘前/盘中等未生成数据的日期直接跳过 async with async_session() as session: await _replace_day(session, MarketDaily, daily_rows, d) _sync_state["done_days"] += 1 # 3) 最新「有数据」交易日的快照（daily_basic，仅 1 次调用） # 用 market_daily 实际最大交易日（今天的数据收盘后才生成，日历最新日会拉到空） async with async_session() as session: latest_dt = await session.scalar(select(func.max(MarketDaily.trade_date))) latest = latest_dt.strftime("%Y%m%d") if latest_dt else None if latest: async with async_session() as session: have_snap = force or latest not in await _existing_dates(session, DailySnapshot) if have_snap: _sync_state["step"] = f"正在同步 {latest} 每日指标" basic_rows = await asyncio.to_thread(_fetch_basic, pro, latest) if basic_rows: async with async_session() as session: await _replace_day(session, DailySnapshot, basic_rows, latest) _sync_state["step"] = "同步完成" except Exception as e: # noqa: BLE001 _sync_state["error"] = f"同步失败：{str(e)[:300]}" _sync_state["step"] = "同步失败" finally: _sync_state["running"] = False _sync_state["finished_at"] = datetime.now() async def start_sync(session: AsyncSession, days: int, force: bool) -> dict: """幂等启动后台同步任务；已在跑则直接返回当前状态。""" global _sync_task async with _sync_lock: if _sync_state["running"] and _sync_task and not _sync_task.done(): return dict(_sync_state) _sync_state.update({ "running": True, "step": "准备同步", "total_days": days, "done_days": 0, "error": None, "started_at": datetime.now(), "finished_at": None, }) _sync_task = asyncio.create_task(_run_sync(days, force)) return dict(_sync_state) async def get_sync_status(session: AsyncSession) -> dict: """合并任务状态 + DB 实况（最新交易日/行数/ready 标志），与 ScreenerSyncStatus DTO 对齐。""" stocks = int(await session.scalar(select(func.count()).select_from(StockBasic)) or 0) daily_rows = int(await session.scalar(select(func.count()).select_from(MarketDaily)) or 0) snap_rows = int(await session.scalar(select(func.count()).select_from(DailySnapshot)) or 0) last_daily = await session.scalar(select(func.max(MarketDaily.trade_date))) n_dates = int(await session.scalar(select(func.count(func.distinct(MarketDaily.trade_date)))) or 0) status = dict(_sync_state) status.update({ "stats": {"stocks": stocks, "daily_rows": daily_rows, "snapshot_rows": snap_rows, "dates": n_dates}, "last_trade_date": last_daily, "last_synced_at": _sync_state.get("finished_at") or _sync_state.get("started_at"), "ready": daily_rows > 0, }) return status