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 hashlib
import json
import time
from datetime import datetime
from typing import Any

from datahub.adapters.base import AdapterError
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.governance.lkg import LastKnownGood
from datahub.timeutil import isoformat, now_shanghai, yyyymmdd

QUOTE_TTL = 60
INDEX_TTL = 60
INTRADAY_TTL = 20
QUOTE_BATCH = 60


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_market_quotes(db: HubDB) -> dict[str, Any]:
    cache_key = "quotes:market"
    cached = _read_cache(db, cache_key)
    if cached is not None:
        return cached
    errors: list[str] = []
    rows: list[dict[str, Any]] = []
    source = ""
    try:
        rows = EastmoneyAdapter().fetch_market_quotes()
        source = "eastmoney:clist"
    except Exception as exc:
        errors.append(f"eastmoney:{exc}")
        try:
            listed = _listed_ts_codes(db)
            if not listed:
                raise AdapterError("no local stock master for tencent market snapshot")
            rows = TencentAdapter().fetch_quotes(listed)
            if len(rows) < 200:
                raise AdapterError(f"Tencent market snapshot too small: {len(rows)}")
            source = "tencent:qt"
        except Exception as backup_exc:
            errors.append(f"tencent:{backup_exc}")
            recovered = _load_quotes_lkg(db, cache_key)
            if recovered is not None:
                return recovered
            raise RealtimeApiError(
                "SOURCE_UNAVAILABLE",
                "market quotes unavailable: " + "；".join(errors),
            ) from backup_exc
    payload = _quote_payload(rows, source, scope="market")
    _write_cache(db, cache_key, payload, QUOTE_TTL, source)
    return payload


def fetch_quotes(db: HubDB, codes: list[str]) -> dict[str, Any]:
    if not codes:
        return fetch_market_quotes(db)
    resolved: list[str] = []
    seen: set[str] = set()
    for code in codes:
        item = resolve_code(db, code) or _guess_ts_code(code)
        if item and item not in seen:
            seen.add(item)
            resolved.append(item)
    if not resolved:
        raise RealtimeApiError("INVALID_ARGUMENT", "no resolvable codes")
    digest = hashlib.sha1(",".join(sorted(resolved)).encode("utf-8")).hexdigest()
    cache_key = f"quotes:{digest}:{len(resolved)}"
    cached = _read_cache(db, cache_key)
    if cached is not None:
        return cached
    minimum = max(1, int(len(resolved) * 0.5))
    errors: list[str] = []
    rows: list[dict[str, Any]] = []
    source = ""
    try:
        rows = _eastmoney_named_quotes(resolved)
        if len(rows) < minimum:
            raise AdapterError(f"Eastmoney named quotes too small: {len(rows)}/{len(resolved)}")
        source = "eastmoney:ulist"
    except Exception as exc:
        errors.append(f"eastmoney:{exc}")
        try:
            rows = TencentAdapter().fetch_quotes(resolved)
            if len(rows) < minimum:
                raise AdapterError(f"Tencent named quotes too small: {len(rows)}/{len(resolved)}")
            source = "tencent:qt"
        except Exception as backup_exc:
            errors.append(f"tencent:{backup_exc}")
            recovered = _load_quotes_lkg(db, cache_key)
            if recovered is not None:
                return recovered
            raise RealtimeApiError(
                "SOURCE_UNAVAILABLE",
                "quotes unavailable: " + "；".join(errors),
            ) from backup_exc
    payload = _quote_payload(rows, source)
    _write_cache(db, cache_key, payload, QUOTE_TTL, source)
    return payload


