from __future__ import annotations from typing import Any, Iterable from datahub.db import HubDB from datahub.timeutil import iter_yyyymmdd, yyyymmdd MISSING_SAMPLE_LIMIT = 10 def coverage_payload( *, kind: str, start: str, end: str, expected: Iterable[str], available: Iterable[str], extra: dict[str, Any] | None = None, ) -> dict[str, Any]: start = yyyymmdd(start) end = yyyymmdd(end) expected_list = sorted({yyyymmdd(item) for item in expected if item}) available_set = {yyyymmdd(item) for item in available if item} missing = [item for item in expected_list if item not in available_set] payload: dict[str, Any] = { "kind": kind, "complete": not missing, "requested_from": start, "requested_to": end, "available_from": min(available_set) if available_set else None, "available_to": max(available_set) if available_set else None, "expected_count": len(expected_list), "available_count": len(available_set), "missing_count": len(missing), "missing_sample": missing[:MISSING_SAMPLE_LIMIT], } if extra: payload.update(extra) return payload def calendar_coverage(db: HubDB, start: str, end: str, exchange: str = "SSE") -> dict[str, Any]: start = yyyymmdd(start) end = yyyymmdd(end) expected = list(iter_yyyymmdd(start, end)) rows = db.fetchall( "SELECT cal_date FROM trade_calendar WHERE exchange = ? AND cal_date >= ? AND cal_date <= ?", (exchange, start, end), ) return coverage_payload( kind="calendar", start=start, end=end, expected=expected, available=(row["cal_date"] for row in rows), extra={"exchange": exchange}, ) def published_range_coverage( db: HubDB, dataset: str, start: str, end: str, ts_code: str = "", table: str = "", ) -> dict[str, Any]: start = yyyymmdd(start) end = yyyymmdd(end) calendar = calendar_coverage(db, start, end) open_rows = db.fetchall( """ SELECT cal_date FROM trade_calendar WHERE exchange = 'SSE' AND is_open = 1 AND cal_date >= ? AND cal_date <= ? ORDER BY cal_date """, (start, end), ) expected_open = [row["cal_date"] for row in open_rows] pubs = db.fetchall( """ SELECT trade_date, active_batch FROM publications WHERE dataset = ? AND trade_date >= ? AND trade_date <= ? ORDER BY trade_date """, (dataset, start, end), ) published_dates = [row["trade_date"] for row in pubs] available = list(published_dates) extra: dict[str, Any] = { "dataset": dataset, "calendar_complete": calendar["complete"], "calendar_missing_count": calendar["missing_count"], } if ts_code and table and pubs: present_code: list[str] = [] for pub in pubs: hit = db.fetchone( f"SELECT 1 AS ok FROM {table} WHERE trade_date = ? AND batch_id = ? AND ts_code = ? LIMIT 1", (pub["trade_date"], pub["active_batch"], ts_code), ) if hit: present_code.append(pub["trade_date"]) available = present_code extra["code"] = ts_code payload = coverage_payload( kind="published_range", start=start, end=end, expected=expected_open, available=available, extra=extra, ) if not calendar["complete"]: payload["complete"] = False payload["calendar_missing_sample"] = calendar["missing_sample"] return payload def point_coverage(trade_date: str, dataset: str = "") -> dict[str, Any]: day = yyyymmdd(trade_date) payload = coverage_payload( kind="point", start=day, end=day, expected=[day], available=[day], extra={"dataset": dataset} if dataset else None, ) return payload