89 lines
2.6 KiB
Python
89 lines
2.6 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",
|
|
"medical_equipment",
|
|
"medical_supply",
|
|
"physician_work_time",
|
|
"rvu",
|
|
"rvu_proposed",
|
|
"zip_carrier_locality",
|
|
),
|
|
"cms": ("ingest_log",),
|
|
}
|
|
|
|
|
|
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
|
|
for schema in args.schemas:
|
|
for table in SCHEMAS[schema]:
|
|
t = time.time()
|
|
df = src.execute(f"SELECT * FROM {schema}.{table}").pl() # noqa: S608
|
|
ctx.save(f"{schema}.{table}", nw.from_native(df), mode="replace")
|
|
back = nw.to_native(ctx.load(f"{schema}.{table}"))
|
|
status = "OK" if back.height == df.height else "MISMATCH"
|
|
print(
|
|
f" {schema}.{table}: {df.height} rows → lake "
|
|
f"({back.height} read back) [{status}] in {time.time() - t:.1f}s"
|
|
)
|
|
if status != "OK":
|
|
return 1
|
|
src.close()
|
|
print("lake publish complete")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|