stack/tests/lead/test_explore.py

"""The public-data explorer reads a mirror and does not write one."""

from __future__ import annotations

from pathlib import Path

import duckdb
import pytest

from lead.explore import catalog, guard, writeup
from lead.explore.catalog import HEAVY_ROWS, NeedsFilter, QueryTimeout, TableCard
from lead.load.frame import ensure_log, record


def _mirror(tmp_path: Path) -> duckdb.DuckDBPyConnection:
    con = duckdb.connect(str(tmp_path / "public.duckdb"))
    ensure_log(con)
    record(
        con,
        source_id="cms_demo",
        release="2024",
        sha256="a" * 64,
        path="/data/cms_demo.csv",
        table="lead_public.cms_demo",
        status="loaded",
        rows=2,
        reason="",
        actor="stack",
    )
    record(
        con,
        source_id="qcew",
        release="2020",
        sha256="b" * 64,
        path="/data/2020_qcew.zip",
        table="bls.qcew",
        status="loaded",
        rows=HEAVY_ROWS,
        reason="",
        actor="stack",
    )
    record(
        con,
        source_id="bls_ces",
        release="ces.csv",
        sha256="c" * 64,
        path="/data/ces.csv",
        table="bls.ces",
        status="loaded",
        rows=2,
        reason="",
        actor="stack",
    )
    record(
        con,
        source_id="fred",
        release="",
        sha256="",
        path="/data/fred.csv",
        table="",
        status="absent",
        rows=0,
        reason="unreachable",
        actor="stack",
    )
    record(
        con,
        source_id="cms_note",
        release="note",
        sha256="d" * 64,
        path="/data/note.txt",
        table="lead_public.cms_note",
        status="loaded",
        rows=99,
        reason="",
        actor="stack",
    )
    record(
        con,
        source_id="cms_note",
        release="note",
        sha256="d" * 64,
        path="/data/note.txt",
        table="lead_public.cms_note",
        status="skipped",
        rows=0,
        reason="kind:note",
        actor="stack",
    )
    con.execute("CREATE SCHEMA IF NOT EXISTS bls")
    con.execute(
        """
        CREATE TABLE lead_public.cms_demo (
            geo_fips VARCHAR, spine_year INTEGER, agi VARCHAR, geo_zip VARCHAR
        )
        """
    )
    con.execute(
        """
        INSERT INTO lead_public.cms_demo VALUES
            ('01001', 2010, '5', NULL),
            ('01003', 2011, '9', '30301')
        """
    )
    con.execute(
        """
        CREATE TABLE bls.qcew (
            geo_fips VARCHAR, spine_year INTEGER, series_id VARCHAR, month3_emplvl VARCHAR
        )
        """
    )
    con.execute(
        """
        INSERT INTO bls.qcew VALUES
            ('01001', 2020, 'X', '10'),
            ('01001', 2021, 'X', '11'),
            ('01003', 2020, 'Y', '9')
        """
    )
    con.execute(
        """
        CREATE TABLE bls.ces (
            series_id VARCHAR, year INTEGER, month INTEGER, value DOUBLE
        )
        """
    )
    con.execute(
        """
        INSERT INTO bls.ces VALUES
            ('CES1', 2020, 1, 10),
            ('CES1', 2022, 1, 14)
        """
    )
    return con


def _cards(con):
    return catalog.tables(con)


def _card(con, name):
    return catalog.card_for(_cards(con), name)


def test_latest_file_wins_and_its_rows_are_not_double_counted(tmp_path):
    con = _mirror(tmp_path)
    files = catalog.latest_files(con)
    note = next(row for row in files if row.source_id == "cms_note")
    assert note.status == "skipped"
    assert note.row_count == 0
    loaded = [row for row in files if row.status == "loaded"]
    assert {row.source_id for row in loaded} == {"cms_demo", "qcew", "bls_ces"}
    assert sum(row.row_count for row in loaded) == 2 + HEAVY_ROWS + 2


def test_a_heavy_table_is_not_aggregated_without_a_year(tmp_path):
    con = _mirror(tmp_path)
    qcew = _card(con, "bls.qcew")
    assert qcew.heavy
    with pytest.raises(NeedsFilter):
        catalog.count_rows(con, qcew, geo=("geo_fips", "01001"))
    with pytest.raises(NeedsFilter):
        catalog.sample(con, qcew, geo=("geo_fips", "01001"))
    with pytest.raises(NeedsFilter):
        catalog.profile(con, qcew)
    got = catalog.sample(con, qcew, geo=("geo_fips", "01001"), year=2020)
    assert [row["month3_emplvl"] for row in got] == ["10"]
    assert catalog.count_rows(con, qcew, geo=("geo_fips", "01001"), year=2020) == 1


