xiaobai-review/backend/data/datahub/bridge.py

from __future__ import annotations

import logging
import sys
from typing import Any, Callable

from backend.data.datahub.client import DatahubClient, DatahubResponse
from backend.data.datahub.compare import compare_rows
from backend.data.datahub.errors import DatahubError
from backend.data.datahub.native import (
    API_TO_DATASET,
    filter_calendar_rows,
    filter_stock_rows,
    project_fields,
    to_native_rows,
    yyyymmdd,
)
from backend.data.datahub.redact import redact_text, redact_value
from backend.data.datahub.settings import DatahubSettings
from backend.data.providers.tushare_client import TushareClient

LOGGER = logging.getLogger("xiaobai.datahub")
ShadowSink = Callable[[dict[str, Any]], None]
EMPTY_FAIL_DATASETS = {"stocks", "daily", "index_daily", "valuation", "moneyflow", "auction"}


def looks_like_heaven(module_name: str, filename: str = "") -> bool:
    """问天调用栈识别。问天未永久冻结，只是本阶段仍走旧 Tushare 链路。"""
    path = filename.replace("\\", "/")
    return module_name.startswith("backend.features.heaven") or "/features/heaven/" in path


def caller_is_heaven(depth: int = 24) -> bool:
    frame = sys._getframe(1)
    for _ in range(depth):
        frame = frame.f_back if frame is not None else None
        if frame is None:
            return False
        name = str(frame.f_globals.get("__name__") or "")
        filename = str(frame.f_code.co_filename or "")
        if looks_like_heaven(name, filename):
            return True
    return False


