stack/src/lead/explore/catalog.py

"""Questions the explorer asks the public-data mirror.

Every function takes an open DuckDB connection. None of them download,
publish, or write. A table whose load log recorded at least
:data:`HEAVY_ROWS` is not aggregated unless the caller passes a year or
a series, because a geography predicate alone still reads the file.
"""

from __future__ import annotations

import re
import threading
from dataclasses import dataclass
from pathlib import Path

import duckdb

from lead.explore.guard import prepare

DEFAULT_DB = "/home/kert/stack/data/lead/public.duckdb"
HEAVY_ROWS = 2_000_000
SAMPLE_CAP = 500
SQL_CAP = 2_000

_IDENT = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")
_HIDDEN = frozenset({"information_schema", "pg_catalog"})
PROVENANCE = (
    "release",
    "content_sha256",
    "source_path",
    "member",
    "geo_fips",
    "geo_zcta",
    "geo_zip",
    "geo_hsa",
    "geo_hrr",
    "geo_cbsa",
    "spine_year",
    "spine_quarter",
    "spine_month",
    "time_key",
    "file_year",
)
YEAR_COLUMNS = ("spine_year", "file_year", "year")
GEO_COLUMNS = (
    "geo_fips",
    "geo_zcta",
    "geo_zip",
    "geo_hsa",
    "geo_hrr",
    "geo_cbsa",
)
_YEAR_ROLES = frozenset(YEAR_COLUMNS) | {"spine_quarter", "spine_month"}
_NUMERIC = frozenset(
    {
        "TINYINT",
        "SMALLINT",
        "INTEGER",
        "BIGINT",
        "HUGEINT",
        "UTINYINT",
        "USMALLINT",
        "UINTEGER",
        "UBIGINT",
        "UHUGEINT",
        "FLOAT",
        "DOUBLE",
        "REAL",
        "DECIMAL",
    }
)


class NeedsFilter(ValueError):
    """A large table needs a year or a series before it is aggregated."""

    def __init__(self, table: str) -> None:
        self.table = table
        super().__init__(f"{table} needs a year or a series before it is read")


class QueryTimeout(TimeoutError):
    """DuckDB interrupted the scan."""

    def __init__(self, seconds: float) -> None:
        self.seconds = seconds
        super().__init__(f"stopped after {seconds:g}s")


@dataclass(frozen=True)
class FileRow:
    source_id: str
    release: str
    path: str
    table_name: str
    status: str
    row_count: int
    reason: str


@dataclass(frozen=True)
class TableCard:
    schema: str
    name: str
    columns: tuple[tuple[str, str], ...]
    loaded_files: int
    loaded_rows: int
    releases: tuple[str, ...]

    @property
    def qualified(self) -> str:
        return f"{self.schema}.{self.name}"

    @property
    def heavy(self) -> bool:
        return self.loaded_rows >= HEAVY_ROWS

    def column_names(self) -> set[str]:
        return {name for name, _type in self.columns}

    def year_column(self) -> str:
        have = self.column_names()
        for name in YEAR_COLUMNS:
            if name in have:
                return name
        return ""

    def measure_columns(self) -> tuple[str, ...]:
        """Numeric columns that are not year or quarter markers."""
        return tuple(
            name
            for name, kind in self.columns
            if name not in _YEAR_ROLES and _base_type(kind) in _NUMERIC
        )


@dataclass(frozen=True)
class SeriesSpan:
    points: int
    first_year: int | None
    last_year: int | None
    first_value: float | None
    last_value: float | None
    change: float | None


def quote(name: str) -> str:
    if _IDENT.fullmatch(name or "") is None:
        raise ValueError(name)
    return '"' + name + '"'


def qualify(schema: str, name: str) -> str:
    return f"{quote(schema)}.{quote(name)}"


def open_mirror(path: str | Path):
    """Open *path* read-only. A missing file raises ``FileNotFoundError``."""
    src = Path(path)
    if not src.is_file():
        raise FileNotFoundError(src)
    con = duckdb.connect(str(src), read_only=True)
    con.execute("SET preserve_insertion_order = false")
    con.execute("SET memory_limit = '8GB'")
    con.execute("SET threads = 4")
    return con


