market_sync/app/tasks/task_share_snapshot.py

"""同步任务：股本快照（雪球 quote_detail）。

写入两张表：
  - stocks.total_share / float_share / share_updated_at（基础信息表）
  - share（独立股本时序表）
"""
from __future__ import annotations

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

from app.core.config import settings
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
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.share")


class SyncShareSnapshot(SyncTask):
    dataset_id = "share_snapshot"

    def _run(
        self,
        *,
        trigger_source: str = "manual",
        codes: list[str] | None = None,
        max_workers: int = 10,
        **kwargs,
    ) -> dict[str, Any]:
        if not settings.xueqiu_token:
            mark_sync_blocked(self.dataset_id, message="未配置 XUEQIU_TOKEN")
            return {"status": "blocked", "message": "未配置 XUEQIU_TOKEN"}

        xq = ds_registry.get("datasource_xueqiu")
        if xq is None:
            return {"status": "error", "message": "雪球数据源未注册"}
        ok, reason = is_source_ready(xq.key)
        if not ok:
            mark_sync_blocked(self.dataset_id, message=f"雪球未就绪: {reason}")
            return {"status": "blocked", "message": f"雪球未就绪: {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": "无股票代码"}

        # 跳过今日已刷过的
        today = time.strftime("%Y-%m-%d")
        existing = {s["code"]: s for s in db_ops.fetch_all_stocks()}
        todo = []
        for c6 in stock_codes:
            old = existing.get(c6, {})
            if str(old.get("share_updated_at", "") or "") == today and old.get("total_share", 0) > 0:
                continue
            todo.append(c6)
        skip = len(stock_codes) - len(todo)
        if not todo:
            return {"status": "ok", "message": f"全部 {skip} 只已最新，跳过"}

        total = len(todo)
        ok_cnt = fail_cnt = 0
        rows_to_share_table: list[dict] = []
        t0 = time.time()
        with ThreadPoolExecutor(max_workers=max_workers) as pool:
            futures = {pool.submit(xq.fetch_share_snapshot, c6): c6 for c6 in todo}
            for i, future in enumerate(as_completed(futures), 1):
                c6 = futures[future]
                try:
                    r = future.result()
                except Exception as e:
                    fail_cnt += 1
                    logger.warning(f"[share {c6}] 异常: {e}")
                    continue
                if r is None:
                    fail_cnt += 1
                    continue
                # 写 stocks
                hermes = self._to_hermes(c6)
                db_ops.update_stock_share_snapshot(
                    code=hermes,
                    total_share=r["total_share"],
                    float_share=r["float_share"],
                    trade_date=r.get("trade_date", today),
                )
                # 也写 share 表
                rows_to_share_table.append({
                    "stock_code": c6,
                    "trade_date": r.get("trade_date", today),
                    "total_share": r["total_share"],
                    "float_share": r["float_share"],
                })
                ok_cnt += 1
                if i % 50 == 0 or i == total:
                    self._progress(
                        message=f"股本 {i}/{total} OK:{ok_cnt} FAIL:{fail_cnt}",
                        current=i, total=total, current_step=c6,
                    )
        # 批量写 share 表
        if rows_to_share_table:
            try:
                db_ops.upsert_share(rows_to_share_table)
            except Exception as e:
                logger.warning(f"写 share 表失败: {e}")

        elapsed = round(time.time() - t0, 1)
        msg = f"股本 {ok_cnt}成 {fail_cnt}败 {skip}跳, {elapsed}s"
        return {
            "status": "ok" if fail_cnt == 0 else "warning",
            "message": msg,
            "ok": ok_cnt, "fail": fail_cnt, "skip": skip,
            "elapsed_sec": elapsed,
        }

    @staticmethod
    def _to_hermes(code6: str) -> str:
        if code6.startswith(("5", "6", "9")):
            return f"SH{code6}"
        if code6.startswith(("4", "8")):
            return f"BJ{code6}"
        return f"SZ{code6}"