def fetch_sector_quote(db: HubDB, code: str, expected_date: str = "") -> dict[str, Any]:
    ts_code = str(code or "").strip().upper()
    if ts_code.isdigit():
        ts_code = f"{ts_code}.SI"
    cache_key = f"sector:{ts_code}"
    cached = _read_cache(db, cache_key)
    if cached is not None:
        return cached
    errors: list[str] = []
    try:
        row = EastmoneyAdapter().fetch_shenwan_quote(ts_code)
        source = str(row.get("source") or "eastmoney_sw")
    except Exception as exc:
        errors.append(f"eastmoney:{exc}")
        recovered = _load_quotes_lkg(db, cache_key)
        if recovered is not None:
            return recovered
        raise RealtimeApiError(
            "SOURCE_UNAVAILABLE",
            "sector quote unavailable: " + "；".join(errors),
        ) from exc
    want = str(expected_date or "").replace("-", "")[:8]
    quote_date = str(row.get("quote_date") or "")
    if want and quote_date and quote_date != want:
        recovered = _load_quotes_lkg(db, cache_key)
        if recovered is not None:
            return recovered
        raise RealtimeApiError("SOURCE_UNAVAILABLE", f"sector quote date {quote_date} != {want}")
    payload = _envelope(
        row,
        {
            "tier": "provisional",
            "trade_date": quote_date or 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_limit_pool(db: HubDB, trade_date: str = "") -> dict[str, Any]:
    day = yyyymmdd(trade_date or now_shanghai())
    cache_key = f"limit-pool:{day}"
    cached = _read_cache(db, cache_key)
    if cached is not None:
        return cached
    try:
        rows = EastmoneyAdapter().fetch_limit_pool(day)
        source = "eastmoney:zt_pool"
    except Exception as exc:
        recovered = _load_quotes_lkg(db, cache_key)
        if recovered is not None:
            return recovered
        raise RealtimeApiError("SOURCE_UNAVAILABLE", f"limit pool unavailable: {exc}") from exc
    payload = _envelope(
        rows,
        {
            "tier": "provisional",
            "trade_date": day,
            "source": source,
            "stale": False,
            "staleness_seconds": 0,
            "published_at": isoformat(now_shanghai()),
        },
    )
    _write_cache(db, cache_key, payload, QUOTE_TTL, source)
    return payload


def _eastmoney_named_quotes(codes: list[str]) -> list[dict[str, Any]]:
    adapter = EastmoneyAdapter()
    rows: list[dict[str, Any]] = []
    for index in range(0, len(codes), QUOTE_BATCH):
        rows.extend(adapter.fetch_quotes(codes[index:index + QUOTE_BATCH]))
    return rows


def _listed_ts_codes(db: HubDB) -> list[str]:
    try:
        rows = db.fetchall(
            "SELECT ts_code FROM stock_master WHERE list_status = 'L' ORDER BY ts_code"
        )
    except Exception:
        return []
    return [str(row.get("ts_code") or "") for row in rows if row.get("ts_code")]


def _quote_payload(
    rows: list[dict[str, Any]],
    source: str,
    scope: str = "",
) -> dict[str, Any]:
    meta: dict[str, Any] = {
        "tier": "provisional",
        "trade_date": yyyymmdd(now_shanghai()),
        "source": source,
        "stale": False,
        "staleness_seconds": 0,
        "published_at": isoformat(now_shanghai()),
        "failover": source.startswith("tencent"),
        "delay_notice": "",
    }
    if scope:
        meta["scope"] = scope
    return _envelope(rows, meta)


def _load_quotes_lkg(db: HubDB, cache_key: str) -> dict[str, Any] | None:
    store = LastKnownGood(db)
    item = store.load(cache_key)
    payload = item.get("payload") if item else None
    if not isinstance(payload, dict):
        return None
    data = payload.get("data")
    if not isinstance(data, list) or not data:
        return None
    stamped = dict(payload)
    meta = dict(stamped.get("meta") or {})
    stored = str((item or {}).get("stored_at") or "")
    try:
        age = max(0, int(time.time() - datetime.fromisoformat(stored).timestamp()))
    except Exception:
        age = QUOTE_TTL
    meta["stale"] = True
    meta["staleness_seconds"] = age
    meta["delay_notice"] = f"主备免费行情均暂不可用，显示 {age} 秒前的真实快照"
    meta["lkg_source"] = str((item or {}).get("source") or meta.get("source") or "")
    stamped["meta"] = meta
    return stamped


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, date)
        source = "eastmoney:trends2"
    except Exception as exc:
        recovered = _load_intraday_lkg(db, ts_code, date)
        if recovered is None:
            raise RealtimeApiError("SOURCE_UNAVAILABLE", f"intraday unavailable: {exc}") from exc
        return recovered
    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 _load_intraday_lkg(db: HubDB, ts_code: str, date: str = "") -> dict[str, Any] | None:
    store = LastKnownGood(db)
    keys = [f"intraday:{ts_code}:{date or 'today'}"]
    if date:
        keys.append(f"intraday:{ts_code}:today")
    for key in keys:
        item = store.load(key)
        payload = _lkg_payload(item)
        if payload is not None:
            return payload
    row = db.fetchone(
        "SELECT * FROM last_known_good WHERE cache_key LIKE ? ORDER BY stored_at DESC LIMIT 1",
        (f"intraday:{ts_code}:%",),
    )
    if not row:
        return None
    try:
        raw = json.loads(row["payload"])
    except json.JSONDecodeError:
        return None
    return _mark_stale(raw) if isinstance(raw, dict) else None


def _lkg_payload(item: dict[str, Any] | None) -> dict[str, Any] | None:
    if not item:
        return None
    payload = item.get("payload")
    return _mark_stale(payload) if isinstance(payload, dict) else None


def _mark_stale(payload: dict[str, Any]) -> dict[str, Any] | None:
    data = payload.get("data")
    if not isinstance(data, dict) or not data.get("points"):
        return None
    stamped = dict(payload)
    meta = dict(stamped.get("meta") or {})
    meta["stale"] = True
    stamped["meta"] = meta
    return stamped


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