def bounded_execute(con, sql: str, params: list | None = None, *, seconds: float = 45):
    """Run *sql* and interrupt it if it is still going after *seconds*."""
    timer = threading.Timer(seconds, con.interrupt)
    timer.start()
    try:
        return con.execute(sql, params or [])
    except duckdb.InterruptException as exc:
        raise QueryTimeout(seconds) from exc
    finally:
        timer.cancel()


def latest_files(con) -> tuple[FileRow, ...]:
    """Newest load-log row for each source and path.

    A rerun that finds the file already loaded does not append a row, so
    the newest row still carries the count from the read that inserted.
    """
    found = con.execute(
        """
        SELECT 1 FROM information_schema.tables
        WHERE table_schema = 'lead_public' AND table_name = 'load_log'
        """
    ).fetchone()
    if found is None:
        return ()
    rows = con.execute(
        """
        SELECT source_id, release, path, table_name, status, row_count, reason
        FROM (
            SELECT source_id, release, path, table_name, status, row_count, reason,
                   row_number() OVER (
                       PARTITION BY source_id, path
                       ORDER BY rowid DESC
                   ) AS _n
            FROM lead_public.load_log
        ) t
        WHERE _n = 1
        ORDER BY source_id, release, path
        """
    ).fetchall()
    return tuple(
        FileRow(
            source_id=str(row[0] or ""),
            release=str(row[1] or ""),
            path=str(row[2] or ""),
            table_name=str(row[3] or ""),
            status=str(row[4] or ""),
            row_count=int(row[5] or 0),
            reason=str(row[6] or ""),
        )
        for row in rows
    )


def status_counts(files: tuple[FileRow, ...]) -> dict[str, int]:
    counts: dict[str, int] = {}
    for row in files:
        counts[row.status] = counts.get(row.status, 0) + 1
    return counts


def tables(con, files: tuple[FileRow, ...] | None = None) -> tuple[TableCard, ...]:
    """Every base table outside the system schemas, with log totals attached."""
    logged = latest_files(con) if files is None else files
    totals: dict[str, list] = {}
    for row in logged:
        if row.status != "loaded" or not row.table_name:
            continue
        slot = totals.setdefault(row.table_name, [0, 0, set()])
        slot[0] += 1
        slot[1] += row.row_count
        if row.release:
            slot[2].add(row.release)
    meta = con.execute(
        """
        SELECT table_schema, table_name, column_name, data_type
        FROM information_schema.columns
        WHERE table_schema NOT IN ('information_schema', 'pg_catalog')
        ORDER BY table_schema, table_name, ordinal_position
        """
    ).fetchall()
    grouped: dict[tuple[str, str], list[tuple[str, str]]] = {}
    for schema, name, column, kind in meta:
        if schema in _HIDDEN:
            continue
        grouped.setdefault((str(schema), str(name)), []).append(
            (str(column), str(kind))
        )
    cards = []
    for schema, name in sorted(grouped):
        qualified_name = f"{schema}.{name}"
        files_n, rows_n, releases = totals.get(qualified_name, [0, 0, set()])
        cards.append(
            TableCard(
                schema=schema,
                name=name,
                columns=tuple(grouped[(schema, name)]),
                loaded_files=files_n,
                loaded_rows=rows_n,
                releases=tuple(sorted(releases)),
            )
        )
    return tuple(cards)


def card_for(
    cards: tuple[TableCard, ...] | list[TableCard], qualified_name: str
) -> TableCard:
    for card in cards:
        if card.qualified == qualified_name:
            return card
    raise KeyError(qualified_name)