class DatahubBridge:
    def __init__(
        self,
        settings: DatahubSettings,
        client: DatahubClient,
        shadow_sink: ShadowSink | None = None,
        heaven_guard: Callable[[], bool] | None = None,
    ) -> None:
        self.settings = settings
        self.client = client
        self.shadow_sink = shadow_sink
        self.heaven_guard = heaven_guard or caller_is_heaven

    def dataset_status(self, trade_date: str) -> list[dict[str, Any]] | None:
        flags = self.settings.flags("status")
        if not flags.read and not flags.shadow:
            return None
        try:
            response = self._require_fresh(self.client.dataset_status(yyyymmdd(trade_date)), "status")
            rows = list(response.data or [])
            if flags.shadow:
                self._emit_shadow(compare_rows("status", [], rows, response.meta))
            if flags.read:
                return rows
            return None
        except Exception as exc:
            self._log_failure("status", exc)
            if flags.shadow:
                self._emit_shadow(compare_rows("status", [], [], {}, self._error_text(exc)))
            return None

    def batches(self, trade_date: str, dataset: str = "") -> list[dict[str, Any]] | None:
        flags = self.settings.flags("status")
        if not flags.read:
            return None
        try:
            response = self._require_fresh(
                self.client.batches(yyyymmdd(trade_date), dataset),
                "status",
            )
            return list(response.data or [])
        except Exception as exc:
            self._log_failure("status", exc)
            return None

    def query(
        self,
        api_name: str,
        params: dict[str, Any] | None,
        fields: str,
        legacy_query: Callable[..., list[dict[str, Any]]],
    ) -> list[dict[str, Any]]:
        dataset = API_TO_DATASET.get(api_name)
        # 问天允许后续纳入 datahub；首批只读接入仍保持旧链路，避免误切。
        if not dataset or self.heaven_guard():
            return legacy_query(api_name, params, fields)
        flags = self.settings.flags(dataset)
        if not flags.read and not flags.shadow:
            return legacy_query(api_name, params, fields)

        hub_rows: list[dict[str, Any]] | None = None
        hub_meta: dict[str, Any] = {}
        hub_error: str | None = None
        hub_canonical: list[dict[str, Any]] = []
        try:
            response = self._fetch_dataset(dataset, params or {})
            hub_canonical = self._extract_rows(dataset, response, params or {})
            hub_rows = to_native_rows(dataset, hub_canonical)
            hub_meta = dict(response.meta)
            self._validate_usable(dataset, hub_rows, response)
        except Exception as exc:
            hub_error = self._error_text(exc)
            self._log_failure(dataset, exc)

        if flags.shadow:
            try:
                legacy_rows = legacy_query(api_name, params, fields)
            except Exception as exc:
                if flags.read and hub_rows is not None and hub_error is None:
                    self._emit_shadow(compare_rows(dataset, [], hub_canonical, hub_meta, self._error_text(exc)))
                    return project_fields(hub_rows, fields)
                raise
            self._emit_shadow(compare_rows(dataset, legacy_rows, hub_canonical, hub_meta, hub_error))
            if flags.read and hub_rows is not None and hub_error is None:
                return project_fields(hub_rows, fields)
            return legacy_rows

        if flags.read and hub_rows is not None and hub_error is None:
            return project_fields(hub_rows, fields)
        return legacy_query(api_name, params, fields)

    def _fetch_dataset(self, dataset: str, params: dict[str, Any]) -> DatahubResponse:
        date = yyyymmdd(params.get("trade_date") or params.get("date"))
        start = yyyymmdd(params.get("start_date") or params.get("from") or date)
        end = yyyymmdd(params.get("end_date") or params.get("to") or date)
        code = str(params.get("ts_code") or params.get("code") or "").strip()
        if dataset == "calendar":
            if not start or not end:
                raise DatahubError("INVALID_ARGUMENT", "calendar requires start_date and end_date")
            return self.client.calendar(start, end)
        if dataset == "stocks":
            return self._paginate(self.client.stocks, {})
        fetchers = {
            "daily": self.client.daily_bars,
            "index_daily": self.client.index_bars,
            "valuation": self.client.valuation,
            "moneyflow": self.client.moneyflow,
            "auction": self.client.auction,
        }
        fetcher = fetchers[dataset]
        query: dict[str, Any] = {}
        if code:
            query["code"] = code
        if date and not (params.get("start_date") or params.get("end_date")):
            query["date"] = date
        else:
            if start:
                query["from"] = start
            if end:
                query["to"] = end
        if dataset == "daily":
            query["adjust"] = "none"
        return self._paginate(fetcher, query)

    def _paginate(self, fetcher: Callable[..., DatahubResponse], params: dict[str, Any]) -> DatahubResponse:
        limit = self.settings.page_limit
        offset = 0
        rows: list[Any] = []
        meta: dict[str, Any] = {}
        schema_version = 1
        while True:
            page = fetcher(**{**params, "limit": limit, "offset": offset})
            meta = dict(page.meta)
            schema_version = page.schema_version
            data = page.data or []
            if not isinstance(data, list):
                raise DatahubError("INTERNAL", "datahub returned a non-list payload")
            rows.extend(data)
            if len(data) < limit:
                break
            offset += limit
            if offset > 200_000:
                break
        return DatahubResponse(data=rows, meta=meta, schema_version=schema_version)

    def _extract_rows(
        self,
        dataset: str,
        response: DatahubResponse,
        params: dict[str, Any],
    ) -> list[dict[str, Any]]:
        rows = [dict(item) for item in (response.data or [])]
        if dataset == "calendar":
            return filter_calendar_rows(rows, params)
        if dataset == "stocks":
            return filter_stock_rows(rows, params)
        return rows

    def _validate_usable(self, dataset: str, rows: list[dict[str, Any]], response: DatahubResponse) -> None:
        meta = response.meta or {}
        stale_seconds = int(meta.get("staleness_seconds") or 0)
        if meta.get("stale") or stale_seconds > self.settings.stale_seconds_max:
            raise DatahubError("STALE", f"{dataset} data is stale")
        if dataset in EMPTY_FAIL_DATASETS and not rows:
            raise DatahubError("EMPTY", f"{dataset} returned no rows")

    def _require_fresh(self, response: DatahubResponse, dataset: str) -> DatahubResponse:
        self._validate_usable(dataset, list(response.data or []) if isinstance(response.data, list) else [], response)
        return response

    def _emit_shadow(self, report: dict[str, Any]) -> None:
        safe = redact_value(report, secrets=self.settings.secrets())
        LOGGER.info("datahub shadow %s", safe)
        if self.shadow_sink is not None:
            self.shadow_sink(report)

    def _log_failure(self, dataset: str, exc: Exception) -> None:
        LOGGER.warning(
            "datahub fallback dataset=%s error=%s",
            dataset,
            redact_text(self._error_text(exc), self.settings.secrets()),
        )

    def _error_text(self, exc: Exception) -> str:
        if isinstance(exc, DatahubError):
            text = f"{exc.code}: {exc.message}"
        else:
            text = str(exc)
        return redact_text(text, self.settings.secrets())


class DatahubAwareTushareClient:
    def __init__(self, legacy: TushareClient, bridge: DatahubBridge) -> None:
        self._legacy = legacy
        self._bridge = bridge

    def query(
        self,
        api_name: str,
        params: dict[str, Any] | None = None,
        fields: str = "",
    ) -> list[dict[str, Any]]:
        return self._bridge.query(api_name, params, fields, self._legacy.query)

    def __getattr__(self, name: str) -> Any:
        return getattr(self._legacy, name)