d08401dc53
Client defaults to /api/v3; _typeFilter on start_export/export with type_filter_tag() for the System-Type tag; jobs() sends the bulk-data headers BCDA requires (Prefer: respond-async, Accept: application/fhir+json — the old call 400'd on v2 and v3) and unwraps the Task Bundle; completed_jobs() lists reusable exports; every non-2xx raises BcdaHttpError carrying the OperationOutcome diagnostics (still an httpx.HTTPStatusError); a poll that hits 404 'Job has been cancelled' raises JobCancelledError; /_version and /_health. Store.adopt(client, job_url) records and downloads a job this store did not start (completed jobs stay downloadable 24 h and the public sandbox entity is shared); export(cancel_retries=N) restarts after server-side cancellation, recording each 'cancelled' transition. bcda.express.v3 flattens C4BB 2.2.0 resources into the existing bcda.* tables, table-driven from the v3 Data Dictionary: v3 moved every BlueButton code off resources/variables/<name> onto fhir/StructureDefinition|CodeSystem/<CODE> and consolidated amounts into adjudication/total/item.adjudication, so the v2 suffix match returned NULL for almost everything. Dispatch is per resource (is_v3); v2 payloads are unchanged. Identity follows the migration guide: bene_id = Patient.id (BENE_SK), current MBI = us-mbi identifier without period.end, clm_id = C4BB 'uc', clm_cntl_num pairs SharedSystem with NCH. Race crosswalked from US Core OMB codes; Coverage is one row per plan+period (Part A/B/C/D + Dual) with v3-native columns; EOB gains System-Type / Final-Action / EOB-TYPE, IDR dates, header totals and PPS/IME/DSH amounts; items gain NDC quantity, revenue-centre indicators, ANSI codes, Part D/DME amounts. flatten_export streams <Type>.N.ndjson parts to Parquet with model-derived Arrow schemas (a Large-entity EOB export is 192 files / 2.2 GB); load_bcda() gathers the newest export dir per resource type (a v3 pull is three single-type jobs). cclf.stg_beneficiary_demographics prefers the Part A row and joins the Dual row; eligibility reads coalesce(dual_01, dual_stus_cd). bcda package at 100% line coverage; tests run on CMS's official v3 samples. refs #1139
657 lines
22 KiB
Python
657 lines
22 KiB
Python
"""Persistent storage and state management for BCDA exports.
|
|
|
|
Wraps :class:`bcda.client.Client` with append-only Parquet state
|
|
tracking and configurable file storage (local, network, S3, Azure).
|
|
|
|
State is kept in two ever-growing Parquet files — ``jobs.parquet``
|
|
and ``files.parquet`` — where each mutation appends a new row.
|
|
The current state of any job or file is the **last row** for its
|
|
key, giving a full audit trail for free.
|
|
|
|
Storage backends are selected by the ``path`` prefix:
|
|
|
|
- Local / network mount: ``"data/bcda"`` or ``"/mnt/share/bcda"``
|
|
- S3: ``"s3://bucket/prefix"``
|
|
- Azure Blob: ``"az://container/prefix"``
|
|
|
|
All backends use ``fsspec`` under the hood. S3 and Azure require
|
|
the appropriate extra (``s3fs``, ``adlfs``) installed by the user.
|
|
|
|
Usage::
|
|
|
|
from bcda.client import Client, JobCancelledError
|
|
from bcda.store import Store
|
|
|
|
store = Store()
|
|
|
|
with Client(client_id="...", client_secret="...") as c:
|
|
key = store.export(c, types=["Patient", "Coverage"])
|
|
|
|
# Resume after interruption
|
|
with Client(client_id="...", client_secret="...") as c:
|
|
store.resume(c, key)
|
|
|
|
# Inspect state
|
|
store.get_job(key)
|
|
store.get_files(key)
|
|
store.list_jobs(status="failed")
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import random
|
|
import string
|
|
import tempfile
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import fsspec
|
|
import pyarrow as pa
|
|
import pyarrow.parquet as pq
|
|
|
|
from bcda.client import Client, JobCancelledError
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
# ── Parquet schemas ──────────────────────────────────────────
|
|
|
|
JOBS_SCHEMA = pa.schema(
|
|
[
|
|
("ts", pa.timestamp("us", tz="UTC")),
|
|
("key", pa.string()),
|
|
("job_url", pa.string()),
|
|
("endpoint", pa.string()),
|
|
("types", pa.string()),
|
|
("since", pa.string()),
|
|
("status", pa.string()),
|
|
("error", pa.string()),
|
|
]
|
|
)
|
|
|
|
FILES_SCHEMA = pa.schema(
|
|
[
|
|
("ts", pa.timestamp("us", tz="UTC")),
|
|
("job_key", pa.string()),
|
|
("resource_type", pa.string()),
|
|
("url", pa.string()),
|
|
("storage_path", pa.string()),
|
|
("file_size", pa.int64()),
|
|
("status", pa.string()),
|
|
("error", pa.string()),
|
|
]
|
|
)
|
|
|
|
|
|
def _now() -> datetime:
|
|
return datetime.now(tz=timezone.utc)
|
|
|
|
|
|
def _generate_key() -> str:
|
|
chars = string.ascii_uppercase + string.digits
|
|
return "".join(random.choices(chars, k=8)) # noqa: S311
|
|
|
|
|
|
def _current_state(table: pa.Table, key_col: str) -> pa.Table:
|
|
"""Deduplicate an append-only table, keeping the last row per key."""
|
|
if table.num_rows == 0:
|
|
return table
|
|
# Sort by ts ascending so the last occurrence per key wins.
|
|
indices = pa.compute.sort_indices(table, sort_keys=[("ts", "ascending")])
|
|
table = table.take(indices)
|
|
# Walk backwards, keeping only the first occurrence of each key.
|
|
keys = table.column(key_col)
|
|
seen: set[str] = set()
|
|
keep: list[int] = []
|
|
for i in range(table.num_rows - 1, -1, -1):
|
|
k = keys[i].as_py()
|
|
if k not in seen:
|
|
seen.add(k)
|
|
keep.append(i)
|
|
keep.reverse()
|
|
return table.take(keep)
|
|
|
|
|
|
class Store:
|
|
"""Persistent BCDA export storage and state manager.
|
|
|
|
Parameters
|
|
----------
|
|
path : str or Path
|
|
Root location for state files, logs, and ndjson exports.
|
|
Accepts local paths, network mounts, or URIs for S3
|
|
(``s3://...``) and Azure (``az://...``).
|
|
storage_options : dict
|
|
Extra arguments forwarded to ``fsspec.filesystem()`` —
|
|
typically credentials for cloud backends.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
path: str | Path | None = None,
|
|
storage_options: dict[str, Any] | None = None,
|
|
) -> None:
|
|
if path is None:
|
|
from conf import path as _path
|
|
|
|
path = str(_path("storage.bcda"))
|
|
self._root = str(path).rstrip("/")
|
|
self._storage_options = storage_options or {}
|
|
self._filesystem: fsspec.AbstractFileSystem | None = None
|
|
|
|
# ── Filesystem ───────────────────────────────────────────
|
|
|
|
def _fs(self) -> fsspec.AbstractFileSystem:
|
|
"""Return (and cache) the fsspec filesystem for ``self._root``."""
|
|
if self._filesystem is None:
|
|
self._filesystem, _ = fsspec.core.url_to_fs(
|
|
self._root,
|
|
**self._storage_options,
|
|
)
|
|
return self._filesystem
|
|
|
|
def _full(self, *parts: str) -> str:
|
|
"""Join parts onto the root path."""
|
|
return "/".join([self._root, *parts])
|
|
|
|
def _ensure_dir(self, path: str) -> None:
|
|
self._fs().mkdirs(path, exist_ok=True)
|
|
|
|
# ── Parquet state I/O ────────────────────────────────────
|
|
|
|
@property
|
|
def _jobs_path(self) -> str:
|
|
return self._full("jobs.parquet")
|
|
|
|
@property
|
|
def _files_path(self) -> str:
|
|
return self._full("files.parquet")
|
|
|
|
def _read_parquet(self, path: str, schema: pa.Schema) -> pa.Table:
|
|
"""Read a parquet file; return empty table if it doesn't exist."""
|
|
fs = self._fs()
|
|
if not fs.exists(path):
|
|
return schema.empty_table()
|
|
with fs.open(path, "rb") as f:
|
|
return pq.read_table(f, schema=schema)
|
|
|
|
def _append_parquet(
|
|
self,
|
|
path: str,
|
|
schema: pa.Schema,
|
|
rows: list[dict],
|
|
) -> None:
|
|
"""Append rows to an ever-growing parquet file."""
|
|
existing = self._read_parquet(path, schema)
|
|
new = pa.Table.from_pylist(rows, schema=schema)
|
|
combined = pa.concat_tables([existing, new])
|
|
self._ensure_dir(self._root)
|
|
fs = self._fs()
|
|
with fs.open(path, "wb") as f:
|
|
pq.write_table(combined, f)
|
|
|
|
def _append_job(self, **fields: Any) -> None:
|
|
fields.setdefault("ts", _now())
|
|
for col in JOBS_SCHEMA.names:
|
|
fields.setdefault(col, "")
|
|
self._append_parquet(self._jobs_path, JOBS_SCHEMA, [fields])
|
|
|
|
def _append_file(self, **fields: Any) -> None:
|
|
fields.setdefault("ts", _now())
|
|
fields.setdefault("file_size", 0)
|
|
for col in FILES_SCHEMA.names:
|
|
fields.setdefault(col, "")
|
|
self._append_parquet(self._files_path, FILES_SCHEMA, [fields])
|
|
|
|
# ── Query state ──────────────────────────────────────────
|
|
|
|
def get_job(self, key: str) -> dict:
|
|
"""Get the current state of a job by its key.
|
|
|
|
Raises ``KeyError`` if no job with that key exists.
|
|
"""
|
|
table = self._read_parquet(self._jobs_path, JOBS_SCHEMA)
|
|
current = _current_state(table, "key")
|
|
for row in current.to_pylist():
|
|
if row["key"] == key:
|
|
return row
|
|
raise KeyError(f"Job not found: {key}")
|
|
|
|
def list_jobs(self, *, status: str | None = None) -> list[dict]:
|
|
"""List all jobs, optionally filtered by status."""
|
|
table = self._read_parquet(self._jobs_path, JOBS_SCHEMA)
|
|
current = _current_state(table, "key")
|
|
rows = current.to_pylist()
|
|
if status:
|
|
rows = [r for r in rows if r["status"] == status]
|
|
return rows
|
|
|
|
def get_files(self, job_key: str) -> list[dict]:
|
|
"""Get current state of all files for a job."""
|
|
table = self._read_parquet(self._files_path, FILES_SCHEMA)
|
|
current = _current_state(table, "url")
|
|
return [r for r in current.to_pylist() if r["job_key"] == job_key]
|
|
|
|
# ── File storage ─────────────────────────────────────────
|
|
|
|
def _store_file(
|
|
self,
|
|
local_path: Path,
|
|
job_key: str,
|
|
filename: str,
|
|
) -> str:
|
|
"""Copy a local file into the storage backend.
|
|
|
|
Returns the storage path.
|
|
"""
|
|
dest = self._full("exports", job_key, filename)
|
|
self._ensure_dir(self._full("exports", job_key))
|
|
fs = self._fs()
|
|
fs.put_file(str(local_path), dest)
|
|
return dest
|
|
|
|
def open(self, storage_path: str, mode: str = "rb"):
|
|
"""Open a file from the storage backend."""
|
|
return self._fs().open(storage_path, mode)
|
|
|
|
# ── Export with state tracking ───────────────────────────
|
|
|
|
def export(
|
|
self,
|
|
client: Client,
|
|
*,
|
|
endpoint: str = "Group/all",
|
|
types: list[str] | None = None,
|
|
since: str | None = None,
|
|
type_filter: str | None = None,
|
|
poll_interval: float = 10.0,
|
|
cancel_retries: int = 0,
|
|
) -> str:
|
|
"""Run a full export with persistent state tracking.
|
|
|
|
Wraps :meth:`Client.start_export`, :meth:`Client.poll`,
|
|
and file downloads, recording every state transition to
|
|
Parquet. Returns the job key.
|
|
|
|
Parameters
|
|
----------
|
|
client : Client
|
|
An authenticated BCDA client instance.
|
|
endpoint : str
|
|
``"Group/all"`` (default), ``"Group/runout"``, or
|
|
``"Patient"``.
|
|
types : list[str] | None
|
|
FHIR resource types to request.
|
|
since : str | None
|
|
``_since`` timestamp filter.
|
|
type_filter : str | None
|
|
v3 ``_typeFilter`` subquery (recorded alongside ``types``).
|
|
poll_interval : float
|
|
Seconds between status polls.
|
|
cancel_retries : int
|
|
How many times to restart the export after the server
|
|
reports it cancelled. On the public sandbox every caller
|
|
of a model entity can ``DELETE`` any job, and sweeps that
|
|
cancel every in-progress job happen several times an hour;
|
|
each cancellation is recorded as a ``cancelled`` transition
|
|
under the same key before the next attempt.
|
|
"""
|
|
key = _generate_key()
|
|
types_str = ",".join(types) if types else ""
|
|
if type_filter:
|
|
types_str = f"{types_str}|_typeFilter={type_filter}"
|
|
log.info("Starting export %s", key, extra={"job_key": key})
|
|
|
|
# 1. Record pending.
|
|
self._append_job(
|
|
key=key,
|
|
endpoint=endpoint,
|
|
types=types_str,
|
|
since=since or "",
|
|
status="pending",
|
|
)
|
|
|
|
try:
|
|
for attempt in range(cancel_retries + 1):
|
|
# 2. Start export.
|
|
job_url = client.start_export(
|
|
endpoint=endpoint,
|
|
types=types,
|
|
since=since,
|
|
type_filter=type_filter,
|
|
)
|
|
self._append_job(
|
|
key=key,
|
|
job_url=job_url,
|
|
endpoint=endpoint,
|
|
types=types_str,
|
|
since=since or "",
|
|
status="polling",
|
|
)
|
|
|
|
# 3. Poll until complete.
|
|
try:
|
|
result = client.poll(job_url, interval=poll_interval)
|
|
break
|
|
except JobCancelledError as exc:
|
|
self._append_job(
|
|
key=key,
|
|
job_url=job_url,
|
|
endpoint=endpoint,
|
|
types=types_str,
|
|
since=since or "",
|
|
status="cancelled",
|
|
error=str(exc),
|
|
)
|
|
if attempt >= cancel_retries:
|
|
raise
|
|
log.warning(
|
|
"Export %s cancelled server-side (%s); restarting (%d/%d)",
|
|
key,
|
|
job_url,
|
|
attempt + 1,
|
|
cancel_retries,
|
|
extra={"job_key": key},
|
|
)
|
|
self._append_job(
|
|
key=key,
|
|
job_url=job_url,
|
|
endpoint=endpoint,
|
|
types=types_str,
|
|
since=since or "",
|
|
status="complete",
|
|
)
|
|
|
|
# 4. Download files.
|
|
self._download_outputs(client, key, result)
|
|
|
|
except Exception as exc:
|
|
log.error(
|
|
"Export %s failed: %s",
|
|
key,
|
|
exc,
|
|
extra={"job_key": key},
|
|
)
|
|
self._append_job(
|
|
key=key,
|
|
endpoint=endpoint,
|
|
types=types_str,
|
|
since=since or "",
|
|
status="failed",
|
|
error=str(exc),
|
|
)
|
|
raise
|
|
|
|
log.info("Export %s finished", key, extra={"job_key": key})
|
|
return key
|
|
|
|
def adopt(
|
|
self,
|
|
client: Client,
|
|
job_url: str,
|
|
*,
|
|
endpoint: str = "Group/all",
|
|
types: list[str] | None = None,
|
|
poll_interval: float = 10.0,
|
|
) -> str:
|
|
"""Record and download a job this store did not start.
|
|
|
|
BCDA keeps a completed job's files for 24 hours, and on the
|
|
public sandbox every caller of a model entity sees the same
|
|
job list — so a fresh export (6 min and 2 GB for the Large
|
|
entity's EOBs) is wasted when ``client.completed_jobs()``
|
|
already lists one. ``adopt`` polls *job_url* (returns at once
|
|
when complete), appends the job under a new key with status
|
|
``complete`` and downloads its outputs exactly as
|
|
:meth:`export` would. Returns the job key.
|
|
"""
|
|
key = _generate_key()
|
|
types_str = ",".join(types) if types else ""
|
|
log.info("Adopting job %s as %s", job_url, key, extra={"job_key": key})
|
|
self._append_job(
|
|
key=key,
|
|
job_url=job_url,
|
|
endpoint=endpoint,
|
|
types=types_str,
|
|
status="polling",
|
|
)
|
|
try:
|
|
result = client.poll(job_url, interval=poll_interval)
|
|
self._append_job(
|
|
key=key,
|
|
job_url=job_url,
|
|
endpoint=endpoint,
|
|
types=types_str,
|
|
status="complete",
|
|
)
|
|
self._download_outputs(client, key, result)
|
|
except Exception as exc:
|
|
log.error("Adopt %s failed: %s", key, exc, extra={"job_key": key})
|
|
self._append_job(
|
|
key=key,
|
|
job_url=job_url,
|
|
endpoint=endpoint,
|
|
types=types_str,
|
|
status="failed",
|
|
error=str(exc),
|
|
)
|
|
raise
|
|
log.info("Adopted %s finished", key, extra={"job_key": key})
|
|
return key
|
|
|
|
def _download_outputs(
|
|
self,
|
|
client: Client,
|
|
job_key: str,
|
|
result: dict,
|
|
) -> None:
|
|
"""Download all output files from a completed job result."""
|
|
outputs = result.get("output", [])
|
|
type_counts: dict[str, int] = {}
|
|
|
|
for entry in outputs:
|
|
url = entry["url"]
|
|
resource_type = entry.get("type", "unknown")
|
|
|
|
count = type_counts.get(resource_type, 0)
|
|
type_counts[resource_type] = count + 1
|
|
if count == 0:
|
|
filename = f"{resource_type}.ndjson"
|
|
else:
|
|
filename = f"{resource_type}.{count}.ndjson"
|
|
|
|
# Record pending.
|
|
self._append_file(
|
|
job_key=job_key,
|
|
resource_type=resource_type,
|
|
url=url,
|
|
status="pending",
|
|
)
|
|
|
|
try:
|
|
self._append_file(
|
|
job_key=job_key,
|
|
resource_type=resource_type,
|
|
url=url,
|
|
status="downloading",
|
|
)
|
|
|
|
# Download to temp, then move to storage.
|
|
with tempfile.NamedTemporaryFile(
|
|
suffix=".ndjson",
|
|
delete=False,
|
|
) as tmp:
|
|
tmp_path = Path(tmp.name)
|
|
|
|
client._download_file(url, tmp_path)
|
|
file_size = tmp_path.stat().st_size
|
|
storage_path = self._store_file(
|
|
tmp_path,
|
|
job_key,
|
|
filename,
|
|
)
|
|
tmp_path.unlink(missing_ok=True)
|
|
|
|
self._append_file(
|
|
job_key=job_key,
|
|
resource_type=resource_type,
|
|
url=url,
|
|
storage_path=storage_path,
|
|
file_size=file_size,
|
|
status="complete",
|
|
)
|
|
log.info(
|
|
"Stored %s (%d bytes)",
|
|
filename,
|
|
file_size,
|
|
extra={
|
|
"job_key": job_key,
|
|
"resource_type": resource_type,
|
|
"file_size": file_size,
|
|
},
|
|
)
|
|
|
|
except Exception as exc:
|
|
tmp_path.unlink(missing_ok=True)
|
|
self._append_file(
|
|
job_key=job_key,
|
|
resource_type=resource_type,
|
|
url=url,
|
|
status="failed",
|
|
error=str(exc),
|
|
)
|
|
log.error(
|
|
"Download failed for %s: %s",
|
|
resource_type,
|
|
exc,
|
|
extra={
|
|
"job_key": job_key,
|
|
"resource_type": resource_type,
|
|
},
|
|
)
|
|
raise
|
|
|
|
# ── Resume / retry ───────────────────────────────────────
|
|
|
|
def resume(
|
|
self,
|
|
client: Client,
|
|
key: str,
|
|
*,
|
|
poll_interval: float = 10.0,
|
|
) -> str:
|
|
"""Resume an incomplete export from its last known state.
|
|
|
|
Picks up where the process left off:
|
|
|
|
- ``pending`` → start the export from scratch.
|
|
- ``polling`` → re-poll the saved ``job_url``.
|
|
- ``complete`` → download any files not yet complete.
|
|
- ``failed`` → retry from the failure point.
|
|
|
|
Parameters
|
|
----------
|
|
client : Client
|
|
An authenticated BCDA client instance.
|
|
key : str
|
|
The job key returned by :meth:`export`.
|
|
poll_interval : float
|
|
Seconds between status polls.
|
|
"""
|
|
job = self.get_job(key)
|
|
status = job["status"]
|
|
log.info(
|
|
"Resuming export %s (status=%s)",
|
|
key,
|
|
status,
|
|
extra={"job_key": key},
|
|
)
|
|
|
|
if status == "pending":
|
|
# Never got as far as start_export; re-run.
|
|
types = job["types"].split(",") if job["types"] else None
|
|
since = job["since"] or None
|
|
job_url = client.start_export(
|
|
endpoint=job["endpoint"],
|
|
types=types,
|
|
since=since,
|
|
)
|
|
self._append_job(
|
|
key=key,
|
|
job_url=job_url,
|
|
endpoint=job["endpoint"],
|
|
types=job["types"],
|
|
since=job["since"],
|
|
status="polling",
|
|
)
|
|
job["job_url"] = job_url
|
|
status = "polling"
|
|
|
|
if status in ("polling", "failed"):
|
|
job_url = job.get("job_url") or job.get("job_url", "")
|
|
if not job_url:
|
|
raise ValueError(f"Job {key} has no job_url — cannot resume polling.")
|
|
try:
|
|
result = client.poll(job_url, interval=poll_interval)
|
|
self._append_job(
|
|
key=key,
|
|
job_url=job_url,
|
|
endpoint=job["endpoint"],
|
|
types=job["types"],
|
|
since=job["since"],
|
|
status="complete",
|
|
)
|
|
status = "complete"
|
|
except Exception as exc:
|
|
self._append_job(
|
|
key=key,
|
|
job_url=job_url,
|
|
endpoint=job["endpoint"],
|
|
types=job["types"],
|
|
since=job["since"],
|
|
status="failed",
|
|
error=str(exc),
|
|
)
|
|
raise
|
|
|
|
if status == "complete":
|
|
# Check which files still need downloading.
|
|
files = self.get_files(key)
|
|
completed_urls = {f["url"] for f in files if f["status"] == "complete"}
|
|
|
|
# We need the job result to know what outputs exist.
|
|
# If we just polled, we have it; otherwise re-poll
|
|
# (a completed job returns 200 immediately).
|
|
if "result" not in dir():
|
|
result = client.poll(
|
|
job["job_url"],
|
|
interval=poll_interval,
|
|
)
|
|
|
|
pending = [
|
|
e for e in result.get("output", []) if e["url"] not in completed_urls
|
|
]
|
|
|
|
if pending:
|
|
log.info(
|
|
"Resuming %d pending file(s) for %s",
|
|
len(pending),
|
|
key,
|
|
extra={"job_key": key},
|
|
)
|
|
# Build a partial result for _download_outputs.
|
|
self._download_outputs(
|
|
client,
|
|
key,
|
|
{"output": pending},
|
|
)
|
|
else:
|
|
log.info(
|
|
"All files already complete for %s",
|
|
key,
|
|
extra={"job_key": key},
|
|
)
|
|
|
|
return key
|