storage/metrics_db.py

import json
import sqlite3
from pathlib import Path
from typing import Optional

DB_PATH = Path("data/metrics.db")


_VALUATION_COLS = [
    ("shares_diluted",     "REAL"),
    ("effective_tax_rate", "REAL"),
    ("interest_expense",   "REAL"),
    ("total_debt",         "REAL"),
    ("dividends_paid",     "REAL"),
    ("buybacks",           "REAL"),
    ("capex",              "REAL"),
    ("stockholders_equity","REAL"),
    ("guidance_period",                "TEXT"),
    ("guidance_revenue_low",           "REAL"),
    ("guidance_revenue_high",          "REAL"),
    ("guidance_eps_low",               "REAL"),
    ("guidance_eps_high",              "REAL"),
    ("guidance_gross_margin_low",      "REAL"),
    ("guidance_gross_margin_high",     "REAL"),
    ("guidance_operating_margin_low",  "REAL"),
    ("guidance_operating_margin_high", "REAL"),
    ("period_basis",         "TEXT"),
    ("report_date",          "TEXT"),
    ("accession",            "TEXT"),
    ("source_url",           "TEXT"),
    ("metric_contexts",      "TEXT"),
    ("quality_warnings",     "TEXT"),
    ("data_quality_status",  "TEXT NOT NULL DEFAULT 'LEGACY_UNVERIFIED'"),
]

_JSON_COLS = frozenset({"metric_contexts", "quality_warnings"})

_GUIDANCE_COLS = [
    "guidance_period",
    "guidance_revenue_low", "guidance_revenue_high",
    "guidance_eps_low", "guidance_eps_high",
    "guidance_gross_margin_low", "guidance_gross_margin_high",
    "guidance_operating_margin_low", "guidance_operating_margin_high",
]


def _create_metrics_table(conn: sqlite3.Connection) -> None:
    conn.execute("""
        CREATE TABLE IF NOT EXISTS metrics (
            ticker           TEXT NOT NULL,
            period           TEXT NOT NULL,
            company_name     TEXT,
            filing_date      TEXT,
            form_type        TEXT,
            revenue          REAL,
            revenue_yoy_pct  REAL,
            eps              REAL,
            gross_margin     REAL,
            operating_margin REAL,
            free_cash_flow   REAL,
            guidance_disclosed INTEGER,
            guidance_text    TEXT,
            ingested_at      TEXT,
            PRIMARY KEY (ticker, period)
        )
    """)


def _ensure_valuation_columns(conn: sqlite3.Connection) -> None:
    """Add optional columns without rebuilding or deleting the table."""
    existing = {row[1] for row in conn.execute("PRAGMA table_info(metrics)").fetchall()}
    for col, typ in _VALUATION_COLS:
        if col not in existing:
            conn.execute(f'ALTER TABLE metrics ADD COLUMN "{col}" {typ}')


def _next_legacy_table_name(conn: sqlite3.Connection) -> str:
    """Return an unused, stable archive name for a pre-composite schema."""
    suffix = 1
    while True:
        candidate = f"metrics_legacy_v{suffix}"
        exists = conn.execute(
            "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?",
            (candidate,),
        ).fetchone()
        if not exists:
            return candidate
        suffix += 1


