xiaobai-review/xiaobai-datahub/datahub/hub.py

from __future__ import annotations

from pathlib import Path
from typing import Any

from datahub.adapters.ifind import IfindAdapter
from datahub.adapters.tushare import TushareAdapter
from datahub.admin_api import AdminAPI
from datahub.auth import AuthService
from datahub.crypto import SecretVault
from datahub.db import HubDB
from datahub.governance.circuit import CircuitBreaker
from datahub.governance.lkg import LastKnownGood
from datahub.governance.ratelimit import TokenBucket
from datahub.pipeline import Pipeline
from datahub.scheduler import Scheduler
from datahub.serving import V1API
from datahub.settings import Settings, load_settings
from datahub.siteauth import SiteAuth, SiteBridge
from datahub.siteconsole import SiteConsole


class Hub:
    def __init__(
        self,
        settings: Settings,
        adapter: TushareAdapter | None = None,
        site_auth: SiteAuth | None = None,
    ) -> None:
        if not settings.encryption_key:
            raise SystemExit("DATAHUB_ENCRYPTION_KEY 未配置")
        self.settings = settings
        self.db = HubDB(settings.db_path)
        # HEL-543: carry the observability kill switch on the db handle so
        # every call site that already threads `db` through (pipeline,
        # realtime_serve, steward, admin_api) picks it up for free with no
        # extra plumbing. Missing this attribute (e.g. a bare HubDB built
        # directly in tests) defaults to enabled — see observability.is_enabled.
        self.db.observability_enabled = settings.observability_enabled
        self.vault = SecretVault(settings.encryption_key)
        self.auth = AuthService(self.db, self.vault, settings.api_token)
        token = settings.tushare_token or self.auth.load_credential("tushare_token")
        if settings.tushare_token:
            self.auth.store_credential("tushare_token", settings.tushare_token)
            token = settings.tushare_token
        refresh = settings.ifind_refresh_token or self.auth.load_credential("ifind_refresh_token") or ""
        access = settings.ifind_access_token or self.auth.load_credential("ifind_access_token") or ""
        if settings.ifind_refresh_token:
            self.auth.store_credential("ifind_refresh_token", settings.ifind_refresh_token)
            refresh = settings.ifind_refresh_token
        if settings.ifind_access_token:
            self.auth.store_credential("ifind_access_token", settings.ifind_access_token)
            access = settings.ifind_access_token
        self.adapter = adapter or TushareAdapter(token)
        self.ifind = IfindAdapter(refresh, access)
        self.pipeline = Pipeline(
            self.db,
            self.adapter,
            settings,
            bucket=TokenBucket(settings.tushare_rate_per_minute),
            breaker=CircuitBreaker(),
        )
        self.lkg = LastKnownGood(self.db)
        self.scheduler = Scheduler(self.db, self.pipeline)
        self.api = V1API(self.db, self.pipeline, settings, ifind=self.ifind)
        self.site_bridge = SiteBridge(settings.review_base_url, settings.hub_admin_token)
        # Tests inject a stub so the console gate can run without a live site.
        self.site_auth = site_auth or SiteAuth(self.site_bridge, settings.encryption_key)
        self.site_console = SiteConsole(self.site_bridge)
        self.admin = AdminAPI(
            self.db,
            self.pipeline,
            self.scheduler,
            self.auth,
            ifind=self.ifind,
            site_auth=self.site_auth,
        )
        self.static_dir = Path(__file__).resolve().parents[1] / "admin"

    def put_tushare_credentials(self, body: dict[str, Any] | None) -> dict[str, Any]:
        """Store and hot-swap the Tushare token without a restart.

        The adapter holds the token in memory, so writing the credential and
        pushing it into the live adapter in one step is what makes the save take
        effect for the next fetch instead of the next deploy.
        """
        from datahub.serving import envelope

        token = str((body or {}).get("tushare_token") or "").strip()
        if not token:
            raise ValueError("Tushare Token 不能为空")
        self.auth.store_credential("tushare_token", token)
        self.adapter.token = token
        return envelope(self.adapter.probe(), {"source": "tushare"})

    def put_ifind_credentials(self, body: dict[str, Any] | None) -> dict[str, Any]:
        from datahub.serving import envelope

        payload = dict(body or {})
        refresh = str(payload.get("ifind_refresh_token") or "").strip()
        access = str(payload.get("ifind_access_token") or "").strip()
        self.auth.store_credential("ifind_refresh_token", refresh)
        self.auth.store_credential("ifind_access_token", access)
        self.ifind.set_credentials(refresh, access)
        return envelope(self.ifind.status(), {"source": "ifind"})

    def start(self) -> None:
        if self.settings.scheduler_enabled:
            self.scheduler.start()

    def stop(self) -> None:
        self.scheduler.stop()


def build_hub(settings: Settings | None = None) -> Hub:
    return Hub(settings or load_settings())