def _where(
    card: TableCard,
    *,
    geo: tuple[str, str] | None,
    year: int | None,
    series: tuple[str, str] | None,
) -> tuple[str, list]:
    names = card.column_names()
    clauses: list[str] = []
    params: list = []
    if geo is not None:
        column, value = geo
        if column not in GEO_COLUMNS or column not in names:
            raise ValueError(column)
        clauses.append(f"{quote(column)} = ?")
        params.append(value)
    if year is not None:
        column = card.year_column()
        if not column:
            raise ValueError("no year column")
        clauses.append(f"{quote(column)} = ?")
        params.append(int(year))
    if series is not None:
        column, value = series
        if column not in names:
            raise ValueError(column)
        clauses.append(f"{quote(column)} = ?")
        params.append(value)
    text = (" WHERE " + " AND ".join(clauses)) if clauses else ""
    return text, params


def sample(
    con,
    card: TableCard,
    *,
    geo: tuple[str, str] | None = None,
    year: int | None = None,
    series: tuple[str, str] | None = None,
    limit: int = 100,
    seconds: float = 20,
) -> list[dict]:
    """Return up to *limit* rows. No ``ORDER BY``: a sort would read the file.

    An unfiltered sample stops after *limit* rows. A filter on a heavy
    table still has to find those rows, so it needs a year or a series.
    """
    if card.heavy and (geo is not None or year is not None or series is not None):
        if year is None and series is None:
            raise NeedsFilter(card.qualified)
    where, params = _where(card, geo=geo, year=year, series=series)
    bound = int(limit)
    if bound < 1 or bound > SAMPLE_CAP:
        raise ValueError("limit")
    sql = f"SELECT * FROM {qualify(card.schema, card.name)}{where} LIMIT {bound}"
    return _dicts(bounded_execute(con, sql, params, seconds=seconds))


def count_rows(
    con,
    card: TableCard,
    *,
    geo: tuple[str, str] | None = None,
    year: int | None = None,
    seconds: float = 45,
) -> int:
    if card.heavy and year is None:
        raise NeedsFilter(card.qualified)
    where, params = _where(card, geo=geo, year=year, series=None)
    sql = f"SELECT count(*) FROM {qualify(card.schema, card.name)}{where}"
    return int(bounded_execute(con, sql, params, seconds=seconds).fetchone()[0])


def series_ids(
    con, card: TableCard, column: str = "series_id", *, limit: int = 500
) -> tuple[str, ...]:
    if card.heavy:
        raise NeedsFilter(card.qualified)
    if column not in card.column_names():
        raise ValueError(column)
    bound = min(int(limit), SAMPLE_CAP)
    sql = (
        f"SELECT DISTINCT {quote(column)} FROM {qualify(card.schema, card.name)} "
        f"WHERE {quote(column)} IS NOT NULL ORDER BY 1 LIMIT {bound}"
    )
    return tuple(
        str(row[0]) for row in bounded_execute(con, sql, seconds=20).fetchall()
    )


def series_points(
    con,
    card: TableCard,
    *,
    series_column: str,
    series_id: str,
    value_column: str,
    seconds: float = 30,
) -> list[dict]:
    names = card.column_names()
    if series_column not in names or value_column not in names:
        raise ValueError("column")
    if not series_id:
        raise NeedsFilter(card.qualified)
    year = card.year_column()
    month = (
        "spine_month" if "spine_month" in names else "month" if "month" in names else ""
    )
    picked = [series_column, value_column]
    for extra in (year, month, "period", "time_key"):
        if extra and extra not in picked and extra in names:
            picked.append(extra)
    cols = ", ".join(quote(name) for name in picked)
    order = f" ORDER BY {quote(year)}" if year else ""
    if month and year:
        order += f", {quote(month)}"
    sql = (
        f"SELECT {cols} FROM {qualify(card.schema, card.name)} "
        f"WHERE {quote(series_column)} = ?{order}"
    )
    return _dicts(bounded_execute(con, sql, [series_id], seconds=seconds))


def span(points: list[dict], *, year_key: str, value_key: str) -> SeriesSpan:
    """Last numeric value minus the first, in year order. Not a model."""
    usable: list[tuple[int, float]] = []
    for row in points:
        raw_year = row.get(year_key)
        raw_value = row.get(value_key)
        if raw_year is None or raw_value is None:
            continue
        try:
            usable.append((int(raw_year), float(raw_value)))
        except (TypeError, ValueError):
            continue
    if not usable:
        return SeriesSpan(0, None, None, None, None, None)
    usable.sort()
    change = usable[-1][1] - usable[0][1] if len(usable) > 1 else None
    return SeriesSpan(
        len(usable),
        usable[0][0],
        usable[-1][0],
        usable[0][1],
        usable[-1][1],
        change,
    )


