xiaobai-review/xiaobai-datahub/datahub/serving.py

from __future__ import annotations

from http import HTTPStatus
from typing import Any
from urllib.parse import parse_qs

from datahub import SCHEMA_VERSION
from datahub.codes import resolve_code
from datahub.db import HubDB
from datahub.normalize import qfq_bar
from datahub.numbers import finite_number
from datahub.pipeline import Pipeline
from datahub.settings import Settings
from datahub.timeutil import isoformat, now_shanghai, session_phase, yyyymmdd

ERROR_STATUS = {
    "UNAUTHORIZED": HTTPStatus.UNAUTHORIZED,
    "INVALID_ARGUMENT": HTTPStatus.BAD_REQUEST,
    "RATE_LIMITED": HTTPStatus.TOO_MANY_REQUESTS,
    "SOURCE_UNAVAILABLE": HTTPStatus.SERVICE_UNAVAILABLE,
    "DATASET_NOT_PUBLISHED": HTTPStatus.NOT_FOUND,
    "STALE_DATA": HTTPStatus.OK,
    "INTERNAL": HTTPStatus.INTERNAL_SERVER_ERROR,
}


class ApiError(Exception):
    def __init__(self, code: str, message: str, retry_after: int | None = None, extra: dict[str, Any] | None = None) -> None:
        super().__init__(message)
        self.code = code
        self.message = message
        self.retry_after = retry_after
        self.extra = extra or {}

    def payload(self) -> dict[str, Any]:
        body: dict[str, Any] = {"code": self.code, "message": self.message}
        if self.retry_after is not None:
            body["retry_after"] = self.retry_after
        body.update(self.extra)
        return {"error": body}

    @property
    def status(self) -> HTTPStatus:
        return ERROR_STATUS.get(self.code, HTTPStatus.INTERNAL_SERVER_ERROR)


def envelope(data: Any, meta: dict[str, Any]) -> dict[str, Any]:
    return {"schema_version": SCHEMA_VERSION, "data": data, "meta": meta}


