stack/src/aco/load/bcda.py

"""Load BCDA FHIR exports into DuckDB.

Flattens raw ndjson files (Patient, ExplanationOfBenefit, Coverage) into
Parquet via the existing ``bcda.express.flatten`` module, then loads the
resulting Parquet files into the ``bcda`` schema in DuckDB.

Usage::

    from aco.load.bcda import load_bcda

    stats = load_bcda()
    # {"bcda.patient": 50, "bcda.explanation_of_benefit": 2021, ...}
"""

from __future__ import annotations

from pathlib import Path


def load_bcda(
    ndjson_dir: str | Path | None = None,
    database: str | None = None,
    *,
    skip_flatten: bool = False,
) -> dict[str, int]:
    """Flatten BCDA ndjson and load into DuckDB.

    Parameters
    ----------
    ndjson_dir : str or Path, optional
        Directory containing raw ndjson files.  Defaults to the latest
        export directory under ``conf.path("storage.bcda")``.
    database : str, optional
        DuckDB database path.  Defaults to ``conf.path("db.aco")``.
    skip_flatten : bool
        If True, skip the flatten step and load existing Parquet files.

    Returns
    -------
    dict[str, int]
        Row counts per table loaded.
    """
    import duckdb

    from conf import path

    db_path = database or str(path("db.aco"))
    store_path = path("storage.bcda")

    if ndjson_dir is None:  # pragma: no cover — tests always provide explicit dir
        ndjson_dir = _find_latest_export(store_path)

    if not skip_flatten:
        from bcda.express.flatten import flatten_export

        flatten_export(ndjson_dir, store_path)

    # Load Parquet files into DuckDB
    flat_dir = store_path / "flat"
    if not flat_dir.exists():
        raise FileNotFoundError(f"No flattened data at {flat_dir}")

    con = duckdb.connect(str(db_path), read_only=False)
    con.execute('CREATE SCHEMA IF NOT EXISTS "bcda"')

    stats: dict[str, int] = {}
    parquet_map = {
        "patient": "bcda.patient",
        "explanation_of_benefit": "bcda.explanation_of_benefit",
        "eob_item": "bcda.eob_item",
        "coverage": "bcda.coverage",
    }

    for file_stem, table_ref in parquet_map.items():
        pq_path = flat_dir / f"{file_stem}.parquet"
        if not pq_path.exists():
            continue

        schema, table = table_ref.split(".")
        qualified = f'"{schema}"."{table}"'
        con.execute(f"DROP TABLE IF EXISTS {qualified}")
        con.execute(
            f"CREATE TABLE {qualified} AS SELECT * FROM read_parquet('{pq_path}')"
        )
        count = con.execute(f"SELECT COUNT(*) FROM {qualified}").fetchone()[0]
        stats[table_ref] = count

    con.close()
    return stats


def _find_latest_export(store_path: Path) -> Path:
    """Find the most recent export directory under the BCDA store."""
    exports_dir = store_path / "exports"
    if not exports_dir.exists():
        raise FileNotFoundError(
            f"No exports directory at {exports_dir}. "
            "Run a BCDA export first or specify --ndjson-dir."
        )
    # Export directories are named by job ID; pick the most recent
    dirs = sorted(
        [d for d in exports_dir.iterdir() if d.is_dir()],
        key=lambda d: d.stat().st_mtime,
        reverse=True,
    )
    if not dirs:
        raise FileNotFoundError(f"No export directories in {exports_dir}")
    return dirs[0]