market_sync/app/tasks/task_tick_trade.py

"""同步任务：当天逐笔交易（mairui 源）。

数据源：mairui `hsrl/zbjy/{code6}/{licence}`
口径：单笔成交 — 时间 / 价 / 量 / 买卖方向
更新：每日 21:00（mairui 当天数据此时才发布；21:00 之前 API 返空或昨日数据）

设计要点：
  1) 起步门控：21:00 之前（未传 force=True）→ 早返 warning
     避免 0 行空跑（5000 只股票 × 1s 限速 ≈ 17 分钟浪费）
  2) 5 worker × 1 RPS/worker = 5 RPS（守住 mairui 上限）
  3) 不做"今日已同步"dedup —— mairui "当天"实际是"最近已发布交易日"，
     trade_date 可能比 wall-clock 早一天；重跑靠 PK + ON CONFLICT 幂等
  4) 单只股票日均 5w-10w tick，CHUNK=2000 upsert 防 driver 撑爆
  5) 数据是"当天 only"，无 start/end 窗口；不存增量概念
"""
from __future__ import annotations

import os
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from typing import Any

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.tick_trade")

# mairui 文档：「更新：每日 21:00」—— 早于此时 mairui 仍返昨日/空数据
MAIRUI_TICK_TRADE_PUBLISH_HOUR = 21
MAIRUI_TICK_TRADE_PUBLISH_MIN = 0
# 21:00 之后再跑（留 5 分钟缓冲，等 mairui 完整入库）
TICK_TRADE_RUN_GATE = (MAIRUI_TICK_TRADE_PUBLISH_HOUR, MAIRUI_TICK_TRADE_PUBLISH_MIN + 5)


def _passes_publish_gate(now: Optional[datetime] = None, *, force: bool = False) -> tuple[bool, str]:
    """判断当前时间是否已过 mairui 21:00 数据发布门控。

    force=True 跳过门控（手动 / 测试用）。
    """
    if force:
        return True, "force=True 跳过门控"
    n = now or datetime.now()
    gate = n.replace(hour=TICK_TRADE_RUN_GATE[0], minute=TICK_TRADE_RUN_GATE[1], second=0, microsecond=0)
    if n < gate:
        return False, f"mairui 逐笔数据 {MAIRUI_TICK_TRADE_PUBLISH_HOUR:02d}:{MAIRUI_TICK_TRADE_PUBLISH_MIN:02d} 发布；当前 {n.strftime('%H:%M')} 早于门控，skip"
    return True, "ok"


class SyncTickTrade(SyncTask):
    dataset_id = "tick_trade"

    def _run(
        self,
        *,
        trigger_source: str = "manual",
        codes: list[str] | None = None,
        max_workers: int = 5,
        force: bool = False,
        **kwargs,
    ) -> dict[str, Any]:
        mr = ds_registry.get("datasource_mairui")
        if mr is None:
            return {"status": "error", "message": "mairui 数据源未注册"}
        ok, reason = is_source_ready(mr.key)
        if not ok:
            mark_sync_blocked(self.dataset_id, message=f"mairui 未就绪: {reason}")
            return {"status": "blocked", "message": f"mairui 未就绪: {reason}"}

        # ── 21:00 门控（force=True 跳过）──
        gate_ok, gate_msg = _passes_publish_gate(force=force)
        if not gate_ok:
            logger.info(f"[tick_trade] 门控未过: {gate_msg}")
            return {"status": "warning", "message": gate_msg, "ok": 0, "fail": 0, "rows": 0, "elapsed_sec": 0.0}

        # ── 取股票列表 ──
        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": "无股票代码"}

        # 注：不做"今日已同步"dedup。原因：
        #   - mairui 的"当天"实际是"最近一个已发布的交易日"（21:00 之前是昨日），
        #     用 today_str 做去重会失效（即使刚跑完，trade_date 还是昨日）。
        #   - PK 含 trade_date，重跑是 ON CONFLICT DO UPDATE，安全且幂等。
        #   - 17 分钟全市场开销可接受；非交易日用 `MARKET_DATA_STOCK_LIMIT=0` 跳。

        total = len(stock_codes)
        if total == 0:
            return {"status": "ok", "message": "无股票需要同步", "ok": 0, "fail": 0, "rows": 0, "elapsed_sec": 0.0}

        # ── 预估耗时：5 RPS 上限 ──
        est_sec = total / 5.0
        logger.info(f"[tick_trade] 计划: {total} 只, 预估 {est_sec:.0f}s ({est_sec/60:.1f}min)")
        self._progress(
            message=f"计划: {total} 只, 预估 {est_sec/60:.1f}min @5RPS",
            current=0, total=total, current_step="planning",
        )

        # ── 执行 ──
        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, mr, c6): c6 for c6 in stock_codes}
            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"[tick_trade {c6}] {res['error']}")
                except Exception as e:
                    fail_cnt += 1
                    logger.warning(f"[tick_trade {c6}] {e}")
                if i % 50 == 0 or i == total:
                    self._progress(
                        message=f"tick_trade 进度 {i}/{total} OK:{ok_cnt} FAIL:{fail_cnt} 行:{rows_total}",
                        current=i, total=total, current_step=c6,
                    )
        elapsed = round(time.time() - t0, 1)
        msg = f"逐笔 {ok_cnt}成 {fail_cnt}败 共{rows_total}行, {elapsed}s"
        return {
            "status": "ok" if fail_cnt == 0 else "warning",
            "message": msg,
            "ok": ok_cnt, "fail": fail_cnt, "rows": rows_total, "elapsed_sec": elapsed,
        }

    def _sync_one(self, mr, code6: str) -> dict:
        try:
            df = mr.fetch_tick_trade(code6)
        except Exception as e:
            return {"status": "fail", "error": str(e)}
        if df is None or df.empty:
            return {"status": "fail", "error": "no data (mairui 21:00 前可能为空)"}
        try:
            rows = [
                {
                    "stock_code": to_hermes(code6),
                    "trade_date": str(r["trade_date"]),
                    "trade_time": r["trade_time"].strftime("%Y-%m-%d %H:%M:%S")
                                 if hasattr(r["trade_time"], "strftime") else str(r["trade_time"]),
                    "price": float(r.get("price") or 0),
                    "volume": float(r.get("volume") or 0),
                    "direction_code": int(r.get("direction_code") or 0),
                    "direction": str(r.get("direction") or ""),
                    "amount": float(r.get("amount") or 0),
                }
                for _, r in df.iterrows()
            ]
            db_ops.upsert_tick_trade(rows)
            return {"status": "ok", "rows": len(rows)}
        except Exception as e:
            return {"status": "fail", "error": f"db write: {e}"}