refresh_etl_scenarios.py

#!/usr/bin/env python3
"""
refresh_etl_scenarios.py - regenerate the ETL-side scenario list.

Reads the working CSV
(`etl/ingestion/scenario_listing/model_run_file_source_working.csv`) and
writes `etl/common/etl_scenarios.py` containing a frozenset of every
short_code the ETL pipeline is intended to process.

Filter:
- Excludes rows whose `download_status` column is `skip` or `retired`.
- Includes every other well-formed row (blank `download_status`, `done`,
  `needs_review`, anything else).

Consumers (`run_all.py`, `verify_all_sections.py`) get the broader
"what the ETL touches" set. For the narrower "what is live on the
public website" set, use `etl/common/active_scenarios.py`, regenerated
by `refresh_active_scenarios.py` from the API.

Usage:
  python etl/ingestion/tools/refresh_etl_scenarios.py
  python etl/ingestion/tools/refresh_etl_scenarios.py --working-csv path/to.csv
  python etl/ingestion/tools/refresh_etl_scenarios.py --dry-run
"""

from __future__ import annotations

# Make `etl.common` importable when this script is invoked directly with
# `python etl/ingestion/tools/refresh_etl_scenarios.py`. See
# etl/common/__init__.py for the rationale.
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[3]))

import argparse  # noqa: E402
import csv  # noqa: E402
import logging  # noqa: E402
from datetime import datetime, timezone  # noqa: E402
from typing import List, Set  # noqa: E402

from etl.ingestion.lib.config import COLUMN_MAP, WORKING_CSV_PATH  # noqa: E402

REPO_ROOT = Path(__file__).resolve().parents[3]
DEFAULT_WORKING_CSV = REPO_ROOT / WORKING_CSV_PATH
DEFAULT_ETL_SCENARIOS_PY = REPO_ROOT / "etl" / "common" / "etl_scenarios.py"

EXCLUDED_DOWNLOAD_STATUSES: Set[str] = {"skip", "retired"}

SHORT_CODE_COL = COLUMN_MAP["short_code"]
DOWNLOAD_STATUS_COL = COLUMN_MAP["download_status"]

log = logging.getLogger("refresh_etl_scenarios")


def _read_working_csv(csv_path: Path) -> List[dict]:
    """Read every row of the working CSV as a dict, preserving order."""
    if not csv_path.exists():
        raise SystemExit(
            f"\nWorking CSV not found: {csv_path}\n"
            f"Bootstrap from the reference copy first. See etl/README.md.\n"
        )
    with csv_path.open(encoding="utf-8-sig") as f:
        reader = csv.DictReader(f)
        if SHORT_CODE_COL not in (reader.fieldnames or []):
            raise SystemExit(
                f"\nWorking CSV is missing required column '{SHORT_CODE_COL}'.\n"
                f"  CSV path: {csv_path}\n"
            )
        return list(reader)


def _select_short_codes(rows: List[dict]) -> List[str]:
    """Return sorted unique short_codes after the download_status filter."""
    keep: Set[str] = set()
    excluded = 0
    blank = 0
    for raw in rows:
        sc = (raw.get(SHORT_CODE_COL) or "").strip().lower()
        if not sc:
            blank += 1
            continue
        status = (raw.get(DOWNLOAD_STATUS_COL) or "").strip().lower()
        if status in EXCLUDED_DOWNLOAD_STATUSES:
            excluded += 1
            continue
        keep.add(sc)
    log.info(
        "Read %d rows; kept %d, excluded %d by download_status, skipped %d with blank short_code",
        len(rows), len(keep), excluded, blank,
    )
    return sorted(keep)


def _render_etl_scenarios_py(short_codes: List[str], csv_path: Path) -> str:
    """Render the contents of etl/common/etl_scenarios.py."""
    now = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
    try:
        rel_csv = csv_path.resolve().relative_to(REPO_ROOT)
        csv_label = str(rel_csv)
    except ValueError:
        csv_label = str(csv_path)
    excluded_list = ", ".join(sorted(f'"{s}"' for s in EXCLUDED_DOWNLOAD_STATUSES))
    if short_codes:
        body = "\n".join(f'    "{code}",' for code in short_codes)
        literal = "frozenset({\n" + body + "\n})"
    else:
        literal = "frozenset()"
    return (
        '"""Scenarios the ETL pipeline is intended to process.\n'
        "\n"
        f"Auto-generated by `etl/ingestion/tools/refresh_etl_scenarios.py`\n"
        f"on {now} from `{csv_label}`. Rows whose `download_status` is in\n"
        f"{{{excluded_list}}} are excluded; everything else is included.\n"
        "Do not edit by hand. Re-run the refresh script to regenerate.\n"
        "\n"
        "Use this set for anything that runs against raw ETL output in S3 or\n"
        "the per-scenario statistics in the DB: `run_all.py --all-scenarios`,\n"
        "`verify_all_sections.py --all-scenarios`.\n"
        "For the narrower curated/public set (what the website serves), import\n"
        "`ACTIVE_SCENARIOS` from `etl.common.active_scenarios` instead.\n"
        '"""\n'
        "\n"
        "from __future__ import annotations\n"
        "\n"
        f"ETL_SCENARIOS: frozenset[str] = {literal}\n"
    )


def main() -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)-5s %(message)s",
        datefmt="%Y-%m-%dT%H:%M:%SZ",
    )
    parser = argparse.ArgumentParser(
        description="Refresh etl/common/etl_scenarios.py from the working CSV."
    )
    parser.add_argument("--working-csv", default=str(DEFAULT_WORKING_CSV),
                        help=f"Path to the working CSV (default: {DEFAULT_WORKING_CSV})")
    parser.add_argument("--etl-py", default=str(DEFAULT_ETL_SCENARIOS_PY),
                        help=f"Path to write the ETL_SCENARIOS Python module (default: {DEFAULT_ETL_SCENARIOS_PY})")
    parser.add_argument("--dry-run", action="store_true",
                        help="Print the generated module to stdout, don't write.")
    args = parser.parse_args()

    csv_path = Path(args.working_csv)
    rows = _read_working_csv(csv_path)
    short_codes = _select_short_codes(rows)
    new_py = _render_etl_scenarios_py(short_codes, csv_path)

    etl_py = Path(args.etl_py)
    current_py = etl_py.read_text() if etl_py.exists() else ""

    if args.dry_run:
        print("--- new etl/common/etl_scenarios.py ---")
        print(new_py)
        print("\n--- diff summary ---")
        print(f"etl_scenarios.py: {'(no change)' if current_py == new_py else '(would change)'}")
        return

    if current_py == new_py:
        log.info("etl/common/etl_scenarios.py already up to date (%d scenarios)", len(short_codes))
        return

    etl_py.parent.mkdir(parents=True, exist_ok=True)
    etl_py.write_text(new_py)
    log.info("Wrote %s (%d scenarios)", etl_py, len(short_codes))


if __name__ == "__main__":
    main()