stack/dev/scripts/hoard_aco_versions.py

"""Backfill every published version of the ACO / CMMI-model datasets.

Reads the data.cms.gov catalog snapshot staged with the LEAD enrichment
download, and for each dataset version in the ACO and payment-model
families downloads every CSV/ZIP/XLSX distribution that is not already
on disk into ``data/public/lead_enrichment/<source_id>/<release>/``
(the layout the LEAD download established: source_id is
``cms_`` + the slugged series title cut to 74 chars; release is the
version's title date with ``-`` -> ``_``).

Polite (one request every --pause seconds), resumable (``.part`` files
with Range), and never overwrites: a file already present is kept, and a
changed upstream copy is stored beside it as ``<stem>.<sha8><ext>``.
Every fetched or already-held file is recorded in
``_catalogs/aco_backfill_manifest.json``.

    uv run --no-sync python dev/scripts/hoard_aco_versions.py [--limit N] [--budget SECS] [--dry-run]
"""

from __future__ import annotations

import argparse
import datetime as dt
import hashlib
import json
import re
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path

UA = "stack-pfs-dossier/1.0 (ci@mail.fhirworx.io)"
ROOT = Path("data/public/lead_enrichment")
CATALOG = ROOT / "_catalogs" / "data_cms_gov_data.json"
MANIFEST = ROOT / "_catalogs" / "aco_backfill_manifest.json"
FAMILY = re.compile(
    r"Accountable Care Organization|Performance Year Financial and Quality"
    r"|Assigned Beneficiaries by County|County-level Aggregate Expenditure.*Assignable"
    r"|Next Generation ACO|Next Generation Accountable|Pioneer ACO|REACH"
    r"|Advance Investment Payment|Prepaid Shared Savings|Kidney Care Choices"
    r"|Comprehensive ESRD|Comprehensive End.Stage|Ambulatory Specialty Model"
    r"|Primary Care First|Making Care Primary|CPC\+",
    re.I,
)
FORMATS = {"CSV", "ZIP", "XLSX"}


def source_id(base: str) -> str:
    slug = re.sub(r"[^a-z0-9]+", "_", base.lower()).strip("_")
    return f"cms_{slug}"[:74]


