stack/dev/scripts/publish_reference_to_lake.py

"""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",
        "zip_carrier_locality",
    ),
}


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())