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.pipeline import OFFICIAL_DATASETS, STOCKS_DATASET, Pipeline
from datahub.scheduler import Scheduler
from datahub.serving import ApiError
from datahub.timeutil import isoformat, now_shanghai, session_phase, yyyymmdd


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

    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),
            "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():
            items.append(
                {
                    "provider": name,
                    "role": "reserved",
                    "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()
        adapter = RESERVED.get(provider)
        if adapter is None:
            raise ApiError("INVALID_ARGUMENT", f"unknown provider: {provider}")
        return adapter.probe()

    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": "stocks_refresh", "at": stocks_times, "title": "股票主档刷新与正式发布（新上市/更名，无变化跳过）"},
                {"id": "history_backfill", "at": "manual", "title": "回补历史日历与指数日 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.
                result = self.pipeline.force_republish_boundary(dataset, day)
                failures = self.pipeline.eod_failures(result)
                if failures:
                    raise ApiError("FAILED_PRECONDITION", "; ".join(failures))
            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