xiaobai-review/xiaobai-datahub/datahub/realtime_serve.py

"""Provisional (盘中观察) serving: quotes, index quotes, intraday points.

Free sources only. Never writes official eod_* tables. Uses rt_cache + LKG.
"""

from __future__ import annotations

import json
import time
from datetime import datetime
from typing import Any

from datahub.adapters.eastmoney import EastmoneyAdapter
from datahub.adapters.tencent import TencentAdapter
from datahub.codes import resolve_code
from datahub.db import HubDB
from datahub.timeutil import isoformat, now_shanghai, yyyymmdd

QUOTE_TTL = 60
INDEX_TTL = 60
INTRADAY_TTL = 20


class RealtimeApiError(RuntimeError):
    def __init__(self, code: str, message: str) -> None:
        super().__init__(message)
        self.code = code
        self.message = message


def _envelope(data: Any, meta: dict[str, Any]) -> dict[str, Any]:
    from datahub import SCHEMA_VERSION

    return {"schema_version": SCHEMA_VERSION, "data": data, "meta": meta}


def fetch_index_quotes(db: HubDB) -> dict[str, Any]:
    cache_key = "indexes:quotes"
    cached = _read_cache(db, cache_key)
    if cached is not None:
        return cached
    eastmoney = EastmoneyAdapter()
    try:
        rows = eastmoney.fetch_indices()
        source = "eastmoney:ulist"
    except Exception:
        rows = TencentAdapter().fetch_indices()
        source = "tencent:qt"
    if len(rows) < 3:
        raise RealtimeApiError("SOURCE_UNAVAILABLE", "index quotes incomplete")
    payload = _envelope(
        rows,
        {
            "tier": "provisional",
            "trade_date": yyyymmdd(now_shanghai()),
            "source": source,
            "stale": False,
            "staleness_seconds": 0,
            "published_at": isoformat(now_shanghai()),
        },
    )
    _write_cache(db, cache_key, payload, INDEX_TTL, source)
    return payload


def fetch_quotes(db: HubDB, codes: list[str]) -> dict[str, Any]:
    if not codes:
        raise RealtimeApiError("INVALID_ARGUMENT", "codes is required")
    resolved: list[str] = []
    for code in codes[:60]:
        item = resolve_code(db, code) or _guess_ts_code(code)
        if item:
            resolved.append(item)
    if not resolved:
        raise RealtimeApiError("INVALID_ARGUMENT", "no resolvable codes")
    cache_key = "quotes:" + ",".join(sorted(resolved))
    cached = _read_cache(db, cache_key)
    if cached is not None:
        return cached
    adapter = EastmoneyAdapter()
    try:
        rows = adapter.fetch_quotes(resolved)
        source = "eastmoney:clist"
    except Exception as exc:
        raise RealtimeApiError("SOURCE_UNAVAILABLE", f"quotes unavailable: {exc}") from exc
    payload = _envelope(
        rows,
        {
            "tier": "provisional",
            "trade_date": yyyymmdd(now_shanghai()),
            "source": source,
            "stale": False,
            "staleness_seconds": 0,
            "published_at": isoformat(now_shanghai()),
        },
    )
    _write_cache(db, cache_key, payload, QUOTE_TTL, source)
    return payload


def fetch_intraday(db: HubDB, code: str, date: str = "") -> dict[str, Any]:
    ts_code = resolve_code(db, code) or _guess_ts_code(code)
    if not ts_code:
        raise RealtimeApiError("INVALID_ARGUMENT", f"ambiguous code: {code}")
    cache_key = f"intraday:{ts_code}:{date or 'today'}"
    cached = _read_cache(db, cache_key)
    if cached is not None:
        return cached
    adapter = EastmoneyAdapter()
    try:
        payload_data = adapter.fetch_intraday(ts_code)
        source = "eastmoney:trends2"
    except Exception as exc:
        raise RealtimeApiError("SOURCE_UNAVAILABLE", f"intraday unavailable: {exc}") from exc
    payload = _envelope(
        payload_data,
        {
            "tier": "provisional",
            "trade_date": yyyymmdd(payload_data.get("trade_date") or date or now_shanghai()),
            "source": source,
            "stale": False,
            "staleness_seconds": 0,
            "published_at": isoformat(now_shanghai()),
        },
    )
    _write_cache(db, cache_key, payload, INTRADAY_TTL, source)
    return payload


def _guess_ts_code(code: str) -> str | None:
    raw = str(code or "").strip().upper()
    if "." in raw:
        return raw
    if len(raw) == 6 and raw.isdigit():
        if raw.startswith(("5", "6", "9")):
            return f"{raw}.SH"
        return f"{raw}.SZ"
    return None


def _read_cache(db: HubDB, cache_key: str) -> dict[str, Any] | None:
    row = db.fetchone("SELECT * FROM rt_cache WHERE cache_key = ?", (cache_key,))
    if not row:
        return None
    expires = str(row.get("expires_at") or "")
    now = isoformat(now_shanghai())
    if expires and expires < now:
        return None
    try:
        payload = json.loads(row["payload"])
    except json.JSONDecodeError:
        return None
    if isinstance(payload, dict) and isinstance(payload.get("meta"), dict):
        stored = str(row.get("stored_at") or "")
        try:
            age = max(0, int(time.time() - datetime.fromisoformat(stored).timestamp()))
        except Exception:
            age = 0
        payload["meta"]["staleness_seconds"] = age
        payload["meta"]["stale"] = age > QUOTE_TTL
    return payload


def _write_cache(db: HubDB, cache_key: str, payload: dict[str, Any], ttl: int, source: str) -> None:
    from datetime import timedelta

    now = now_shanghai()
    stored = isoformat(now)
    expires = isoformat(now + timedelta(seconds=ttl))
    db.execute(
        """
        INSERT INTO rt_cache(cache_key, payload, source, stored_at, expires_at)
        VALUES (?,?,?,?,?)
        ON CONFLICT(cache_key) DO UPDATE SET
          payload=excluded.payload, source=excluded.source,
          stored_at=excluded.stored_at, expires_at=excluded.expires_at
        """,
        (cache_key, json.dumps(payload, ensure_ascii=False), source, stored, expires),
    )
    db.execute(
        """
        INSERT INTO last_known_good(cache_key, payload, source, stored_at)
        VALUES (?,?,?,?)
        ON CONFLICT(cache_key) DO UPDATE SET
          payload=excluded.payload, source=excluded.source, stored_at=excluded.stored_at
        """,
        (cache_key, json.dumps(payload, ensure_ascii=False), source, stored),
    )