llm-quant/app/data/schema.py

"""Database schema management for the investment assistant."""
from __future__ import annotations

import sqlite3
from dataclasses import dataclass
from typing import Iterable

from app.utils.db import db_session


SCHEMA_STATEMENTS: Iterable[str] = (
    """
    CREATE TABLE IF NOT EXISTS stock_basic (
      ts_code TEXT PRIMARY KEY,
      symbol TEXT,
      name TEXT,
      area TEXT,
      industry TEXT,
      market TEXT,
      exchange TEXT,
      list_status TEXT,
      list_date TEXT,
      delist_date TEXT
    );
    """,
    """
    CREATE TABLE IF NOT EXISTS daily (
      ts_code TEXT,
      trade_date TEXT,
      open REAL,
      high REAL,
      low REAL,
      close REAL,
      pre_close REAL,
      change REAL,
      pct_chg REAL,
      vol REAL,
      amount REAL,
      PRIMARY KEY (ts_code, trade_date)
    );
    """,
    """
    CREATE TABLE IF NOT EXISTS daily_basic (
      ts_code TEXT,
      trade_date TEXT,
      close REAL,
      turnover_rate REAL,
      turnover_rate_f REAL,
      volume_ratio REAL,
      pe REAL,
      pe_ttm REAL,
      pb REAL,
      ps REAL,
      ps_ttm REAL,
      dv_ratio REAL,
      dv_ttm REAL,
      total_share REAL,
      float_share REAL,
      free_share REAL,
      total_mv REAL,
      circ_mv REAL,
      PRIMARY KEY (ts_code, trade_date)
    );
    """,
    """
    CREATE TABLE IF NOT EXISTS adj_factor (
      ts_code TEXT,
      trade_date TEXT,
      adj_factor REAL,
      PRIMARY KEY (ts_code, trade_date)
    );
    """,
    """
    CREATE TABLE IF NOT EXISTS suspend (
      ts_code TEXT,
      suspend_date TEXT,
      resume_date TEXT,
      suspend_type TEXT,
      ann_date TEXT,
      suspend_timing TEXT,
      resume_timing TEXT,
      reason TEXT,
      PRIMARY KEY (ts_code, suspend_date)
    );
    """,
    """
    CREATE TABLE IF NOT EXISTS trade_calendar (
      exchange TEXT,
      cal_date TEXT PRIMARY KEY,
      is_open INTEGER,
      pretrade_date TEXT
    );
    """,
    """
    CREATE TABLE IF NOT EXISTS stk_limit (
      ts_code TEXT,
      trade_date TEXT,
      up_limit REAL,
      down_limit REAL,
      PRIMARY KEY (ts_code, trade_date)
    );
    """,
    """
    CREATE TABLE IF NOT EXISTS news (
      id TEXT PRIMARY KEY,
      ts_code TEXT,
      pub_time TEXT,
      source TEXT,
      title TEXT,
      summary TEXT,
      url TEXT,
      entities TEXT,
      sentiment REAL,
      heat REAL
    );
    """,
    """
    CREATE INDEX IF NOT EXISTS idx_news_time ON news(pub_time DESC);
    """,
    """
    CREATE INDEX IF NOT EXISTS idx_news_code ON news(ts_code, pub_time DESC);
    """,
    """
    CREATE TABLE IF NOT EXISTS heat_daily (
      scope TEXT,
      key TEXT,
      trade_date TEXT,
      heat REAL,
      top_topics TEXT,
      PRIMARY KEY (scope, key, trade_date)
    );
    """,
    """
    CREATE TABLE IF NOT EXISTS bt_config (
      id TEXT PRIMARY KEY,
      name TEXT,
      start_date TEXT,
      end_date TEXT,
      universe TEXT,
      params TEXT
    );
    """,
    """
    CREATE TABLE IF NOT EXISTS bt_trades (
      cfg_id TEXT,
      ts_code TEXT,
      trade_date TEXT,
      side TEXT,
      price REAL,
      qty REAL,
      reason TEXT,
      PRIMARY KEY (cfg_id, ts_code, trade_date, side)
    );
    """,
    """
    CREATE TABLE IF NOT EXISTS bt_nav (
      cfg_id TEXT,
      trade_date TEXT,
      nav REAL,
      ret REAL,
      pos_count INTEGER,
      turnover REAL,
      dd REAL,
      info TEXT,
      PRIMARY KEY (cfg_id, trade_date)
    );
    """,
    """
    CREATE TABLE IF NOT EXISTS bt_report (
      cfg_id TEXT PRIMARY KEY,
      summary TEXT
    );
    """,
    """
    CREATE TABLE IF NOT EXISTS run_log (
      ts TEXT PRIMARY KEY,
      stage TEXT,
      level TEXT,
      msg TEXT
    );
    """,
    """
    CREATE TABLE IF NOT EXISTS agent_utils (
      trade_date TEXT,
      ts_code TEXT,
      agent TEXT,
      action TEXT,
      utils TEXT,
      feasible TEXT,
      weight REAL,
      PRIMARY KEY (trade_date, ts_code, agent)
    );
    """,
    """
    CREATE TABLE IF NOT EXISTS alloc_log (
      trade_date TEXT,
      ts_code TEXT,
      target_weight REAL,
      clipped_weight REAL,
      reason TEXT,
      PRIMARY KEY (trade_date, ts_code)
    );
    """
)


@dataclass
class MigrationResult:
    executed: int
    skipped: bool = False


def _schema_exists() -> bool:
    try:
        with db_session(read_only=True) as conn:
            cursor = conn.execute(
                "SELECT name FROM sqlite_master WHERE type='table' AND name='news'"
            )
            return cursor.fetchone() is not None
    except sqlite3.OperationalError:
        return False


def initialize_database() -> MigrationResult:
    """Create tables and indexes required by the application."""

    if _schema_exists():
        return MigrationResult(executed=0, skipped=True)

    executed = 0
    with db_session() as conn:
        cursor = conn.cursor()
        for statement in SCHEMA_STATEMENTS:
            cursor.executescript(statement)
            executed += 1
    return MigrationResult(executed=executed)