market_sync/app/sources/baostock.py

"""Baostock 数据源。

提供：
  - 全市场股票基础信息（query_stock_basic）
  - 日 K 线（query_history_k_data_plus，免费，但 QPS 限制需要信号量）
  - 股票-行业映射（query_stock_industry）

免费、无 token，但需 bs.login() / bs.logout() 维护会话。
"""
from __future__ import annotations

import threading
from typing import Any, Optional

import pandas as pd

from app.core.datasource.base import DataSource
from app.core.datasource.utils import (
    code6_to_baostock,
    normalize_kline,
)


class BaostockSource(DataSource):
    key = "datasource_baostock"
    name = "Baostock"
    provides = ["kline_daily", "stock_basic", "industry"]
    requires_credential = False

    # 进程内只 login 一次
    _login_lock = threading.Lock()
    _logged_in = False
    _semaphore = threading.BoundedSemaphore(3)  # QPS 限流

    def _ensure_login(self) -> None:
        if self._logged_in:
            return
        with self._login_lock:
            if self._logged_in:
                return
            import baostock as bs
            bs.login()
            self._logged_in = True

    def is_available(self) -> tuple[bool, str]:
        try:
            self._ensure_login()
            return True, "ok"
        except Exception as e:
            return False, f"Baostock 登录失败: {e}"

    def health_check(self) -> dict[str, Any]:
        """连通性测试：复用 _ensure_login 避免重复 login 阻塞。"""
        try:
            self._ensure_login()
            import baostock as bs
            rs = bs.query_stock_basic(code="sh.600036")
            if rs.error_code != "0":
                return {"success": False, "message": f"Baostock 查询失败: {rs.error_msg}"}
            rows = []
            while rs.next():
                rows.append(rs.get_row_data())
            return {"success": True, "message": f"连接成功，查询到 {len(rows)} 条股票信息"}
        except Exception as e:
            return {"success": False, "message": f"Baostock 连接失败: {e}"}

    def fetch_stock_basic(self) -> list[dict[str, Any]]:
        """全市场 A 股基础信息。

        baostock query_stock_basic 返回字段（按当前接口）：
          code, code_name, ipoDate, outDate, type, status
        按字段名索引，防止 Baostock 改顺序。
        """
        import baostock as bs
        self._ensure_login()
        try:
            rs = bs.query_stock_basic()
            if rs.error_code != "0":
                raise RuntimeError(f"Baostock query_stock_basic 失败: {rs.error_msg}")
            fields = list(rs.fields)
            idx = {name: i for i, name in enumerate(fields)}
            rows = []
            while rs.next():
                row = rs.get_row_data()
                bs_code = row[idx["code"]]
                name = row[idx["code_name"]]
                stype = row[idx["type"]]
                status = row[idx["status"]]
                ipo = row[idx.get("ipoDate", -1)] if "ipoDate" in idx else ""
                if stype != "1" or status != "1":
                    continue
                exchange = bs_code.split(".")[0].upper()
                if exchange not in ("SH", "SZ", "BJ"):
                    continue
                code6 = bs_code.split(".")[-1]
                rows.append({
                    "code": f"{exchange}{code6}",
                    "name": name,
                    "exchange": exchange,
                    "list_date": ipo,
                    "listing_status": "st" if (name.startswith("ST") or name.startswith("*ST")) else "normal",
                })
            return rows
        except Exception:
            return []

    def fetch_kline_daily(self, code6: str, start: str, end: str) -> pd.DataFrame:
        import baostock as bs
        self._ensure_login()
        bs_code = code6_to_baostock(code6)
        with self._semaphore:
            try:
                rs = bs.query_history_k_data_plus(
                    bs_code,
                    "date,open,high,low,close,volume",
                    start_date=start, end_date=end,
                    frequency="d", adjustflag="2",  # 前复权
                )
                if rs.error_code != "0":
                    return pd.DataFrame()
                rows = []
                while rs.next():
                    rows.append(rs.get_row_data())
                if not rows:
                    return pd.DataFrame()
                df = pd.DataFrame(rows, columns=["trade_date", "open", "high", "low", "close", "volume"])
                return normalize_kline(df)
            except Exception:
                return pd.DataFrame()

    def fetch_industry_map(self) -> list[dict[str, Any]]:
        """全市场股票-行业映射。"""
        import baostock as bs
        self._ensure_login()
        try:
            rs = bs.query_stock_industry(code="", date="")
            if rs.error_code != "0":
                raise RuntimeError(f"query_stock_industry 失败: {rs.error_msg}")
            fields = list(rs.fields)  # 实际字段顺序：['updateDate', 'code', 'code_name', 'industry', 'industryClassification']
            idx = {name: i for i, name in enumerate(fields)}
            rows = []
            while rs.next():
                row = rs.get_row_data()
                industry = row[idx.get("industry", 3)] if "industry" in idx else ""
                if not industry:
                    continue
                bs_code = row[idx.get("code", 1)] if "code" in idx else ""
                code6 = bs_code.split(".")[-1]
                classification = row[idx.get("industryClassification", 4)] if "industryClassification" in idx else ""
                update = row[idx.get("updateDate", 0)] if "updateDate" in idx else ""
                rows.append({
                    "code": code6,
                    "industry_name": industry.strip(),
                    "industry_classification": (classification or "").strip(),
                    "update_date": update or "",
                })
            return rows
        except Exception:
            return []