stack/src/lead/load/frame.py

"""Shared DuckDB writes for the public-data mirror.

The log DDL lives here, not in ``lead.ledger.schema``, so the Delta
schema string stays free of a catalog token. Observation tables are
created when a file is loaded. This module does not download and does
not open a URL.
"""

from __future__ import annotations

import hashlib
import re
from datetime import datetime, timezone
from pathlib import Path

from cms.datacms import sanitize_column_name
from lead.load.errors import HashMismatch
from lead.load.spine_map import SpineMap, file_quarter, file_year, spine_columns

LOAD_LOG_DDL = """
CREATE SCHEMA IF NOT EXISTS lead_public;
CREATE TABLE IF NOT EXISTS lead_public.load_log (
    source_id VARCHAR,
    release VARCHAR,
    sha256 VARCHAR,
    path VARCHAR,
    table_name VARCHAR,
    status VARCHAR,
    row_count BIGINT,
    reason VARCHAR,
    actor VARCHAR,
    created_at VARCHAR
)
"""

_IDENT = re.compile(r"[a-z_][a-z0-9_]*")
_RESERVED = frozenset(
    {
        "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",
    }
)
_INTS = frozenset({"spine_year", "spine_quarter", "spine_month", "file_year"})
_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",
)


def ensure_log(con) -> None:
    for statement in LOAD_LOG_DDL.split(";"):
        text = statement.strip()
        if text:
            con.execute(text)


def qualify(table: str) -> str:
    """``schema.name``, defaulting a bare name to ``lead_public``."""
    parts = table.split(".")
    if len(parts) == 1:
        parts = ["lead_public", parts[0]]
    if len(parts) != 2 or any(_IDENT.fullmatch(part) is None for part in parts):
        raise ValueError(table)
    return ".".join(parts)


def hash_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def require_hash(path: Path, expected: str) -> str:
    """Hash *path* and require it to equal *expected*.

    Raises:
        ValueError: *expected* is blank.
        HashMismatch: the bytes are not that digest. Nothing was inserted.
    """
    wanted = (expected or "").strip().lower()
    if not wanted:
        raise ValueError("sha256 is required")
    actual = hash_file(path)
    if actual != wanted:
        raise HashMismatch(path, wanted, actual)
    return actual


def already_loaded(con, *, source_id: str, path: str, sha256: str) -> bool:
    ensure_log(con)
    row = con.execute(
        """
        SELECT 1 FROM lead_public.load_log
        WHERE source_id = ? AND path = ? AND sha256 = ? AND status = 'loaded'
        """,
        [source_id, path, sha256.strip().lower()],
    ).fetchone()
    return row is not None


def record(
    con,
    *,
    source_id: str,
    release: str,
    sha256: str,
    path: str,
    table: str,
    status: str,
    rows: int,
    reason: str,
    actor: str,
) -> None:
    ensure_log(con)
    con.execute(
        """
        INSERT INTO lead_public.load_log (
            source_id, release, sha256, path, table_name, status,
            row_count, reason, actor, created_at
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """,
        [
            source_id,
            release,
            sha256,
            path,
            table,
            status,
            rows,
            reason,
            actor,
            datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        ],
    )


def latest(con, *, source_id: str, path: str) -> tuple[str, str]:
    row = con.execute(
        """
        SELECT status, reason FROM lead_public.load_log
        WHERE source_id = ? AND path = ?
        ORDER BY rowid DESC
        LIMIT 1
        """,
        [source_id, path],
    ).fetchone()
    if row is None:
        return "", ""
    return str(row[0]), str(row[1] or "")


def quote_ident(name: str) -> str:
    return '"' + name.replace('"', '""') + '"'


def unique_names(raw_names: list[str]) -> list[str]:
    """SQL-safe, unique, non-reserved names in header order."""
    seen: dict[str, int] = {}
    out: list[str] = []
    for raw in raw_names:
        name = sanitize_column_name(raw) or "col"
        if name in _RESERVED:
            name = "src_" + name
        seen[name] = seen.get(name, 0) + 1
        if seen[name] > 1:
            name = f"{name}_{seen[name]}"
        out.append(name)
    return out


