From 96b74c545cd20f5ba60d9506ff9629ac2212d0df Mon Sep 17 00:00:00 2001 From: kert Date: Thu, 12 Nov 2026 22:56:31 -0500 Subject: [PATCH] feat(hoard): backfill every published ACO/model dataset version from the data.cms.gov catalog --- dev/scripts/hoard_aco_versions.py | 217 ++++++++++++++++++++++++++++++ 1 file changed, 217 insertions(+) create mode 100644 dev/scripts/hoard_aco_versions.py diff --git a/dev/scripts/hoard_aco_versions.py b/dev/scripts/hoard_aco_versions.py new file mode 100644 index 0000000..923bd54 --- /dev/null +++ b/dev/scripts/hoard_aco_versions.py @@ -0,0 +1,217 @@ +"""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///`` +(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 ``.``. +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", +} + + +def register(dry: bool) -> int: + """One bib source item per newly fetched dataset version (file attached). + + Same tags as import_lead_corpus (``source:lead-corpus`` + + ``lead-dataset:``) plus ``doc:dataset`` and ``program:``. + Items match on URL, so re-running is safe. + """ + from collections import Counter + + manifest = json.loads(MANIFEST.read_text()) + rows = [m for m in manifest if m["status"] == "fetched"] + stats: Counter = Counter() + if not dry: + from bib.item import Source + from bib.store import Store + + store = Store() + col = store.ensure_collections( + {"LEAD corpus": {"Dataset documentation": {}}} + )["Dataset documentation"] + for m in rows: + prog = next((v for k, v in PROGRAM.items() if k in m["source_id"]), "") + tags = ["source:lead-corpus", f"lead-dataset:{m['source_id']}", "doc:dataset"] + tags += [f"program:{prog}"] if prog else [] + title = f"{m['series']} : {m['release'].replace('_', '-')}" + if dry: + print("would register", title, tags) + stats["dry"] += 1 + continue + item = Source( + title=title, url=m["url"], institution="CMS / data.cms.gov", + date_published=m["release"].replace("_", "-")[:10], + extra=f"sha256: {m['sha256']}", doc_type="dataset", + ) + key, status = store.upsert_status(item, tags=tags, collection=col) + store.attach_file(key, ROOT / m["source_id"] / m["release"] / m["name"]) + stats[status] += 1 + print("register:", dict(stats)) + 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())