"""同步任务:股票-行业映射(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, }