def _typed(name: str) -> str:
    if name in _INTS:
        return "INTEGER"
    return "VARCHAR"


def ensure_table(con, table: str, columns: list[str]) -> None:
    schema, name = qualify(table).split(".")
    con.execute(f"CREATE SCHEMA IF NOT EXISTS {schema}")
    found = con.execute(
        """
        SELECT 1 FROM information_schema.tables
        WHERE table_schema = ? AND table_name = ?
        """,
        [schema, name],
    ).fetchone()
    qualified = f"{schema}.{name}"
    if found is None:
        defs = ", ".join(f"{column} {_typed(column)}" for column in columns)
        con.execute(f"CREATE TABLE {qualified} ({defs})")
        return
    have = {row[0] for row in con.execute(f"DESCRIBE {qualified}").fetchall()}
    for column in columns:
        if column not in have:
            con.execute(f"ALTER TABLE {qualified} ADD COLUMN {column} {_typed(column)}")


def _five(column: str, *, drop_state_total: bool) -> str:
    state = f" AND right(trim({column}), 3) <> '000'" if drop_state_total else ""
    return (
        f"CASE WHEN regexp_full_match(trim({column}), '[0-9]{{5}}'){state} "
        f"THEN trim({column}) END"
    )


def _compose(state: str, county: str) -> str:
    return (
        f"CASE WHEN regexp_full_match(trim({state}), '[0-9]{{1,2}}') "
        f"AND regexp_full_match(trim({county}), '[0-9]{{1,3}}') "
        f"AND lpad(trim({county}), 3, '0') <> '000' "
        f"THEN lpad(trim({state}), 2, '0') || lpad(trim({county}), 3, '0') END"
    )


def _int_expr(column: str, low: int, high: int) -> str:
    cast = f"TRY_CAST(TRY_CAST(trim({column}) AS DOUBLE) AS INTEGER)"
    return f"CASE WHEN {cast} BETWEEN {low} AND {high} THEN {cast} END"


def _fips_expr(spine: SpineMap) -> str:
    parts: list[str] = []
    if spine.fips:
        parts.append(_five(spine.fips, drop_state_total=True))
    if spine.county_fips:
        parts.append(_five(spine.county_fips, drop_state_total=True))
        if spine.state_fips:
            parts.append(_compose(spine.state_fips, spine.county_fips))
    if not parts:
        return "NULL"
    if len(parts) == 1:
        return parts[0]
    return "COALESCE(" + ", ".join(parts) + ")"


def _text_expr(column: str) -> str:
    return f"NULLIF(trim({column}), '')"


def spine_projection(
    columns: list[str],
    *,
    file_name: str,
    when: tuple[int, int] | None,
) -> tuple[str, str, str, str]:
    """Return SQL for year, quarter-before-month, month, and file_year.

    Quarter-from-month is applied by the caller, which can see the month.
    """
    spine = spine_columns(columns)
    named_year = file_year(file_name)
    named_quarter = file_quarter(file_name)
    if spine.year:
        year_sql = _int_expr(spine.year, 1900, 2100)
    elif when is not None:
        year_sql = str(int(when[0]))
    elif named_year is not None:
        year_sql = str(named_year)
    else:
        year_sql = "NULL"
    quarter_parts: list[str] = []
    if spine.quarter:
        quarter_parts.append(_int_expr(spine.quarter, 1, 4))
    if spine.period:
        quarter_parts.append(
            "CASE WHEN regexp_full_match(trim(period), 'Q0[1-4]') "
            "THEN TRY_CAST(substr(trim(period), 3) AS INTEGER) END"
        )
    quarter_sql = (
        "COALESCE(" + ", ".join(quarter_parts) + ")" if quarter_parts else "NULL"
    )
    month_parts: list[str] = []
    if spine.month:
        month_parts.append(_int_expr(spine.month, 1, 12))
    if spine.period:
        month_parts.append(
            "CASE WHEN regexp_full_match(trim(period), 'M(0[1-9]|1[0-2])') "
            "THEN TRY_CAST(substr(trim(period), 2) AS INTEGER) END"
        )
    if when is not None and not month_parts:
        month_parts.append(str(int(when[1])))
    month_sql = "COALESCE(" + ", ".join(month_parts) + ")" if month_parts else "NULL"
    file_year_sql = "NULL" if named_year is None else str(named_year)
    # Stash the filename quarter in the quarter SQL when the file has no
    # quarter column and no month. A month still wins, in the outer query.
    if named_quarter is not None and quarter_sql == "NULL" and month_sql == "NULL":
        quarter_sql = str(named_quarter)
    return year_sql, quarter_sql, month_sql, file_year_sql


