218 lines
8.3 KiB
Python
218 lines
8.3 KiB
Python
"""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",
|
|
}
|
|
|
|
|
|
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:<source_id>``) plus ``doc:dataset`` and ``program:<x>``.
|
|
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())
|