def profile(
    con, card: TableCard, *, year: int | None = None, seconds: float = 45
) -> dict:
    """One scan: row count, year range, and non-null counts for spine columns."""
    if card.heavy and year is None:
        raise NeedsFilter(card.qualified)
    if (
        card.loaded_rows == 0
        and card.qualified != "lead_public.load_log"
        and year is None
    ):
        raise NeedsFilter(card.qualified)
    names = card.column_names()
    spine = [name for name in PROVENANCE if name in names]
    pieces = ["count(*)"]
    for name in spine:
        pieces.append(f"count({quote(name)})")
    year_name = card.year_column()
    if year_name:
        pieces.append(f"min({quote(year_name)})")
        pieces.append(f"max({quote(year_name)})")
    where, params = ("", [])
    if year is not None:
        if not year_name:
            raise ValueError("no year column")
        where, params = f" WHERE {quote(year_name)} = ?", [int(year)]
    sql = f"SELECT {', '.join(pieces)} FROM {qualify(card.schema, card.name)}{where}"
    row = bounded_execute(con, sql, params, seconds=seconds).fetchone()
    out: dict = {"rows": int(row[0])}
    cursor = 1
    for name in spine:
        out[name] = int(row[cursor])
        cursor += 1
    if year_name:
        out["year_min"] = row[cursor]
        out["year_max"] = row[cursor + 1]
        out["year_column"] = year_name
    return out


def place_summary(
    con,
    cards: tuple[TableCard, ...] | list[TableCard],
    *,
    geo_column: str,
    geo_value: str,
    year: int | None = None,
    seconds: float = 45,
) -> list[dict]:
    """Count *geo_value* in every table that carries *geo_column*.

    A heavy table is not scanned until *year* is set. A timeout on one
    table is recorded on that row and the rest still run.
    """
    if geo_column not in GEO_COLUMNS:
        raise ValueError(geo_column)
    rows: list[dict] = []
    for card in cards:
        if geo_column not in card.column_names():
            continue
        entry = {
            "table": card.qualified,
            "rows": None,
            "year_column": card.year_column(),
            "note": "",
        }
        use_year = year if card.year_column() else None
        if year is not None and not card.year_column():
            entry["note"] = "no year column"
        if card.heavy and use_year is None:
            entry["note"] = "needs a year"
            rows.append(entry)
            continue
        try:
            entry["rows"] = count_rows(
                con,
                card,
                geo=(geo_column, geo_value),
                year=use_year,
                seconds=seconds,
            )
        except QueryTimeout:
            entry["note"] = "timed out"
        rows.append(entry)
    return rows


def run_select(
    con,
    sql: str,
    cards: tuple[TableCard, ...] | list[TableCard],
    *,
    cap: int = SQL_CAP,
    seconds: float = 45,
) -> tuple[list[str], list[tuple]]:
    """Run one guarded read. A heavy table without ``WHERE`` is refused."""
    statement = prepare(sql, cap=cap)
    if statement.lower().startswith("select"):
        folded = sql.replace('"', "")
        for card in cards:
            if not card.heavy:
                continue
            pattern = (
                rf"(?i)(?<![A-Za-z0-9_]){re.escape(card.qualified)}(?![A-Za-z0-9_])"
            )
            if re.search(pattern, folded) and re.search(r"(?is)\bwhere\b", sql) is None:
                raise NeedsFilter(card.qualified)
    cursor = bounded_execute(con, statement, seconds=seconds)
    headers = [item[0] for item in cursor.description] if cursor.description else []
    return headers, cursor.fetchall()


def _dicts(cursor) -> list[dict]:
    headers = [item[0] for item in cursor.description]
    return [dict(zip(headers, row, strict=True)) for row in cursor.fetchall()]


def _base_type(kind: str) -> str:
    return kind.split("(", 1)[0].upper()