market_sync/app/tasks/task_industry_sector.py

"""同步任务：股票-行业映射（Baostock）。

数据流：Baostock query_stock_industry() → 写 industry / sectors / stock_sector_map 三张表。"""
from __future__ import annotations

import time
from typing import Any

from sqlalchemy import update

from app.core.db import ops as db_ops
from app.core.db.models import Stock
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 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.industry_sector")


def _build_sector_key(taxonomy: str, sector_name: str) -> str:
    tax = (taxonomy or "").strip() or "default"
    name = (sector_name or "").strip()
    return f"{tax}::{name}"


class SyncIndustrySector(SyncTask):
    dataset_id = "industry_sector"

    def _run(self, *, trigger_source: str = "manual", **kwargs) -> dict[str, Any]:
        bs = ds_registry.get("datasource_baostock")
        if bs is None:
            return {"status": "error", "message": "Baostock 数据源未注册"}
        ok, reason = is_source_ready(bs.key)
        if not ok:
            mark_sync_blocked(self.dataset_id, message=f"Baostock 未就绪: {reason}")
            return {"status": "blocked", "message": f"Baostock 未就绪: {reason}"}

        self._progress(message="拉取股票-行业映射...")
        t0 = time.time()
        try:
            raw_rows = bs.fetch_industry_map()
        except Exception as e:
            return {"status": "error", "message": f"拉取失败: {e}"}

        if not raw_rows:
            return {"status": "warning", "message": "Baostock 未返回行业数据"}

        # 合并：sectors + stock_sector_map + industry
        sectors_rows: dict[str, dict] = {}
        stock_sector_rows: list[dict] = []
        industry_rows: list[dict] = []

        for r in raw_rows:
            code6 = str(r["code"]).zfill(6)
            name = (r.get("industry_name") or "").strip()
            if not name:
                continue
            classification = (r.get("industry_classification") or "").strip()
            sector_key = _build_sector_key(classification, name)

            sectors_rows[sector_key] = {
                "sector_key": sector_key,
                "sector_name": name,
                "taxonomy": classification,
                "level": "",
                "source": "baostock_query_stock_industry",
                "enabled": 1,
            }
            stock_sector_rows.append({"stock_code": to_hermes(code6), "sector_key": sector_key})
            industry_rows.append({
                "code": to_hermes(code6),
                "industry_name": name,
                "industry_classification": classification,
                "update_date": r.get("update_date", ""),
            })

        # 写库（全量替换）
        try:
            db_ops.replace_all_sectors(list(sectors_rows.values()))
            db_ops.replace_all_stock_sector_map(stock_sector_rows)
            db_ops.replace_all_industries(industry_rows)
        except Exception as e:
            return {"status": "error", "message": f"写库失败: {e}"}

        # 回填 stocks.industry 字段（让 stocks 表也能直接看到）
        # 走 ORM update(Stock) + IN (SHxxx, SZxxx, BJxxx, ...)
        try:
            # 按 industry_name 分组，对每组用一个 IN 批量更新
            by_industry: dict[str, list[str]] = {}
            for r in industry_rows:
                name = r["industry_name"]
                code = r["code"]
                if name not in by_industry:
                    by_industry[name] = []
                by_industry[name].append(code)
            for ind_name, codes in by_industry.items():
                # 6 位 code × 3 个交易所前缀（SH/SZ/BJ）
                code_variants = [
                    f"{prefix}{c}" for c in codes for prefix in ("SH", "SZ", "BJ")
                ]
                with db_ops.get_session() as s:
                    s.execute(
                        update(Stock)
                        .where(Stock.code.in_(code_variants))
                        .values(industry=ind_name)
                    )
        except Exception as e:
            logger.warning(f"回填 stocks.industry 失败（不影响主任务）: {e}")

        elapsed = round(time.time() - t0, 1)
        msg = f"行业 {len(sectors_rows)}个 映射{len(stock_sector_rows)}条, {elapsed}s"
        return {
            "status": "ok",
            "message": msg,
            "sectors": len(sectors_rows),
            "stock_sector_map": len(stock_sector_rows),
            "elapsed_sec": elapsed,
        }