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)}