def insert_frame(
    con,
    table: str,
    temp: str,
    columns: list[str],
    *,
    release: str,
    sha256: str,
    source_path: str,
    member: str,
    file_name: str,
    when: tuple[int, int] | None,
    member_column: str = "",
) -> int:
    """Append *temp* into *table* with spine and provenance columns.

    *member_column*, when set, is a column of *temp* and supplies each
    row's member name. *member* is used for every row otherwise.
    """
    qualified = qualify(table)
    year_sql, quarter_sql, month_sql, file_year_sql = spine_projection(
        columns, file_name=file_name, when=when
    )
    spine = spine_columns(columns)
    count = con.execute(f"SELECT COUNT(*) FROM {temp}").fetchone()[0]
    carried = columns if not member_column else [*columns, member_column]
    select_cols = ", ".join(carried)
    con.execute(
        f"""
        CREATE OR REPLACE TEMP TABLE _lead_spine AS
        SELECT {select_cols},
               ({year_sql}) AS spine_year,
               ({quarter_sql}) AS quarter_from_column,
               ({month_sql}) AS spine_month,
               ({file_year_sql}) AS file_year,
               ({_fips_expr(spine)}) AS geo_fips,
               ({_five(spine.zcta, drop_state_total=False) if spine.zcta else "NULL"}) AS geo_zcta,
               ({_zip_expr(spine.zip) if spine.zip else "NULL"}) AS geo_zip,
               ({_text_expr(spine.hsa) if spine.hsa else "NULL"}) AS geo_hsa,
               ({_text_expr(spine.hrr) if spine.hrr else "NULL"}) AS geo_hrr,
               ({_text_expr(spine.cbsa) if spine.cbsa else "NULL"}) AS geo_cbsa
        FROM {temp}
        """
    )
    out_columns = [
        *columns,
        *_PROVENANCE,
    ]
    ensure_table(con, qualified, out_columns)
    member_sql = member_column if member_column else "?"
    params: list[object] = [release, sha256, source_path]
    if not member_column:
        params.append(member)
    con.execute(
        f"""
        INSERT INTO {qualified} ({", ".join(out_columns)})
        SELECT {", ".join(columns)},
               ?,
               ?,
               ?,
               {member_sql},
               geo_fips, geo_zcta, geo_zip, geo_hsa, geo_hrr, geo_cbsa,
               spine_year,
               COALESCE(
                   quarter_from_column,
                   CASE WHEN spine_month BETWEEN 1 AND 12
                        THEN 1 + (spine_month - 1) // 3 END
               ),
               spine_month,
               CASE
                   WHEN spine_year IS NOT NULL
                    AND COALESCE(
                        quarter_from_column,
                        CASE WHEN spine_month BETWEEN 1 AND 12
                             THEN 1 + (spine_month - 1) // 3 END
                    ) BETWEEN 1 AND 4
                   THEN spine_year || 'Q' || COALESCE(
                        quarter_from_column,
                        CASE WHEN spine_month BETWEEN 1 AND 12
                             THEN 1 + (spine_month - 1) // 3 END
                   )
               END,
               file_year
        FROM _lead_spine
        """,
        params,
    )
    return int(count)


def _zip_expr(column: str) -> str:
    return (
        f"CASE WHEN regexp_full_match(trim({column}), '[0-9]{{5}}') THEN trim({column}) "
        f"WHEN regexp_full_match(trim({column}), '[0-9]{{9}}') THEN left(trim({column}), 5) END"
    )