from __future__ import annotations import sqlite3 import threading from collections.abc import Iterator from contextlib import contextmanager from pathlib import Path from typing import Any from datahub.timeutil import isoformat SCHEMA = """ CREATE TABLE IF NOT EXISTS schema_migrations ( version INTEGER PRIMARY KEY, applied_at TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS credentials ( name TEXT PRIMARY KEY, encrypted_payload TEXT NOT NULL, last4 TEXT, updated_at TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS hub_admin ( id INTEGER PRIMARY KEY, username TEXT NOT NULL UNIQUE, password_salt TEXT NOT NULL, password_hash TEXT NOT NULL, password_must_change INTEGER NOT NULL DEFAULT 1, failed_attempts INTEGER NOT NULL DEFAULT 0, locked_until TEXT, created_at TEXT NOT NULL, updated_at TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS hub_sessions ( token_hash TEXT PRIMARY KEY, csrf_token TEXT NOT NULL, expires_at TEXT NOT NULL, created_at TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS api_tokens ( token_hash TEXT PRIMARY KEY, name TEXT NOT NULL, last4 TEXT NOT NULL, created_at TEXT NOT NULL, revoked_at TEXT ); CREATE TABLE IF NOT EXISTS trade_calendar ( exchange TEXT NOT NULL, cal_date TEXT NOT NULL, is_open INTEGER NOT NULL, pretrade_date TEXT, fetched_at TEXT NOT NULL, PRIMARY KEY (exchange, cal_date) ); CREATE TABLE IF NOT EXISTS stock_master ( ts_code TEXT PRIMARY KEY, symbol TEXT, name TEXT, area TEXT, industry TEXT, market TEXT, list_status TEXT, list_date TEXT, updated_at TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS eod_bars ( ts_code TEXT NOT NULL, trade_date TEXT NOT NULL, open REAL, high REAL, low REAL, close REAL, pct_chg REAL, volume REAL, amount REAL, adj_factor REAL, batch_id TEXT NOT NULL, PRIMARY KEY (ts_code, trade_date, batch_id) ) WITHOUT ROWID; CREATE TABLE IF NOT EXISTS eod_valuation ( ts_code TEXT NOT NULL, trade_date TEXT NOT NULL, turnover_rate REAL, volume_ratio REAL, total_mv REAL, circ_mv REAL, pe_ttm REAL, pb REAL, ps_ttm REAL, dv_ttm REAL, batch_id TEXT NOT NULL, PRIMARY KEY (ts_code, trade_date, batch_id) ) WITHOUT ROWID; CREATE TABLE IF NOT EXISTS eod_moneyflow ( ts_code TEXT NOT NULL, trade_date TEXT NOT NULL, buy_sm_amount REAL, sell_sm_amount REAL, buy_md_amount REAL, sell_md_amount REAL, buy_lg_amount REAL, sell_lg_amount REAL, buy_elg_amount REAL, sell_elg_amount REAL, net_mf_amount REAL, batch_id TEXT NOT NULL, PRIMARY KEY (ts_code, trade_date, batch_id) ) WITHOUT ROWID; CREATE TABLE IF NOT EXISTS eod_auction ( ts_code TEXT NOT NULL, trade_date TEXT NOT NULL, volume REAL, price REAL, amount REAL, pre_close REAL, turnover_rate REAL, volume_ratio REAL, float_share REAL, batch_id TEXT NOT NULL, PRIMARY KEY (ts_code, trade_date, batch_id) ) WITHOUT ROWID; CREATE TABLE IF NOT EXISTS eod_index_bars ( ts_code TEXT NOT NULL, trade_date TEXT NOT NULL, open REAL, high REAL, low REAL, close REAL, pct_chg REAL, volume REAL, amount REAL, batch_id TEXT NOT NULL, PRIMARY KEY (ts_code, trade_date, batch_id) ) WITHOUT ROWID; CREATE TABLE IF NOT EXISTS staging_bars ( ts_code TEXT NOT NULL, trade_date TEXT NOT NULL, batch_id TEXT NOT NULL, open REAL, high REAL, low REAL, close REAL, pct_chg REAL, volume REAL, amount REAL, adj_factor REAL, PRIMARY KEY (batch_id, ts_code, trade_date) ); CREATE TABLE IF NOT EXISTS staging_valuation ( ts_code TEXT NOT NULL, trade_date TEXT NOT NULL, batch_id TEXT NOT NULL, turnover_rate REAL, volume_ratio REAL, total_mv REAL, circ_mv REAL, pe_ttm REAL, pb REAL, ps_ttm REAL, dv_ttm REAL, PRIMARY KEY (batch_id, ts_code, trade_date) ); CREATE TABLE IF NOT EXISTS staging_moneyflow ( ts_code TEXT NOT NULL, trade_date TEXT NOT NULL, batch_id TEXT NOT NULL, buy_sm_amount REAL, sell_sm_amount REAL, buy_md_amount REAL, sell_md_amount REAL, buy_lg_amount REAL, sell_lg_amount REAL, buy_elg_amount REAL, sell_elg_amount REAL, net_mf_amount REAL, PRIMARY KEY (batch_id, ts_code, trade_date) ); CREATE TABLE IF NOT EXISTS staging_auction ( ts_code TEXT NOT NULL, trade_date TEXT NOT NULL, batch_id TEXT NOT NULL, volume REAL, price REAL, amount REAL, pre_close REAL, turnover_rate REAL, volume_ratio REAL, float_share REAL, PRIMARY KEY (batch_id, ts_code, trade_date) ); CREATE TABLE IF NOT EXISTS staging_index_bars ( ts_code TEXT NOT NULL, trade_date TEXT NOT NULL, batch_id TEXT NOT NULL, open REAL, high REAL, low REAL, close REAL, pct_chg REAL, volume REAL, amount REAL, PRIMARY KEY (batch_id, ts_code, trade_date) ); CREATE TABLE IF NOT EXISTS publications ( dataset TEXT NOT NULL, trade_date TEXT NOT NULL, active_batch TEXT NOT NULL, prev_batch TEXT, state TEXT NOT NULL, published_at TEXT NOT NULL, PRIMARY KEY (dataset, trade_date) ); CREATE TABLE IF NOT EXISTS publication_history ( dataset TEXT NOT NULL, trade_date TEXT NOT NULL, batch_id TEXT NOT NULL, published_at TEXT NOT NULL, generation INTEGER NOT NULL, PRIMARY KEY (dataset, trade_date, batch_id) ); CREATE TABLE IF NOT EXISTS batches ( batch_id TEXT PRIMARY KEY, dataset TEXT NOT NULL, trade_date TEXT NOT NULL, state TEXT NOT NULL, attempt INTEGER DEFAULT 0, rows_in INTEGER, rows_out INTEGER, quality_json TEXT, started_at TEXT, finished_at TEXT, error TEXT ); CREATE TABLE IF NOT EXISTS src_health ( provider TEXT NOT NULL, endpoint_class TEXT NOT NULL, state TEXT NOT NULL, last_ok_at TEXT, last_error TEXT, consec_failures INTEGER DEFAULT 0, opened_at TEXT, cooldown_until TEXT, PRIMARY KEY (provider, endpoint_class) ); CREATE TABLE IF NOT EXISTS src_calls ( id INTEGER PRIMARY KEY AUTOINCREMENT, provider TEXT NOT NULL, endpoint TEXT NOT NULL, ok INTEGER NOT NULL, latency_ms INTEGER, error TEXT, created_at TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS job_runs ( id INTEGER PRIMARY KEY AUTOINCREMENT, job_id TEXT NOT NULL, state TEXT NOT NULL, started_at TEXT, finished_at TEXT, rows_in INTEGER, rows_out INTEGER, error TEXT, attempt INTEGER DEFAULT 1, detail TEXT ); CREATE TABLE IF NOT EXISTS audit_log ( id INTEGER PRIMARY KEY AUTOINCREMENT, actor TEXT NOT NULL, action TEXT NOT NULL, target TEXT, detail TEXT, created_at TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS rt_cache ( cache_key TEXT PRIMARY KEY, payload TEXT NOT NULL, source TEXT NOT NULL, stored_at TEXT NOT NULL, expires_at TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS last_known_good ( cache_key TEXT PRIMARY KEY, payload TEXT NOT NULL, source TEXT NOT NULL, stored_at TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS diff_reports ( id INTEGER PRIMARY KEY AUTOINCREMENT, trade_date TEXT NOT NULL, metric TEXT NOT NULL, left_source TEXT, right_source TEXT, left_value REAL, right_value REAL, deviation REAL, sample_count INTEGER, created_at TEXT NOT NULL ); CREATE INDEX IF NOT EXISTS idx_batches_date ON batches(trade_date, dataset); CREATE INDEX IF NOT EXISTS idx_job_runs_job ON job_runs(job_id, started_at); CREATE INDEX IF NOT EXISTS idx_src_calls_created ON src_calls(created_at); CREATE INDEX IF NOT EXISTS idx_eod_bars_date ON eod_bars(trade_date, batch_id); CREATE INDEX IF NOT EXISTS idx_calendar_open ON trade_calendar(is_open, cal_date); """ DATASET_TABLES = { "daily": ("eod_bars", "staging_bars"), "valuation": ("eod_valuation", "staging_valuation"), "moneyflow": ("eod_moneyflow", "staging_moneyflow"), "auction": ("eod_auction", "staging_auction"), "index_daily": ("eod_index_bars", "staging_index_bars"), } class ManagedConnection(sqlite3.Connection): def __exit__(self, exc_type, exc_value, traceback): try: return super().__exit__(exc_type, exc_value, traceback) finally: self.close() class HubDB: def __init__(self, path: Path, timeout_seconds: float = 20) -> None: self.path = Path(path) self.path.parent.mkdir(parents=True, exist_ok=True) self.timeout_seconds = timeout_seconds self._write_lock = threading.RLock() self.initialize() def connect(self) -> sqlite3.Connection: connection = sqlite3.connect( self.path, timeout=self.timeout_seconds, factory=ManagedConnection, ) connection.row_factory = sqlite3.Row connection.execute("PRAGMA journal_mode=WAL") connection.execute("PRAGMA foreign_keys=ON") connection.execute("PRAGMA busy_timeout=20000") connection.execute("PRAGMA synchronous=NORMAL") return connection def initialize(self) -> None: with self.connect() as connection: connection.executescript(SCHEMA) row = connection.execute( "SELECT version FROM schema_migrations ORDER BY version DESC LIMIT 1" ).fetchone() if row is None: connection.execute( "INSERT INTO schema_migrations(version, applied_at) VALUES (1, ?)", (isoformat(),), ) @contextmanager def write(self) -> Iterator[sqlite3.Connection]: with self._write_lock: with self.connect() as connection: yield connection def fetchall(self, sql: str, params: tuple[Any, ...] = ()) -> list[dict[str, Any]]: with self.connect() as connection: rows = connection.execute(sql, params).fetchall() return [dict(row) for row in rows] def fetchone(self, sql: str, params: tuple[Any, ...] = ()) -> dict[str, Any] | None: with self.connect() as connection: row = connection.execute(sql, params).fetchone() return dict(row) if row else None def execute(self, sql: str, params: tuple[Any, ...] = ()) -> None: with self.write() as connection: connection.execute(sql, params) def executemany(self, sql: str, rows: list[tuple[Any, ...]]) -> None: with self.write() as connection: connection.executemany(sql, rows) def backup_to(self, dest: Path) -> None: dest.parent.mkdir(parents=True, exist_ok=True) with self.connect() as source, sqlite3.connect(dest) as target: source.backup(target) def vacuum(self) -> None: with self.connect() as connection: connection.execute("VACUUM")