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

from __future__ import annotations

import json
import logging
import urllib.error
import urllib.parse
import urllib.request
from dataclasses import dataclass, field
from typing import Any, Callable

from backend.data.datahub.errors import DatahubError
from backend.data.datahub.redact import redact_text
from backend.data.datahub.settings import DatahubSettings

LOGGER = logging.getLogger("xiaobai.datahub")
UrlOpen = Callable[..., Any]


@dataclass(frozen=True)
class DatahubResponse:
    data: Any
    meta: dict[str, Any] = field(default_factory=dict)
    schema_version: int = 1
    status: int = 200


class DatahubClient:
    def __init__(
        self,
        settings: DatahubSettings,
        urlopen: UrlOpen = urllib.request.urlopen,
    ) -> None:
        self.settings = settings
        self._urlopen = urlopen

    def health(self) -> DatahubResponse:
        return self.get("/v1/health")

    def calendar(self, start: str, end: str) -> DatahubResponse:
        return self.get("/v1/calendar", {"from": start, "to": end})

    def stocks(self, updated_since: str = "", limit: int | None = None, offset: int = 0) -> DatahubResponse:
        params: dict[str, Any] = {"offset": offset, "limit": limit or self.settings.page_limit}
        if updated_since:
            params["updated_since"] = updated_since
        return self.get("/v1/stocks", params)

    def daily_bars(self, **params: Any) -> DatahubResponse:
        return self.get("/v1/bars/daily", params)

    def index_bars(self, **params: Any) -> DatahubResponse:
        return self.get("/v1/indexes/bars", params)

    def valuation(self, **params: Any) -> DatahubResponse:
        return self.get("/v1/valuation", params)

    def moneyflow(self, **params: Any) -> DatahubResponse:
        return self.get("/v1/moneyflow", params)

    def auction(self, **params: Any) -> DatahubResponse:
        return self.get("/v1/auction", params)

    def dataset_status(self, date: str) -> DatahubResponse:
        return self.get("/v1/datasets/status", {"date": date})

    def batches(self, date: str, dataset: str = "") -> DatahubResponse:
        params: dict[str, Any] = {"date": date}
        if dataset:
            params["dataset"] = dataset
        return self.get("/v1/batches", params)

    def get(self, path: str, params: dict[str, Any] | None = None) -> DatahubResponse:
        if not self.settings.token:
            raise DatahubError("NOT_CONFIGURED", "DATAHUB_TOKEN is not configured")
        query = {
            key: value
            for key, value in (params or {}).items()
            if value is not None and value != ""
        }
        url = self.settings.base_url + path
        if query:
            url = f"{url}?{urllib.parse.urlencode(query)}"
        attempts = 1 + max(0, self.settings.retries)
        last_error: DatahubError | None = None
        for attempt in range(attempts):
            try:
                return self._request(url)
            except DatahubError as exc:
                last_error = exc
                if exc.code not in {"TIMEOUT", "UNAVAILABLE"} or attempt + 1 >= attempts:
                    raise
                LOGGER.warning(
                    "datahub retry %s/%s %s",
                    attempt + 1,
                    attempts,
                    redact_text(str(exc), self.settings.secrets()),
                )
        raise last_error or DatahubError("INTERNAL", "datahub request failed")

    def _request(self, url: str) -> DatahubResponse:
        request = urllib.request.Request(
            url,
            headers={
                "Accept": "application/json",
                "X-Datahub-Token": self.settings.token,
                "User-Agent": "XiaobaiReviewDatahub/1.0",
            },
            method="GET",
        )
        try:
            with self._urlopen(request, timeout=self.settings.timeout_seconds) as response:
                status = int(getattr(response, "status", 200) or 200)
                raw = response.read().decode("utf-8")
        except TimeoutError as exc:
            raise DatahubError("TIMEOUT", "datahub request timed out") from exc
        except urllib.error.HTTPError as exc:
            body = _read_error_body(exc)
            raise _http_error(exc.code, body, self.settings.secrets()) from exc
        except urllib.error.URLError as exc:
            reason = redact_text(str(getattr(exc, "reason", exc)), self.settings.secrets())
            if "timed out" in reason.lower():
                raise DatahubError("TIMEOUT", "datahub request timed out") from exc
            raise DatahubError("UNAVAILABLE", f"datahub unavailable: {reason}") from exc
        payload = _parse_json(raw, self.settings.secrets())
        return _as_response(payload, status, self.settings.secrets())


def _parse_json(raw: str, secrets: tuple[str, ...]) -> dict[str, Any]:
    try:
        payload = json.loads(raw)
    except json.JSONDecodeError as exc:
        raise DatahubError("INTERNAL", "datahub returned invalid json") from exc
    if not isinstance(payload, dict):
        raise DatahubError("INTERNAL", "datahub returned a non-object payload")
    return payload


def _as_response(payload: dict[str, Any], status: int, secrets: tuple[str, ...]) -> DatahubResponse:
    error = payload.get("error")
    if isinstance(error, dict):
        raise _mapped_error(str(error.get("code") or "INTERNAL"), str(error.get("message") or "datahub error"), status)
    if status >= 400:
        raise DatahubError("UNAVAILABLE", f"datahub http {status}", status)
    return DatahubResponse(
        data=payload.get("data"),
        meta=dict(payload.get("meta") or {}),
        schema_version=int(payload.get("schema_version") or 1),
        status=status,
    )


def _http_error(status: int, payload: dict[str, Any], secrets: tuple[str, ...]) -> DatahubError:
    error = payload.get("error") if isinstance(payload.get("error"), dict) else {}
    code = str((error or {}).get("code") or "")
    message = str((error or {}).get("message") or payload.get("message") or f"datahub http {status}")
    message = redact_text(message, secrets)
    if status == 401 or code == "UNAUTHORIZED":
        return DatahubError("UNAUTHORIZED", message, status)
    if status == 404 or code == "DATASET_NOT_PUBLISHED":
        return DatahubError("DATASET_NOT_PUBLISHED", message, status)
    if status == 400 or code == "INVALID_ARGUMENT":
        return DatahubError("INVALID_ARGUMENT", message, status)
    if status in {429, 503} or code in {"RATE_LIMITED", "SOURCE_UNAVAILABLE"}:
        return DatahubError("UNAVAILABLE", message, status)
    return DatahubError(code or "INTERNAL", message, status)


def _mapped_error(code: str, message: str, status: int) -> DatahubError:
    if code == "STALE_DATA":
        return DatahubError("STALE", message, status)
    if code in {"UNAUTHORIZED", "DATASET_NOT_PUBLISHED", "INVALID_ARGUMENT"}:
        return DatahubError(code, message, status)
    if code in {"RATE_LIMITED", "SOURCE_UNAVAILABLE"}:
        return DatahubError("UNAVAILABLE", message, status)
    return DatahubError(code or "INTERNAL", message, status)


def _read_error_body(exc: urllib.error.HTTPError) -> dict[str, Any]:
    try:
        raw = exc.read().decode("utf-8")
        payload = json.loads(raw)
        return payload if isinstance(payload, dict) else {"message": raw}
    except Exception:
        return {"message": str(exc)}