def test_a_quoted_value_is_a_parameter(tmp_path):
    con = _mirror(tmp_path)
    demo = _card(con, "lead_public.cms_demo")
    got = catalog.sample(con, demo, geo=("geo_fips", "01001' OR '1'='1"))
    assert got == []
    plain = catalog.sample(con, demo, geo=("geo_fips", "01001"))
    assert [row["agi"] for row in plain] == ["5"]


def test_an_unfiltered_sample_of_a_heavy_table_is_capped(tmp_path):
    con = _mirror(tmp_path)
    got = catalog.sample(con, _card(con, "bls.qcew"), limit=2)
    assert len(got) == 2


def test_series_span_is_last_minus_first_in_year_order(tmp_path):
    con = _mirror(tmp_path)
    ces = _card(con, "bls.ces")
    points = catalog.series_points(
        con, ces, series_column="series_id", series_id="CES1", value_column="value"
    )
    got = catalog.span(points, year_key="year", value_key="value")
    assert (got.first_year, got.last_year, got.change) == (2020, 2022, 4)


def test_place_summary_skips_a_heavy_scan_until_a_year_is_set(tmp_path):
    con = _mirror(tmp_path)
    early = catalog.place_summary(
        con, _cards(con), geo_column="geo_fips", geo_value="01001"
    )
    by_table = {row["table"]: row for row in early}
    assert by_table["bls.qcew"]["rows"] is None
    assert by_table["bls.qcew"]["note"] == "needs a year"
    assert by_table["lead_public.cms_demo"]["rows"] == 1
    later = catalog.place_summary(
        con, _cards(con), geo_column="geo_fips", geo_value="01001", year=2020
    )
    assert {row["table"]: row["rows"] for row in later}["bls.qcew"] == 1


def test_sql_is_one_capped_read():
    wrapped = guard.prepare("SELECT * FROM bls.qcew")
    assert wrapped.endswith("LIMIT 2000")
    assert "AS _lead_q" in wrapped
    assert (
        guard.prepare("DESCRIBE lead_public.load_log")
        == "DESCRIBE lead_public.load_log"
    )
    with pytest.raises(guard.SqlRejected):
        guard.prepare("SELECT 1; SELECT 2")
    with pytest.raises(guard.SqlRejected):
        guard.prepare("DROP TABLE bls.qcew")
    with pytest.raises(guard.SqlRejected):
        guard.prepare("EXPLAIN ANALYZE SELECT 1")


def test_a_heavy_table_without_where_is_refused(tmp_path):
    con = _mirror(tmp_path)
    cards = _cards(con)
    with pytest.raises(NeedsFilter):
        catalog.run_select(con, "SELECT * FROM bls.qcew", cards)
    headers, rows = catalog.run_select(
        con, "SELECT geo_fips FROM bls.qcew WHERE spine_year = 2020", cards, cap=10
    )
    assert headers == ["geo_fips"]
    assert len(rows) == 2


def test_a_long_scan_is_interrupted():
    con = duckdb.connect()
    with pytest.raises(QueryTimeout):
        catalog.bounded_execute(
            con, "SELECT count(*) FROM range(1000000000000)", seconds=0.3
        )


def test_open_mirror_does_not_write(tmp_path):
    con = _mirror(tmp_path)
    path = tmp_path / "public.duckdb"
    con.close()
    ro = catalog.open_mirror(path)
    with pytest.raises(duckdb.Error, match="(?i)read-only"):
        ro.execute("CREATE TABLE lead_public.nope (a VARCHAR)")
    ro.close()


def test_the_writeup_counts_the_log_and_states_the_contract(tmp_path):
    con = _mirror(tmp_path)
    sections = writeup.build(con)
    text = "\n".join(
        part for section in sections for part in (section.heading, *section.body)
    )
    headings = [section.heading for section in sections]
    assert headings == [
        "Registration",
        "The mirror",
        "What the load recorded",
        "Largest tables",
        "Families",
        "How a row is keyed",
        "Geography and time",
        "Gaps",
        "What this file is not",
        "Exploring it",
    ]
    assert "10.17605/OSF.IO/8WQ3M" in text
    assert "https://osf.io/8wq3m/" in text
    assert "local DuckDB mirror" in text
    assert "read-only" in text
    assert "does not download" in text
    assert "Unity Catalog" in text
    assert "evidence ledger" in text
    assert "geo_zip" in text and "geo_zcta" in text
    assert "000" in text
    assert "time_key" in text
    assert "month 13" in text
    assert "beneficiary" in text
    assert "FRED is absent" in text
    assert "econ.fred" in text
    recorded = next(
        section for section in sections if section.heading == "What the load recorded"
    )
    assert "3 files loaded" in recorded.body[0]
    assert f"{2 + HEAVY_ROWS + 2:,} rows recorded" in recorded.body[0]
    assert ("loaded", 3) in recorded.rows
    assert ("absent", 1) in recorded.rows
    assert ("skipped", 1) in recorded.rows
    families = next(section for section in sections if section.heading == "Families")
    by_name = {row[0]: row for row in families.rows}
    assert by_name["Bureau of Labor Statistics"][1] == 2
    assert "QCEW" in by_name["Bureau of Labor Statistics"][4]
    assert by_name["CMS"][1] == 1
    largest = next(
        section for section in sections if section.heading == "Largest tables"
    )
    assert largest.rows[0][0] == "bls.qcew"
    assert largest.rows[0][3] == "yes"
    gaps = next(section for section in sections if section.heading == "Gaps")
    assert ("fred", "fred.csv", "unreachable") in gaps.rows
    spine = next(
        section for section in sections if section.heading == "How a row is keyed"
    )
    fips = next(row for row in spine.rows if row[0] == "geo_fips")
    assert "000" in fips[1]


