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 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,
            "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")
        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": "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]:
        self._dangerous(password, confirm, f"{dataset}:{trade_date}")
        if dataset == "reference":
            result = self.pipeline.ingest_reference(trade_date)
        else:
            result = self.pipeline.run_dataset(dataset, trade_date)
        self.pipeline.audit(actor, "backfill", f"{dataset}:{trade_date}", 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