"""同步任务：5 分钟 K 线（mairui）。 设计要点： 1) 启动时**一次 SQL** 拿全表 {stock_code: max(bar_time)} 快照 2) 内存里给每只股票算 (start, end, windows_needed) - DB 没数据 → start = mairui 历史深度起点 (2023-06-14) - DB 有数据 → start = max(2023-06-14, max(bar_time) - 2 天) - 兜底：start >= end → 跳过这只 3) 输出"计划"（全量/增量分类、预估请求数、预估耗时）再开始执行 4) 5 worker × 1 RPS/worker = 5 RPS 总（守住 mairui 上限） 5) 窗口大小 1 年（mairui 一次最多 ~11600 条） akshare 时代写法：start 写死 2020-01-01 + 4 天窗口，137 小时跑完全市场。 本版：增量时通常只补最近 1-2 天，~30 分钟。 """ from __future__ import annotations import os import time from concurrent.futures import ThreadPoolExecutor, as_completed from datetime import datetime, timedelta, timezone from typing import Any, Optional from app.core.db import ops as db_ops from app.core.datasource.base import registry as ds_registry from app.core.datasource.registry import is_source_ready from app.core.datasource.utils import is_a_share_code, to_code6, to_hermes from app.core.sync.base import SyncTask from app.core.sync.registry import mark_sync_blocked from app.core.utils.logging import get_logger logger = get_logger("sync.kline_5min") # mairui 5min K 实测历史深度起点（实测 2023-06-14 才有数据，更早就 0） MAIRUI_5MIN_FLOOR = datetime(2023, 6, 14, tzinfo=timezone.utc) # 单次请求窗口大小：1 年（mairui 单次最多 ~11640 条 5min） WINDOW_DAYS = 365 # 增量时往前多取的天数（防交易日历边界漏当天） INCREMENT_OVERLAP_DAYS = 2 class SyncKline5Min(SyncTask): dataset_id = "kline_5min" def _plan( self, stock_codes: list[str], end: datetime, snapshots: dict[str, Optional[datetime]], ) -> list[tuple[str, datetime, datetime]]: """为每只股票算 (start, end) — 已排除 start >= end 的"无需同步"。 Returns: [(code6, start, end), ...] """ plans = [] for code6 in stock_codes: latest = snapshots.get(code6) if latest is None: # DB 里没数据 → 全量从 mairui 历史深度起点 start = MAIRUI_5MIN_FLOOR else: # DB 有数据 → 增量：从 max(bar_time) - overlap 到 end start = max(MAIRUI_5MIN_FLOOR, latest - timedelta(days=INCREMENT_OVERLAP_DAYS)) if start < end: plans.append((code6, start, end)) return plans def _run( self, *, trigger_source: str = "manual", codes: list[str] | None = None, max_workers: int = 5, **kwargs, ) -> dict[str, Any]: primary = ds_registry.get("datasource_mairui") if primary is None: return {"status": "error", "message": "5min 数据源 mairui 未注册"} ok, reason = is_source_ready(primary.key) if not ok: mark_sync_blocked(self.dataset_id, message=f"{primary.key} 未就绪: {reason}") return {"status": "blocked", "message": f"{primary.key} 未就绪: {reason}"} if codes: stock_codes = [to_code6(c) for c in codes] else: stock_codes = [c for c in db_ops.iter_stock_codes(active_only=True) if is_a_share_code(c)] limit_raw = os.environ.get("MARKET_DATA_STOCK_LIMIT", "").strip() if limit_raw.isdigit() and int(limit_raw) > 0: stock_codes = stock_codes[: int(limit_raw)] if not stock_codes: return {"status": "error", "message": "无股票代码"} # ── 1) 一次 SQL 拿全表快照 ── logger.info(f"[5min] 读取 DB 快照 ({len(stock_codes)} 只待查)…") snapshots = db_ops.get_kline_5min_snapshots() logger.info(f"[5min] DB 里已有 {len(snapshots)} 只的 kline_5min 数据") # ── 2) 内存算计划 ── # PG TIMESTAMPTZ 是 timezone-aware，所以 end 也必须 tz-aware 否则比较会炸 end = datetime.now(timezone.utc) plans = self._plan(stock_codes, end, snapshots) if not plans: msg = f"5min K线 — 全部 {len(stock_codes)} 只都已最新（无需同步）" logger.info(msg) return {"status": "ok", "message": msg, "ok": 0, "fail": 0, "rows": 0, "elapsed_sec": 0.0} # 分类统计 full_sync = [p for p in plans if snapshots.get(p[0]) is None] incr_sync = [p for p in plans if snapshots.get(p[0]) is not None] total_windows = sum( max(1, (p[2] - p[1]).days // WINDOW_DAYS + 1) for p in plans ) # mairui 5 RPS 上限（5 worker × 1 RPS） est_sec = total_windows / 5.0 logger.info( f"[5min] 计划: 总 {len(plans)} 只 (全量 {len(full_sync)} / 增量 {len(incr_sync)})，" f"总请求 {total_windows}，按 5 RPS 预估 {est_sec:.0f}s ({est_sec/60:.1f}min)" ) self._progress( message=f"计划: {len(plans)} 只 (全量 {len(full_sync)}/增量 {len(incr_sync)}) " f"约 {total_windows} 个请求, 预估 {est_sec/60:.1f}min", current=0, total=len(plans), current_step="planning", ) # ── 3) 执行 ── t0 = time.time() ok_cnt = fail_cnt = 0 rows_total = 0 with ThreadPoolExecutor(max_workers=max_workers) as pool: futures = {pool.submit(self._sync_one, primary, p[0], p[1], p[2]): p[0] for p in plans} for i, future in enumerate(as_completed(futures), 1): c6 = futures[future] try: res = future.result() if res["status"] == "ok": ok_cnt += 1 rows_total += res["rows"] else: fail_cnt += 1 if res.get("error"): logger.warning(f"[5min {c6}] {res['error']}") except Exception as e: fail_cnt += 1 logger.warning(f"[5min {c6}] {e}") if i % 20 == 0 or i == len(plans): self._progress( message=f"5min 进度 {i}/{len(plans)} OK:{ok_cnt} FAIL:{fail_cnt}", current=i, total=len(plans), current_step=c6, ) elapsed = round(time.time() - t0, 1) msg = ( f"5min K线 {ok_cnt}成 {fail_cnt}败 共{rows_total}行, " f"全量{len(full_sync)}/增量{len(incr_sync)}, {elapsed}s" ) return { "status": "ok" if fail_cnt == 0 else "warning", "message": msg, "ok": ok_cnt, "fail": fail_cnt, "rows": rows_total, "full_sync": len(full_sync), "incr_sync": len(incr_sync), "total_requests": total_windows, "elapsed_sec": elapsed, } def _sync_one(self, primary, code6: str, start: datetime, end: datetime) -> dict: all_rows: list[dict] = [] try: cur = start while cur < end: w_end = min(cur + timedelta(days=WINDOW_DAYS), end) df = primary.fetch_kline_5min( code6, cur.strftime("%Y-%m-%d"), w_end.strftime("%Y-%m-%d"), ) if df is not None and not df.empty: hermes = to_hermes(code6) for _, r in df.iterrows(): bar_time = r["bar_time"] all_rows.append({ "stock_code": hermes, "bar_time": bar_time.strftime("%Y-%m-%d %H:%M:%S") if hasattr(bar_time, "strftime") else str(bar_time), "open": float(r.get("open") or 0), "high": float(r.get("high") or 0), "low": float(r.get("low") or 0), "close": float(r.get("close") or 0), "volume": float(r.get("volume") or 0), "amount": float(r.get("amount") or 0), "turnover_rate": float(r.get("turnover_rate") or 0), }) cur = w_end + timedelta(days=1) except Exception as e: return {"status": "fail", "error": str(e)} if not all_rows: return {"status": "fail", "error": "no data"} try: CHUNK = 5000 for i in range(0, len(all_rows), CHUNK): db_ops.upsert_kline_5min(all_rows[i: i + CHUNK]) except Exception as e: return {"status": "fail", "error": f"db write: {e}"} return {"status": "ok", "rows": len(all_rows)}