def _migrate_ticker_primary_key(conn: sqlite3.Connection, columns: list[sqlite3.Row]) -> None:
    """Copy a ticker-only schema to the composite schema, retaining its archive.

    This function is called inside the transaction opened by ``init_db``.  The
    original table is deliberately retained rather than dropped, so unknown
    legacy columns remain recoverable as well as all copied metric rows.
    """
    legacy_columns = {row[1] for row in columns}
    if "ticker" not in legacy_columns:
        raise RuntimeError("Cannot migrate metrics table without a ticker column")

    archive = _next_legacy_table_name(conn)
    conn.execute(f'ALTER TABLE metrics RENAME TO "{archive}"')
    _create_metrics_table(conn)
    _ensure_valuation_columns(conn)

    target_columns = [
        row[1] for row in conn.execute("PRAGMA table_info(metrics)").fetchall()
    ]
    copied_columns: list[str] = []
    select_expressions: list[str] = []
    for column in target_columns:
        if column == "period":
            copied_columns.append(column)
            if column in legacy_columns:
                select_expressions.append(
                    "COALESCE(NULLIF(TRIM(CAST(\"period\" AS TEXT)), ''), 'LEGACY')"
                )
            else:
                select_expressions.append("'LEGACY'")
        elif column in legacy_columns:
            copied_columns.append(column)
            if column == "data_quality_status":
                select_expressions.append(
                    "COALESCE(\"data_quality_status\", 'LEGACY_UNVERIFIED')"
                )
            else:
                select_expressions.append(f'"{column}"')

    quoted_columns = ", ".join(f'"{column}"' for column in copied_columns)
    selected_values = ", ".join(select_expressions)
    conn.execute(
        f'INSERT INTO metrics ({quoted_columns}) '
        f'SELECT {selected_values} FROM "{archive}"'
    )


def init_db() -> None:
    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
    with sqlite3.connect(DB_PATH) as conn:
        # Keep the schema transition atomic: a failed copy rolls back the rename
        # and leaves the legacy database untouched.
        conn.execute("BEGIN IMMEDIATE")
        rows = conn.execute("PRAGMA table_info(metrics)").fetchall()
        if rows:
            pk_cols = {row[1] for row in rows if row[5] > 0}
            if "period" not in pk_cols:
                _migrate_ticker_primary_key(conn, rows)

        _create_metrics_table(conn)
        _ensure_valuation_columns(conn)


_UPSERT_PARAMS = [
    "ticker", "period", "company_name", "filing_date", "form_type",
    "revenue", "revenue_yoy_pct", "eps", "gross_margin",
    "operating_margin", "free_cash_flow",
    "guidance_disclosed", "guidance_text", "ingested_at",
] + [col for col, _ in _VALUATION_COLS]


def upsert_metrics(data: dict) -> None:
    # Partial rows are valid — absent optional columns bind as NULL.
    data = {key: data.get(key) for key in _UPSERT_PARAMS}
    data["data_quality_status"] = data.get("data_quality_status") or "LEGACY_UNVERIFIED"
    for key in _JSON_COLS:
        value = data.get(key)
        if value is not None and not isinstance(value, str):
            data[key] = json.dumps(value, separators=(",", ":"), sort_keys=True)
    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
    with sqlite3.connect(DB_PATH) as conn:
        conn.execute("""
            INSERT INTO metrics (
                ticker, period, company_name, filing_date, form_type,
                revenue, revenue_yoy_pct, eps, gross_margin,
                operating_margin, free_cash_flow,
                guidance_disclosed, guidance_text, ingested_at,
                shares_diluted, effective_tax_rate, interest_expense, total_debt,
                dividends_paid, buybacks, capex, stockholders_equity,
                guidance_period,
                guidance_revenue_low, guidance_revenue_high,
                guidance_eps_low, guidance_eps_high,
                guidance_gross_margin_low, guidance_gross_margin_high,
                guidance_operating_margin_low, guidance_operating_margin_high,
                period_basis, report_date, accession, source_url,
                metric_contexts, quality_warnings, data_quality_status
            ) VALUES (
                :ticker, :period, :company_name, :filing_date, :form_type,
                :revenue, :revenue_yoy_pct, :eps, :gross_margin,
                :operating_margin, :free_cash_flow,
                :guidance_disclosed, :guidance_text, :ingested_at,
                :shares_diluted, :effective_tax_rate, :interest_expense, :total_debt,
                :dividends_paid, :buybacks, :capex, :stockholders_equity,
                :guidance_period,
                :guidance_revenue_low, :guidance_revenue_high,
                :guidance_eps_low, :guidance_eps_high,
                :guidance_gross_margin_low, :guidance_gross_margin_high,
                :guidance_operating_margin_low, :guidance_operating_margin_high,
                :period_basis, :report_date, :accession, :source_url,
                :metric_contexts, :quality_warnings, :data_quality_status
            )
            ON CONFLICT(ticker, period) DO UPDATE SET
                company_name=excluded.company_name,
                filing_date=excluded.filing_date,
                form_type=excluded.form_type,
                revenue=excluded.revenue,
                revenue_yoy_pct=excluded.revenue_yoy_pct,
                eps=excluded.eps,
                gross_margin=excluded.gross_margin,
                operating_margin=excluded.operating_margin,
                free_cash_flow=excluded.free_cash_flow,
                guidance_disclosed=excluded.guidance_disclosed,
                guidance_text=excluded.guidance_text,
                ingested_at=excluded.ingested_at,
                shares_diluted=excluded.shares_diluted,
                effective_tax_rate=excluded.effective_tax_rate,
                interest_expense=excluded.interest_expense,
                total_debt=excluded.total_debt,
                dividends_paid=excluded.dividends_paid,
                buybacks=excluded.buybacks,
                capex=excluded.capex,
                stockholders_equity=excluded.stockholders_equity,
                guidance_period=excluded.guidance_period,
                guidance_revenue_low=excluded.guidance_revenue_low,
                guidance_revenue_high=excluded.guidance_revenue_high,
                guidance_eps_low=excluded.guidance_eps_low,
                guidance_eps_high=excluded.guidance_eps_high,
                guidance_gross_margin_low=excluded.guidance_gross_margin_low,
                guidance_gross_margin_high=excluded.guidance_gross_margin_high,
                guidance_operating_margin_low=excluded.guidance_operating_margin_low,
                guidance_operating_margin_high=excluded.guidance_operating_margin_high,
                period_basis=excluded.period_basis,
                report_date=excluded.report_date,
                accession=excluded.accession,
                source_url=excluded.source_url,
                metric_contexts=excluded.metric_contexts,
                quality_warnings=excluded.quality_warnings,
                data_quality_status=excluded.data_quality_status
        """, data)