def test_writeup_page_renders_the_recorded_count(tmp_path, monkeypatch):
    con = _mirror(tmp_path)
    path = str(tmp_path / "public.duckdb")
    con.close()
    monkeypatch.setenv("LEAD_PUBLIC_DB", path)
    from streamlit.testing.v1 import AppTest

    at = AppTest.from_string(
        "from lead.explore.views import writeup\nwriteup()\n",
        default_timeout=60,
    )
    at.run()
    assert not at.exception
    text = "\n".join(item.value for item in at.markdown)
    assert "3 files loaded" in text
    assert "FRED is absent" in text


def test_explore_command_points_streamlit_at_the_mirror(tmp_path, monkeypatch):
    db = tmp_path / "public.duckdb"
    db.write_bytes(b"")
    seen: dict = {}

    def fake_call(cmd, env=None, **_kwargs):
        seen["cmd"] = cmd
        seen["env"] = env
        return 0

    monkeypatch.setattr("subprocess.call", fake_call)
    from typer.testing import CliRunner

    from cli.lead import app

    result = CliRunner().invoke(
        app, ["explore", "--db", str(db), "--port", "8765", "--no-headless"]
    )
    assert result.exit_code == 0
    assert "http://localhost:8765" in result.output
    assert seen["env"]["LEAD_PUBLIC_DB"] == str(db)
    cmd = seen["cmd"]
    assert cmd[1:3] == ["-m", "lead.explore"]
    assert cmd[cmd.index("--server.port") + 1] == "8765"
    assert cmd[cmd.index("--server.headless") + 1] == "false"


def test_a_missing_mirror_and_a_bad_identifier_are_refused(tmp_path):
    with pytest.raises(FileNotFoundError):
        catalog.open_mirror(tmp_path / "missing.duckdb")
    with pytest.raises(ValueError):
        catalog.quote("1nope")
    assert catalog.quote("geo_fips") == '"geo_fips"'
    bare = duckdb.connect()
    assert catalog.latest_files(bare) == ()
    with pytest.raises(KeyError):
        catalog.card_for((), "lead_public.missing")


