xiaobai-review/xiaobai-datahub/datahub/admin_api.py

from __future__ import annotations

import json
from typing import Any

from datahub.adapters import RESERVED
from datahub.auth import AuthService
from datahub.db import HubDB
from datahub import lineage as lineage_module
from datahub.pipeline import OFFICIAL_DATASETS, STOCKS_DATASET, Pipeline
from datahub.scheduler import Scheduler
from datahub.serving import ApiError
from datahub import source_catalog
from datahub.timeutil import isoformat, now_shanghai, session_phase, yyyymmdd


class AdminAPI:
    def __init__(self, db: HubDB, pipeline: Pipeline, scheduler: Scheduler, auth: AuthService, ifind: Any = None) -> None:
        self.db = db
        self.pipeline = pipeline
        self.scheduler = scheduler
        self.auth = auth
        self.ifind = ifind

    def overview(self) -> dict[str, Any]:
        today = yyyymmdd(now_shanghai())
        cal = self.db.fetchone(
            "SELECT is_open FROM trade_calendar WHERE exchange = 'SSE' AND cal_date = ?",
            (today,),
        )
        is_open = bool(cal and int(cal["is_open"]) == 1)
        pubs = self.db.fetchall("SELECT * FROM publications WHERE trade_date = ?", (today,))
        failed = self.db.fetchall(
            "SELECT * FROM batches WHERE trade_date = ? AND state IN ('failed','staged')",
            (today,),
        )
        calls = self.db.fetchall(
            "SELECT * FROM src_calls ORDER BY id DESC LIMIT 20",
        )
        return {
            "trade_date": today,
            "session_phase": session_phase(now_shanghai(), is_open),
            "is_open_day": is_open,
            "eod_status": self.scheduler.eod_status(today),
            "revision_status": self.scheduler.revision_status(today),
            "publications": pubs,
            "anomalies": failed,
            "recent_calls": _public_calls(calls),
            "source_count": len(self.db.fetchall("SELECT provider FROM src_health")),
        }

    def sources(self) -> dict[str, Any]:
        health = {f"{row['provider']}:{row['endpoint_class']}": row for row in self.db.fetchall("SELECT * FROM src_health")}
        items = [
            {
                "provider": "tushare",
                "role": "official",
                "health": health.get("tushare:pro") or {"state": "unknown"},
                "credential": self.auth.credential_status("tushare_token") or {"configured": bool(self.pipeline.adapter.token)},
            }
        ]
        for name, adapter in RESERVED.items():
            if name == "ifind":
                live = self.ifind or adapter
                cred = self.auth.credential_status("ifind_refresh_token") or {
                    "configured": bool(getattr(live, "configured", False)),
                    "last4": "",
                    "updated_at": "",
                }
                items.append(
                    {
                        "provider": name,
                        "role": "licensed",
                        "health": live.probe(),
                        "credential": cred,
                    }
                )
                continue
            items.append(
                {
                    "provider": name,
                    "role": "reserved" if name in {"ths", "xgb", "akshare"} else "free",
                    "health": adapter.probe(),
                    "credential": {"configured": False, "last4": "", "updated_at": ""},
                }
            )
        # Prefer encrypted last4 if stored
        cred = self.auth.credential_status("tushare_token")
        if cred.get("configured"):
            items[0]["credential"] = cred
        elif self.pipeline.adapter.token:
            from datahub.crypto import mask_secret

            items[0]["credential"] = {"configured": True, "last4": mask_secret(self.pipeline.adapter.token), "updated_at": ""}
        return {"items": items}

    def probe(self, provider: str) -> dict[str, Any]:
        if provider == "tushare":
            return self.pipeline.adapter.probe()
        if provider == "ifind":
            adapter = self.ifind or RESERVED.get("ifind")
            if adapter is None:
                raise ApiError("INVALID_ARGUMENT", "unknown provider: ifind")
            return adapter.probe()
        adapter = RESERVED.get(provider)
        if adapter is None:
            raise ApiError("INVALID_ARGUMENT", f"unknown provider: {provider}")
        return adapter.probe()

    # ------------------------------------------------------------------
    # HEL-543: read-only side-channel status/catalog/lineage. These never
    # change routing, credentials, or adapters; they only read the
    # provider_call_log/provider_health tables (observability.py) plus the
    # static registries in source_catalog.py / lineage.py.
    # ------------------------------------------------------------------
    def providers_status(self, provider: str = "", limit: int = 50) -> dict[str, Any]:
        limit = max(1, min(int(limit or 50), 200))
        health_sql = "SELECT * FROM provider_health"
        params: tuple[Any, ...] = ()
        if provider:
            health_sql += " WHERE provider = ?"
            params = (provider,)
        health_sql += " ORDER BY provider, interface"
        health = self.db.fetchall(health_sql, params)
        calls_sql = "SELECT * FROM provider_call_log"
        if provider:
            calls_sql += " WHERE provider = ?"
        calls_sql += " ORDER BY id DESC LIMIT ?"
        recent = self.db.fetchall(calls_sql, (*params, limit))
        return {"health": health, "recent_calls": recent}

    def source_catalog(self) -> dict[str, Any]:
        return {"items": source_catalog.snapshot(self.db, self.auth)}

    def lineage(self, trade_date: str = "") -> dict[str, Any]:
        day = yyyymmdd(trade_date) if trade_date else yyyymmdd(now_shanghai())
        return {"trade_date": day, "items": lineage_module.snapshot(self.db, day)}

    def lineage_affected(self, provider: str = "", interface: str = "") -> dict[str, Any]:
        if not provider:
            raise ApiError("INVALID_ARGUMENT", "provider is required")
        return {"provider": provider, "interface": interface, "items": lineage_module.affected(self.db, provider, interface)}

    def jobs(self) -> dict[str, Any]:
        runs = self.db.fetchall("SELECT * FROM job_runs ORDER BY id DESC LIMIT 100")
        stocks_times = "/".join(self.pipeline.settings.stocks_refresh_times) or "20:00"
        return {
            "jobs": [
                {"id": "precheck", "at": "08:45", "title": "盘前预检"},
                {"id": "eod_a", "at": "15:05", "title": "盘后批 A daily/valuation/moneyflow/auction"},
                {"id": "eod_b", "at": "15:10", "title": "盘后批 B index_daily"},
                {"id": "eod_retry", "at": "15:15-23:30", "title": "盘后未出数自动重试（每 30 分钟，成功即停）"},
                {"id": "eod_revise", "at": "20:00-23:20", "title": "估值发布后复核（轻量比对，有修订才整组原子追补）"},
                {"id": "stocks_refresh", "at": stocks_times, "title": "股票主档刷新与正式发布（新上市/更名，无变化跳过）"},
                {"id": "history_backfill", "at": "manual", "title": "回补历史日历、个股日 K 与指数日 K"},
                {"id": "cleanup", "at": "00:30", "title": "清理 staging / 日志"},
                {"id": "backup", "at": "00:40", "title": "SQLite 备份"},
            ],
            "runs": runs,
        }

    def run_job(self, job_id: str, trade_date: str) -> dict[str, Any]:
        return self.scheduler.run_job(job_id, yyyymmdd(trade_date or now_shanghai()))

    def batches(self, date: str, dataset: str = "") -> dict[str, Any]:
        trade_date = yyyymmdd(date or now_shanghai())
        if dataset:
            rows = self.db.fetchall(
                "SELECT * FROM batches WHERE trade_date = ? AND dataset = ? ORDER BY started_at",
                (trade_date, dataset),
            )
        else:
            rows = self.db.fetchall(
                "SELECT * FROM batches WHERE trade_date = ? ORDER BY started_at",
                (trade_date,),
            )
        pubs = self.db.fetchall("SELECT * FROM publications WHERE trade_date = ?", (trade_date,))
        return {"trade_date": trade_date, "batches": rows, "publications": pubs}

    def datasets(self, date: str) -> dict[str, Any]:
        trade_date = yyyymmdd(date or now_shanghai())
        pubs = self.db.fetchall("SELECT * FROM publications WHERE trade_date = ?", (trade_date,))
        diffs = self.db.fetchall(
            "SELECT * FROM diff_reports WHERE trade_date = ? ORDER BY id",
            (trade_date,),
        )
        return {"trade_date": trade_date, "publications": pubs, "diff_reports": diffs}

    def audit(self) -> dict[str, Any]:
        return {"items": self.db.fetchall("SELECT * FROM audit_log ORDER BY id DESC LIMIT 200")}

    def rollback(self, dataset: str, trade_date: str, password: str, confirm: str, actor: str) -> dict[str, Any]:
        self._dangerous(password, confirm, f"{dataset}:{trade_date}")
        result = self.pipeline.rollback(dataset, trade_date, actor=actor)
        return result

    def backfill(self, dataset: str, trade_date: str, password: str, confirm: str, actor: str) -> dict[str, Any]:
        day = yyyymmdd(trade_date or now_shanghai())
        if dataset == "history":
            self._dangerous(password, confirm, "history:full")
            result = self.pipeline.backfill_history(day)
        else:
            self._dangerous(password, confirm, f"{dataset}:{day}")
            if dataset == "reference":
                result = self.pipeline.ingest_reference(day)
            elif dataset in OFFICIAL_DATASETS or dataset == STOCKS_DATASET:
                # Manual same-day republish must rebuild the full A/B boundary.
                # Gate failures and mid-switch exceptions both surface as
                # FAILED_PRECONDITION so the admin API never leaks raw
                # transaction errors to the client.
                try:
                    result = self.pipeline.force_republish_boundary(dataset, day)
                    failures = self.pipeline.eod_failures(result)
                    if failures:
                        raise ApiError("FAILED_PRECONDITION", "; ".join(failures))
                except ApiError:
                    raise
                except Exception as exc:
                    raise ApiError("FAILED_PRECONDITION", str(exc)) from exc
            else:
                raise ApiError("INVALID_ARGUMENT", f"unsupported backfill dataset: {dataset}")
        self.pipeline.audit(actor, "backfill", f"{dataset}:{day}", json.dumps({"ok": True}))
        return result

    def _dangerous(self, password: str, confirm: str, expected: str) -> None:
        if not self.auth.confirm_password(password):
            raise ApiError("UNAUTHORIZED", "二次确认密码错误")
        if confirm.strip() != expected:
            raise ApiError("INVALID_ARGUMENT", f"确认词必须为 {expected}")


def _public_calls(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    out = []
    for row in rows:
        out.append(
            {
                "id": row["id"],
                "provider": row["provider"],
                "endpoint": row["endpoint"],
                "ok": bool(row["ok"]),
                "latency_ms": row["latency_ms"],
                "error": row["error"],
                "created_at": row["created_at"],
            }
        )
    return out