def sha256(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for b in iter(lambda: f.read(1 << 20), b""):
            h.update(b)
    return h.hexdigest()


def plan(catalog: dict) -> list[dict]:
    out = []
    for ds in catalog["dataset"]:
        title = ds["title"]
        if not FAMILY.search(title):
            continue
        base, _, ver = title.partition(" : ")
        if not ver:
            continue
        release = re.sub(r"[^0-9a-z]", "_", ver.strip().lower())
        for dist in ds.get("distribution", []):
            url = dist.get("downloadURL", "")
            if dist.get("format") not in FORMATS or not url:
                continue
            name = urllib.parse.unquote(Path(urllib.parse.urlparse(url).path).name)
            out.append(
                {
                    "source_id": source_id(base),
                    "series": base,
                    "release": release,
                    "url": url,
                    "name": name,
                }
            )
    return out


def fetch(url: str, dest: Path) -> int:
    part = dest.with_name(dest.name + ".part")
    have = part.stat().st_size if part.exists() else 0
    req = urllib.request.Request(url, headers={"User-Agent": UA})
    if have:
        req.add_header("Range", f"bytes={have}-")
    try:
        resp = urllib.request.urlopen(req, timeout=60)
    except urllib.error.HTTPError as e:
        if e.code == 416 and have:  # already complete
            return have
        raise
    mode = "ab" if have and resp.status == 206 else "wb"
    with part.open(mode) as f:
        while chunk := resp.read(1 << 20):
            f.write(chunk)
    return part.stat().st_size


PROGRAM = {
    "accountable_care_organization": "mssp", "performance_year_financial": "mssp",
    "number_of_accountable_care": "mssp", "county_level_aggregate": "mssp",
    "next_generation": "ngaco", "pioneer": "pioneer", "reach": "reach",
    "comprehensive_end_stage": "cec", "kidney_care": "kcc",
    "advance_investment": "mssp", "prepaid_shared": "mssp", "ambulatory_specialty": "asm",
}


def register(dry: bool) -> int:
    """One bib source item per on-disk dataset version (file attached).

    Tags as import_lead_corpus (``source:lead-corpus``,
    ``lead-dataset:<source_id>``) plus ``doc:dataset``,
    ``source:cms-website`` and ``program:<x>``. Items match on URL (the
    upsert key), so re-running is safe; dry-run reports what exists.
    """
    from collections import Counter

    from bib.item import Source
    from bib.store import Store

    store = Store()
    col = store.ensure_collections(
        {"LEAD corpus": {"Dataset documentation": {}}}
    )["Dataset documentation"] if not dry else None
    have = {r[0] for r in store._con().execute("SELECT url FROM items WHERE url <> ''")}
    before: Counter = Counter()
    total: Counter = Counter()
    for m in json.loads(MANIFEST.read_text()):
        total[m["source_id"]] += 1
        before[m["source_id"]] += m["url"] in have
        if m["url"] in have and dry:
            continue
        prog = next((v for k, v in PROGRAM.items() if k in m["source_id"]), "")
        tags = ["source:lead-corpus", "source:cms-website",
                f"lead-dataset:{m['source_id']}", "doc:dataset"]
        tags += [f"program:{prog}"] if prog else []
        if dry:
            print("would register", m["series"], m["release"])
            continue
        item = Source(
            title=f"{m['series']} : {m['release'].replace('_', '-')}", url=m["url"],
            institution="CMS / data.cms.gov",
            date_published=m["release"].replace("_", "-")[:10],
            extra=f"sha256: {m['sha256']}", doc_type="dataset",
        )
        key, _ = store.upsert_status(item, tags=tags, collection=col)
        store.attach_file(key, ROOT / m["source_id"] / m["release"] / m["name"])
    for sid in sorted(total):
        print(f"{sid}: bib before {before[sid]}/{total[sid]}")
    return 0


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--limit", type=int, default=0, help="stop after N downloads")
    ap.add_argument("--budget", type=float, default=0, help="stop starting new downloads after SECS")
    ap.add_argument("--pause", type=float, default=2.0)
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--register", action="store_true", help="register fetched files in bib (honours --dry-run)")
    ap.add_argument("--only", default="", help="substring of source_id/release to restrict to")
    args = ap.parse_args()

    if args.register:
        return register(args.dry_run)
    items = plan(json.loads(CATALOG.read_text()))
    if args.only:
        items = [i for i in items if args.only in f"{i['source_id']}/{i['release']}"]
    manifest = json.loads(MANIFEST.read_text()) if MANIFEST.exists() else []
    known = {(m["source_id"], m["release"], m["url"]) for m in manifest}
    t0, done, failed, held = time.time(), 0, [], 0
    for it in items:
        folder = ROOT / it["source_id"] / it["release"]
        dest = folder / it["name"]
        key = (it["source_id"], it["release"], it["url"])
        if dest.exists():
            held += 1
            if key not in known and not args.dry_run:
                manifest.append({**it, "sha256": sha256(dest), "bytes": dest.stat().st_size,
                                 "fetched_at": None, "status": "already_held"})
                known.add(key)
            continue
        if args.dry_run:
            print("would fetch", dest)
            continue
        if (args.limit and done >= args.limit) or (args.budget and time.time() - t0 > args.budget):
            break
        folder.mkdir(parents=True, exist_ok=True)
        try:
            size = fetch(it["url"], dest)
        except Exception as e:  # noqa: BLE001 — record and move on
            failed.append((str(dest), repr(e)))
            print("FAIL", dest, e, file=sys.stderr)
            time.sleep(args.pause)
            continue
        part = dest.with_name(dest.name + ".part")
        digest = sha256(part)
        part.rename(dest)  # dest did not exist: never an overwrite
        manifest.append({**it, "sha256": digest, "bytes": size,
                         "fetched_at": dt.datetime.now(dt.UTC).isoformat(timespec="seconds"),
                         "status": "fetched"})
        MANIFEST.write_text(json.dumps(manifest, indent=1))
        done += 1
        print(f"ok {size:>12,} {dest}")
        time.sleep(args.pause)
    if not args.dry_run:
        MANIFEST.write_text(json.dumps(manifest, indent=1))
    print(f"fetched {done}, already held {held}, failed {len(failed)}, planned {len(items)}")
    return 1 if failed else 0


if __name__ == "__main__":
    raise SystemExit(main())