112 lines
3.4 KiB
Python
112 lines
3.4 KiB
Python
"""M4 pilot (#513): publish OPPS reference tables to the DuckLake lakehouse.
|
|
|
|
Reads the OPPS tables from the local DuckDB (read replica when present)
|
|
and writes them to DuckLake via ``aco.lake.DuckLakeContext.save`` — the
|
|
Context write path this pilot exists to exercise. Config comes from
|
|
``stack.toml [lake.ducklake]``; the postgres catalog password from the
|
|
``POSTGRES_PASSWORD`` env var.
|
|
|
|
The catalog host and RustFS are compose-internal, so this runs inside a
|
|
data-network container:
|
|
|
|
docker exec -e POSTGRES_PASSWORD=... notebooks \\
|
|
env PYTHONPATH=/home/kert/src uv run --project /home/kert/workspace \\
|
|
python /tmp/publish_opps_to_lake.py
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import time
|
|
|
|
SCHEMAS: dict[str, tuple[str, ...]] = {
|
|
"opps": ("addendum_b", "apc_weight", "skin_sub_addendum_b"),
|
|
"pfs": (
|
|
"carrier_locality",
|
|
"clinical_labor",
|
|
"gpci",
|
|
"gpci_proposed",
|
|
"medical_equipment",
|
|
"medical_supply",
|
|
"physician_work_time",
|
|
"rvu",
|
|
"rvu_proposed",
|
|
"zip_carrier_locality",
|
|
),
|
|
"cms": ("ingest_log",),
|
|
}
|
|
|
|
# (monolith source, lake destination) — tables whose lake name differs
|
|
# from their monolith name. The MSSP designated-primary-care-service
|
|
# list (42 CFR 425.400(c), used for ACO beneficiary assignment) lives
|
|
# monolith-side as a tuva/dbt intermediate; publish it under a stable
|
|
# reference name the notebooks can cite (#630). Published whenever the
|
|
# destination's schema is selected via --schemas.
|
|
RENAMED: tuple[tuple[str, str], ...] = (
|
|
(
|
|
"cms_provider_attribution._primary_care_hcpcs_codes",
|
|
"cms.primary_care_service_code",
|
|
),
|
|
)
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument(
|
|
"--schemas",
|
|
nargs="*",
|
|
default=list(SCHEMAS),
|
|
choices=list(SCHEMAS),
|
|
help="reference schemas to publish",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
import narwhals as nw
|
|
|
|
from aco.lake import DuckLakeContext
|
|
from conf import cfg, connect
|
|
|
|
pw = os.environ.get("POSTGRES_PASSWORD", "")
|
|
dsn = cfg.lake.ducklake.catalog
|
|
if dsn.startswith("postgres:") and pw:
|
|
dsn = f"{dsn} password={pw}"
|
|
|
|
ctx = DuckLakeContext(
|
|
catalog_dsn=dsn,
|
|
data_path=cfg.lake.ducklake.data_path,
|
|
s3_endpoint=cfg.lake.ducklake.s3_endpoint,
|
|
read_only=False,
|
|
)
|
|
|
|
src = connect.duckdb("aco") # read-only; resolves to the replica
|
|
|
|
def _publish(source: str, dest: str) -> bool:
|
|
t = time.time()
|
|
df = src.execute(f"SELECT * FROM {source}").pl() # noqa: S608
|
|
ctx.save(dest, nw.from_native(df), mode="replace")
|
|
back = nw.to_native(ctx.load(dest))
|
|
ok = back.height == df.height
|
|
label = source if source == dest else f"{source} → {dest}"
|
|
print(
|
|
f" {label}: {df.height} rows → lake "
|
|
f"({back.height} read back) [{'OK' if ok else 'MISMATCH'}] "
|
|
f"in {time.time() - t:.1f}s"
|
|
)
|
|
return ok
|
|
|
|
for schema in args.schemas:
|
|
for table in SCHEMAS[schema]:
|
|
if not _publish(f"{schema}.{table}", f"{schema}.{table}"):
|
|
return 1
|
|
for source, dest in RENAMED:
|
|
if dest.split(".", 1)[0] in args.schemas and not _publish(source, dest):
|
|
return 1
|
|
src.close()
|
|
print("lake publish complete")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|