xiaobai-review/xiaobai-datahub/datahub/db.py

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