def test_filters_series_and_a_profile_stay_inside_the_cap(tmp_path):
    con = _mirror(tmp_path)
    ces = _card(con, "bls.ces")
    qcew = _card(con, "bls.qcew")
    demo = _card(con, "lead_public.cms_demo")
    assert ces.year_column() == "year"
    assert "value" in ces.measure_columns()
    plain = TableCard("lead_public", "note", (("body", "VARCHAR"),), 1, 1, ())
    assert plain.year_column() == ""
    measured = TableCard(
        "lead_public",
        "note",
        (("amount", "DECIMAL(10,2)"), ("year", "INTEGER"), ("spine_month", "INTEGER")),
        1,
        1,
        (),
    )
    assert measured.measure_columns() == ("amount",)
    got = catalog.sample(con, ces, series=("series_id", "CES1"))
    assert {row["year"] for row in got} == {2020, 2022}
    with pytest.raises(ValueError, match="nope"):
        catalog.sample(con, ces, series=("nope", "CES1"))
    with pytest.raises(ValueError, match="geo_hsa"):
        catalog.sample(con, ces, geo=("geo_hsa", "1"))
    with pytest.raises(ValueError, match="limit"):
        catalog.sample(con, ces, limit=0)
    with pytest.raises(ValueError, match="limit"):
        catalog.sample(con, ces, limit=catalog.SAMPLE_CAP + 1)
    assert catalog.series_ids(con, ces) == ("CES1",)
    with pytest.raises(NeedsFilter):
        catalog.series_ids(con, qcew)
    with pytest.raises(ValueError, match="nope"):
        catalog.series_ids(con, ces, column="nope")
    with pytest.raises(ValueError, match="column"):
        catalog.series_points(
            con, ces, series_column="nope", series_id="CES1", value_column="value"
        )
    with pytest.raises(NeedsFilter):
        catalog.series_points(
            con, ces, series_column="series_id", series_id="", value_column="value"
        )
    empty = catalog.span(
        [{"year": None, "value": 1}, {"year": "x", "value": "no"}],
        year_key="year",
        value_key="value",
    )
    assert empty.points == 0
    one = catalog.span([{"year": 2020, "value": 3}], year_key="year", value_key="value")
    assert one.change is None
    profiled = catalog.profile(con, demo)
    assert profiled["rows"] == 2
    assert profiled["year_column"] == "spine_year"
    assert profiled["geo_fips"] == 2
    con.execute("CREATE TABLE lead_public.note (body VARCHAR)")
    with pytest.raises(NeedsFilter):
        catalog.profile(con, _card(con, "lead_public.note"))
    con.execute("INSERT INTO lead_public.note VALUES ('a')")
    record(
        con,
        source_id="cms_note_body",
        release="1",
        sha256="e" * 64,
        path="/data/note-body.csv",
        table="lead_public.note",
        status="loaded",
        rows=1,
        reason="",
        actor="stack",
    )
    note = _card(con, "lead_public.note")
    assert catalog.profile(con, note)["rows"] == 1
    with pytest.raises(ValueError, match="year"):
        catalog.profile(con, note, year=2020)
    with pytest.raises(ValueError, match="year"):
        catalog.sample(con, note, year=2020)
    con.execute("CREATE TABLE lead_public.county (geo_fips VARCHAR)")
    con.execute("INSERT INTO lead_public.county VALUES ('01001')")
    looked = catalog.place_summary(
        con,
        (_card(con, "lead_public.county"), demo),
        geo_column="geo_fips",
        geo_value="01001",
        year=2010,
    )
    by_table = {row["table"]: row for row in looked}
    assert by_table["lead_public.county"]["note"] == "no year column"
    assert by_table["lead_public.county"]["rows"] == 1
    assert by_table["lead_public.cms_demo"]["rows"] == 1


def test_a_timed_out_place_is_recorded_and_the_guard_rejects_the_rest(
    tmp_path, monkeypatch
):
    con = _mirror(tmp_path)

    def _timeout(*_args, **_kwargs):
        raise QueryTimeout(1)

    monkeypatch.setattr(catalog, "count_rows", _timeout)
    rows = catalog.place_summary(
        con, _cards(con), geo_column="geo_fips", geo_value="01001", year=2020
    )
    assert {row["note"] for row in rows} == {"timed out"}
    with pytest.raises(ValueError):
        catalog.place_summary(con, (), geo_column="zip", geo_value="1")
    with pytest.raises(guard.SqlRejected, match="empty"):
        guard.prepare("  ")
    with pytest.raises(guard.SqlRejected, match="start with"):
        guard.prepare("hello")
    with pytest.raises(guard.SqlRejected, match="cap"):
        guard.prepare("SELECT 1", cap=0)


def test_the_writeup_names_every_family_and_an_empty_file():
    assert writeup.family_of("irs_soi") == "irs"
    assert writeup.family_of("census_acs") == "census"
    assert writeup.family_of("cdc_places") == "cdc"
    assert writeup.family_of("ahrq_ccsr") == "ahrq"
    assert writeup.family_of("nlm_rxnorm") == "nlm"
    assert writeup.family_of("hrsa_hpsa") == "hrsa"
    assert writeup.family_of("ecqi_measures") == "ecqi"
    assert writeup.family_of("ma_penetration") == "ma"
    assert writeup.family_of("ma") == "ma"
    assert writeup.family_of("fred") == "fred"
    assert writeup.family_of("noaa") == "other"
    recorded = writeup._recorded({"loaded": 1, "partial": 2}, 5, 1)
    assert ("partial", 2) in recorded.rows
    bare = duckdb.connect()
    mirror = next(
        section for section in writeup.build(bare) if section.heading == "The mirror"
    )
    assert "no load log and no observation tables" in mirror.body[1]
    bare.execute("CREATE TABLE notes (body VARCHAR)")
    mirror = next(
        section for section in writeup.build(bare) if section.heading == "The mirror"
    )
    assert "no load log" in mirror.body[1]
    present = writeup._gaps(
        (),
        (TableCard("econ", "fred", (), 1, 1, ()),),
    )
    assert present.body[0] == "`econ.fred` is present in this file."
    other = writeup._gaps(
        (
            catalog.FileRow(
                "fred", "fred.csv", "/data/fred.csv", "", "skipped", 0, "later"
            ),
        ),
        (),
    )
    assert "skipped" in other.body[0]
    assert "reason later" in other.body[0]
    none = writeup._gaps((), ())
    assert none.body[0] == (
        "The load log has no FRED row, and `econ.fred` is not in this file."
    )
    assert "No file's latest status is skipped." in none.body