class V1API:
    def __init__(self, db: HubDB, pipeline: Pipeline, settings: Settings) -> None:
        self.db = db
        self.pipeline = pipeline
        self.settings = settings

    def handle(self, path: str, query: dict[str, list[str]]) -> dict[str, Any]:
        q = {key: values[-1] if values else "" for key, values in query.items()}
        if path == "/v1/health":
            return self.health()
        if path == "/v1/calendar":
            return self.calendar(q.get("from") or "", q.get("to") or "")
        if path == "/v1/stocks":
            return self.stocks(q.get("updated_since") or "", q)
        if path == "/v1/bars/daily":
            return self.daily_bars(q)
        if path == "/v1/indexes/bars":
            return self.index_bars(q)
        if path == "/v1/valuation":
            return self.valuation(q)
        if path == "/v1/moneyflow":
            return self.moneyflow(q)
        if path == "/v1/auction":
            return self.auction(q)
        if path == "/v1/datasets/status":
            return self.dataset_status(q.get("date") or "")
        if path == "/v1/batches":
            return self.batches(q.get("date") or "", q.get("dataset") or "")
        raise ApiError("INVALID_ARGUMENT", f"unknown endpoint: {path}")

    def health(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 cal["is_open"] == 1)
        sources = self.db.fetchall("SELECT * FROM src_health")
        return envelope(
            {
                "status": "ok",
                "session_phase": session_phase(now_shanghai(), is_open),
                "trade_date": today,
                "is_open_day": is_open,
                "sources": [
                    {
                        "provider": row["provider"],
                        "endpoint_class": row["endpoint_class"],
                        "state": row["state"],
                        "last_ok_at": row["last_ok_at"],
                        "consec_failures": row["consec_failures"],
                    }
                    for row in sources
                ],
            },
            {"tier": "official", "trade_date": today, "source": "datahub", "stale": False, "staleness_seconds": 0},
        )

    def calendar(self, start: str, end: str) -> dict[str, Any]:
        start = yyyymmdd(start or add_default(-30))
        end = yyyymmdd(end or add_default(5))
        rows = self.db.fetchall(
            """
            SELECT cal_date, is_open, pretrade_date,
                   (SELECT MAX(cal_date) FROM trade_calendar t2
                     WHERE t2.exchange = 'SSE' AND t2.is_open = 1 AND t2.cal_date < t1.cal_date) AS prev_open
            FROM trade_calendar t1
            WHERE exchange = 'SSE' AND cal_date >= ? AND cal_date <= ?
            ORDER BY cal_date
            """,
            (start, end),
        )
        items = [
            {
                "cal_date": row["cal_date"],
                "is_open": bool(row["is_open"]),
                "pretrade_date": row["pretrade_date"],
                "prev_open": row["prev_open"],
            }
            for row in rows
        ]
        return envelope(items, self._official_meta("calendar", end if items else start, source="tushare:trade_cal"))

    def stocks(self, updated_since: str, q: dict[str, str]) -> dict[str, Any]:
        limit, offset = self._page(q)
        if updated_since:
            rows = self.db.fetchall(
                "SELECT * FROM stock_master WHERE updated_at >= ? ORDER BY ts_code LIMIT ? OFFSET ?",
                (updated_since, limit, offset),
            )
        else:
            rows = self.db.fetchall(
                "SELECT * FROM stock_master ORDER BY ts_code LIMIT ? OFFSET ?",
                (limit, offset),
            )
        return envelope(rows, self._official_meta("stocks", yyyymmdd(), source="tushare:stock_basic"))

    def daily_bars(self, q: dict[str, str]) -> dict[str, Any]:
        return self._published_rows(
            dataset="daily",
            table="eod_bars",
            q=q,
            source="tushare:daily",
            adjust=q.get("adjust") or "none",
        )

    def index_bars(self, q: dict[str, str]) -> dict[str, Any]:
        return self._published_rows(
            dataset="index_daily",
            table="eod_index_bars",
            q=q,
            source="tushare:index_daily",
            default_code="000001.SH",
        )

    def valuation(self, q: dict[str, str]) -> dict[str, Any]:
        return self._published_rows(dataset="valuation", table="eod_valuation", q=q, source="tushare:daily_basic")

    def moneyflow(self, q: dict[str, str]) -> dict[str, Any]:
        return self._published_rows(dataset="moneyflow", table="eod_moneyflow", q=q, source="tushare:moneyflow")

    def auction(self, q: dict[str, str]) -> dict[str, Any]:
        return self._published_rows(dataset="auction", table="eod_auction", q=q, source="tushare:stk_auction")

    def dataset_status(self, date: str) -> dict[str, Any]:
        trade_date = yyyymmdd(date or now_shanghai())
        datasets = ("daily", "valuation", "moneyflow", "auction", "index_daily")
        items = []
        for dataset in datasets:
            pub = self.db.fetchone(
                "SELECT * FROM publications WHERE dataset = ? AND trade_date = ?",
                (dataset, trade_date),
            )
            batch = None
            if pub:
                batch = self.db.fetchone("SELECT * FROM batches WHERE batch_id = ?", (pub["active_batch"],))
            items.append(
                {
                    "dataset": dataset,
                    "trade_date": trade_date,
                    "state": (pub or {}).get("state") or "unpublished",
                    "batch_id": (pub or {}).get("active_batch"),
                    "published_at": (pub or {}).get("published_at"),
                    "rows_out": (batch or {}).get("rows_out"),
                    "quality": _parse_json((batch or {}).get("quality_json")),
                }
            )
        return envelope(items, self._official_meta("status", trade_date, source="datahub"))

    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,),
            )
        return envelope(rows, self._official_meta("batches", trade_date, source="datahub"))

    def _published_rows(
        self,
        dataset: str,
        table: str,
        q: dict[str, str],
        source: str,
        adjust: str = "none",
        default_code: str = "",
    ) -> dict[str, Any]:
        trade_date = q.get("date") or q.get("trade_date") or ""
        code = q.get("code") or default_code
        start = q.get("from") or ""
        end = q.get("to") or ""
        if trade_date:
            trade_date = yyyymmdd(trade_date)
            start = end = trade_date
        if not start or not end:
            if not trade_date:
                raise ApiError("INVALID_ARGUMENT", "date or from/to is required")
        else:
            start = yyyymmdd(start)
            end = yyyymmdd(end)
        ts_code = ""
        if code:
            resolved = resolve_code(self.db, code)
            if resolved is None:
                raise ApiError("INVALID_ARGUMENT", f"ambiguous code: {code}")
            ts_code = resolved
        # For a range, use per-date published batch. Single-date is the common path.
        if start == end:
            pub = self.db.fetchone(
                "SELECT * FROM publications WHERE dataset = ? AND trade_date = ?",
                (dataset, start),
            )
            if not pub:
                raise ApiError(
                    "DATASET_NOT_PUBLISHED",
                    f"{dataset} {start} 尚未发布",
                    extra={"expected_at": "15:05+08:00"},
                )
            limit, offset = self._page(q)
            sql = f"SELECT * FROM {table} WHERE trade_date = ? AND batch_id = ?"
            params: list[Any] = [start, pub["active_batch"]]
            if ts_code:
                sql += " AND ts_code = ?"
                params.append(ts_code)
            sql += " ORDER BY ts_code LIMIT ? OFFSET ?"
            params.extend([limit, offset])
            rows = [dict(row) for row in self.db.fetchall(sql, tuple(params))]
            if adjust == "qfq" and dataset == "daily":
                rows = self._apply_qfq(rows)
            meta = {
                "tier": "official",
                "trade_date": start,
                "published_at": pub["published_at"],
                "source": source,
                "batch_id": pub["active_batch"],
                "stale": False,
                "staleness_seconds": 0,
                "state": pub["state"],
            }
            return envelope(rows, meta)
        # multi-day: walk published dates
        pubs = self.db.fetchall(
            "SELECT * FROM publications WHERE dataset = ? AND trade_date >= ? AND trade_date <= ? ORDER BY trade_date",
            (dataset, start, end),
        )
        if not pubs:
            raise ApiError("DATASET_NOT_PUBLISHED", f"{dataset} {start}-{end} 尚未发布")
        rows: list[dict[str, Any]] = []
        limit, offset = self._page(q)
        for pub in pubs:
            sql = f"SELECT * FROM {table} WHERE trade_date = ? AND batch_id = ?"
            params = [pub["trade_date"], pub["active_batch"]]
            if ts_code:
                sql += " AND ts_code = ?"
                params.append(ts_code)
            sql += " ORDER BY ts_code"
            rows.extend(self.db.fetchall(sql, tuple(params)))
        sliced = rows[offset: offset + limit]
        if adjust == "qfq" and dataset == "daily":
            sliced = self._apply_qfq(sliced)
        last = pubs[-1]
        return envelope(
            sliced,
            {
                "tier": "official",
                "trade_date": last["trade_date"],
                "published_at": last["published_at"],
                "source": source,
                "batch_id": last["active_batch"],
                "stale": False,
                "staleness_seconds": 0,
            },
        )

    def _apply_qfq(self, rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
        by_code: dict[str, list[dict[str, Any]]] = {}
        for row in rows:
            by_code.setdefault(str(row["ts_code"]), []).append(row)
        out: list[dict[str, Any]] = []
        for code, group in by_code.items():
            latest = None
            factors = [finite_number(item.get("adj_factor")) for item in group]
            factors = [item for item in factors if item]
            if factors:
                latest = max(factors)
            else:
                extra = self.db.fetchone(
                    "SELECT MAX(adj_factor) AS f FROM eod_bars WHERE ts_code = ?",
                    (code,),
                )
                latest = finite_number((extra or {}).get("f"), 1.0)
            out.extend(qfq_bar(item, latest) for item in group)
        return out

    def _page(self, q: dict[str, str]) -> tuple[int, int]:
        try:
            limit = int(q.get("limit") or self.settings.list_limit_default)
            offset = int(q.get("offset") or 0)
        except ValueError as exc:
            raise ApiError("INVALID_ARGUMENT", "limit/offset must be integers") from exc
        limit = max(1, min(limit, self.settings.list_limit_max))
        offset = max(0, offset)
        return limit, offset

    def _official_meta(self, dataset: str, trade_date: str, source: str) -> dict[str, Any]:
        pub = self.db.fetchone(
            "SELECT * FROM publications WHERE dataset = ? AND trade_date = ?",
            (dataset, trade_date),
        )
        return {
            "tier": "official",
            "trade_date": trade_date,
            "published_at": (pub or {}).get("published_at"),
            "source": source,
            "batch_id": (pub or {}).get("active_batch"),
            "stale": False,
            "staleness_seconds": 0,
        }


def add_default(days: int) -> str:
    from datetime import timedelta

    return (now_shanghai() + timedelta(days=days)).strftime("%Y%m%d")


def parse_query(raw: str) -> dict[str, list[str]]:
    return parse_qs(raw, keep_blank_values=True)


def _parse_json(raw: Any) -> Any:
    if not raw:
        return None
    if isinstance(raw, dict):
        return raw
    import json

    try:
        return json.loads(str(raw))
    except json.JSONDecodeError:
        return None