def prune_old_metrics(ticker: str, form_type: str, keep_n: int) -> None:
    """Delete rows beyond the keep_n most recent for a given ticker and form type."""
    with sqlite3.connect(DB_PATH) as conn:
        conn.execute("""
            DELETE FROM metrics
            WHERE ticker = ? AND form_type = ?
              AND period NOT IN (
                  SELECT period FROM metrics
                  WHERE ticker = ? AND form_type = ?
                  ORDER BY filing_date DESC
                  LIMIT ?
              )
        """, (ticker, form_type, ticker, form_type, keep_n))


def _decode_row(row: sqlite3.Row) -> dict:
    result = dict(row)
    defaults = {"metric_contexts": {}, "quality_warnings": []}
    for key, default in defaults.items():
        raw = result.get(key)
        if isinstance(raw, str):
            try:
                result[key] = json.loads(raw)
            except json.JSONDecodeError:
                result[key] = default
        elif raw is None:
            result[key] = default
    result.setdefault("data_quality_status", "LEGACY_UNVERIFIED")
    return result


def get_metrics(ticker: str) -> Optional[dict]:
    """Return the most recent filing metrics for a ticker."""
    if not DB_PATH.exists():
        return None
    with sqlite3.connect(DB_PATH) as conn:
        conn.row_factory = sqlite3.Row
        row = conn.execute(
            "SELECT * FROM metrics WHERE ticker = ? ORDER BY filing_date DESC LIMIT 1",
            (ticker.upper(),),
        ).fetchone()
        return _decode_row(row) if row else None


def get_all_metrics(ticker: str) -> list[dict]:
    """Return all filing metrics for a ticker, most recent first."""
    if not DB_PATH.exists():
        return []
    with sqlite3.connect(DB_PATH) as conn:
        conn.row_factory = sqlite3.Row
        rows = conn.execute(
            "SELECT * FROM metrics WHERE ticker = ? ORDER BY filing_date DESC",
            (ticker.upper(),),
        ).fetchall()
        return [_decode_row(r) for r in rows]