fix(ingest): M1 duckdb concurrency — lock preflight + per-year merge (closes #508, #509)
All checks were successful
CI / lint (push) Successful in 28s
CI / notebooks-smoke (push) Successful in 1m24s
Deploy / notebooks (push) Has been skipped
Deploy / zotero (push) Has been skipped
Deploy / docs (push) Has been skipped
Deploy / api (push) Has been skipped
Deploy / mc (push) Has been skipped
Infra CI / notebooks (push) Successful in 45s
Infra CI / zotero (push) Successful in 13s
Infra CI / docs (push) Successful in 1m11s
Infra CI / api (push) Successful in 1m1s
Infra CI / mc (push) Successful in 13s
Deploy / report (push) Successful in 12s
CI / test (push) Successful in 13m12s
All checks were successful
CI / lint (push) Successful in 28s
CI / notebooks-smoke (push) Successful in 1m24s
Deploy / notebooks (push) Has been skipped
Deploy / zotero (push) Has been skipped
Deploy / docs (push) Has been skipped
Deploy / api (push) Has been skipped
Deploy / mc (push) Has been skipped
Infra CI / notebooks (push) Successful in 45s
Infra CI / zotero (push) Successful in 13s
Infra CI / docs (push) Successful in 1m11s
Infra CI / api (push) Successful in 1m1s
Infra CI / mc (push) Successful in 13s
Deploy / report (push) Successful in 12s
CI / test (push) Successful in 13m12s
#508: conf.connect.duckdb_batch() — short-lived RW connection with exponential-backoff retry on the single-writer lock; on exhaustion raises naming the holder PIDs/commands via lsof ('close the notebook tab or stop that PID') instead of a raw IOException. Wired into ingest_opps.py and ingest_asp.py. Connection-hygiene doc at docs/docs/duckdb-concurrency.md. #509: ingest_opps.py --year no longer wipes other years. The load step now merges: existing other-year rows are read back and unioned with the new frames before the rebuild (pandas concat also unions columns, so CMS format drift stays handled). skin_sub_addendum_b is derived from the merged addendum_b table instead of this run's frames. Validated against the real store: --year 2026 run preserved identical per-year counts for 2014-2025 across apc_weight (9 years), addendum_b and skin_sub_addendum_b (13 years) while rebuilding 2026.
This commit is contained in:
@@ -58,7 +58,6 @@ import re
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
|
||||
import duckdb
|
||||
import pandas as pd
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -329,28 +328,31 @@ def main() -> None:
|
||||
|
||||
# Load into DuckDB
|
||||
print(f"\nLoading into DuckDB at {DUCKDB_PATH} ...")
|
||||
con = duckdb.connect(str(DUCKDB_PATH))
|
||||
con.execute("CREATE SCHEMA IF NOT EXISTS skin_subs")
|
||||
con.execute("DROP TABLE IF EXISTS skin_subs.asp_quarterly")
|
||||
if not skin_df.empty:
|
||||
con.execute(
|
||||
"""
|
||||
CREATE TABLE skin_subs.asp_quarterly AS
|
||||
SELECT * FROM read_csv_auto(?, header=true)
|
||||
""",
|
||||
[str(OUTPUT_CSV)],
|
||||
)
|
||||
count = con.execute("SELECT count(*) FROM skin_subs.asp_quarterly").fetchone()[
|
||||
0
|
||||
]
|
||||
print(f" Loaded {count} rows into skin_subs.asp_quarterly")
|
||||
else:
|
||||
print(" WARNING: No skin substitute data found in ASP files")
|
||||
print(
|
||||
" This may be expected — skin subs may not appear in main ASP pricing files"
|
||||
)
|
||||
print(" They may be in separate NOC or tissue coding files")
|
||||
con.close()
|
||||
# Lock preflight + retry instead of a raw IOException when a notebook
|
||||
# kernel holds the single-writer file (#508).
|
||||
from conf.connect import duckdb_batch
|
||||
|
||||
with duckdb_batch("aco") as con:
|
||||
con.execute("CREATE SCHEMA IF NOT EXISTS skin_subs")
|
||||
con.execute("DROP TABLE IF EXISTS skin_subs.asp_quarterly")
|
||||
if not skin_df.empty:
|
||||
con.execute(
|
||||
"""
|
||||
CREATE TABLE skin_subs.asp_quarterly AS
|
||||
SELECT * FROM read_csv_auto(?, header=true)
|
||||
""",
|
||||
[str(OUTPUT_CSV)],
|
||||
)
|
||||
count = con.execute(
|
||||
"SELECT count(*) FROM skin_subs.asp_quarterly"
|
||||
).fetchone()[0]
|
||||
print(f" Loaded {count} rows into skin_subs.asp_quarterly")
|
||||
else:
|
||||
print(" WARNING: No skin substitute data found in ASP files")
|
||||
print(
|
||||
" This may be expected — skin subs may not appear in main ASP pricing files"
|
||||
)
|
||||
print(" They may be in separate NOC or tissue coding files")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -270,60 +270,79 @@ def ingest_all(con: duckdb.DuckDBPyConnection, year_filter: str = "") -> None:
|
||||
# Load into DuckDB
|
||||
print("\n--- Loading into DuckDB ---")
|
||||
|
||||
def _load(table: str, frames: list[pd.DataFrame]) -> None:
|
||||
"""Replace `table` — but with --year, merge instead of wipe.
|
||||
|
||||
A year-filtered run only parses that year's files, so a bare
|
||||
DROP/CREATE here used to erase every other year (#509). When
|
||||
filtering, keep the existing table's other-year rows and rebuild
|
||||
from their union; pandas concat also unions columns, so CMS
|
||||
format drift between years stays handled either way.
|
||||
"""
|
||||
merged = pd.concat(frames, ignore_index=True)
|
||||
if year_filter:
|
||||
exists = con.execute(
|
||||
"SELECT 1 FROM information_schema.tables "
|
||||
"WHERE table_schema = 'opps' AND table_name = ?",
|
||||
[table],
|
||||
).fetchone()
|
||||
if exists:
|
||||
other_years = con.execute(
|
||||
f"SELECT * FROM opps.{table} WHERE year != ?", # noqa: S608
|
||||
[int(year_filter)],
|
||||
).df()
|
||||
merged = pd.concat([other_years, merged], ignore_index=True)
|
||||
con.execute(f"DROP TABLE IF EXISTS opps.{table}")
|
||||
con.execute(f"CREATE TABLE opps.{table} AS SELECT * FROM merged") # noqa: S608
|
||||
print(f" opps.{table}: {len(merged)} rows ({merged['year'].nunique()} years)")
|
||||
|
||||
if apc_frames:
|
||||
all_apc = pd.concat(apc_frames, ignore_index=True)
|
||||
con.execute("DROP TABLE IF EXISTS opps.apc_weight")
|
||||
con.execute("CREATE TABLE opps.apc_weight AS SELECT * FROM all_apc")
|
||||
print(
|
||||
f" opps.apc_weight: {len(all_apc)} rows "
|
||||
f"({all_apc['year'].nunique()} years)"
|
||||
)
|
||||
_load("apc_weight", apc_frames)
|
||||
|
||||
if addb_frames:
|
||||
all_addb = pd.concat(addb_frames, ignore_index=True)
|
||||
con.execute("DROP TABLE IF EXISTS opps.addendum_b")
|
||||
con.execute("CREATE TABLE opps.addendum_b AS SELECT * FROM all_addb")
|
||||
print(
|
||||
f" opps.addendum_b: {len(all_addb)} rows "
|
||||
f"({all_addb['year'].nunique()} years)"
|
||||
)
|
||||
_load("addendum_b", addb_frames)
|
||||
|
||||
if wage_frames:
|
||||
all_wage = pd.concat(wage_frames, ignore_index=True)
|
||||
con.execute("DROP TABLE IF EXISTS opps.wage_index")
|
||||
con.execute("CREATE TABLE opps.wage_index AS SELECT * FROM all_wage")
|
||||
print(
|
||||
f" opps.wage_index: {len(all_wage)} rows "
|
||||
f"({all_wage['year'].nunique()} years)"
|
||||
)
|
||||
_load("wage_index", wage_frames)
|
||||
|
||||
# Skin sub specific: extract skin sub codes from Addendum B history
|
||||
# Skin sub specific: extract skin sub codes from Addendum B history.
|
||||
# Derived from the (merged) table rather than this run's frames, so a
|
||||
# --year run keeps the other years here too.
|
||||
if addb_frames:
|
||||
print("\n--- Skin substitute history in Addendum B ---")
|
||||
all_addb = pd.concat(addb_frames, ignore_index=True)
|
||||
if "hcpcs" in all_addb.columns:
|
||||
skin_codes = all_addb[
|
||||
all_addb["hcpcs"]
|
||||
.astype(str)
|
||||
.str.match(r"^Q4\d{2,3}$|^C527[1-8]$", na=False)
|
||||
]
|
||||
if not skin_codes.empty:
|
||||
con.execute("DROP TABLE IF EXISTS opps.skin_sub_addendum_b")
|
||||
con.execute(
|
||||
"CREATE TABLE opps.skin_sub_addendum_b AS SELECT * FROM skin_codes"
|
||||
)
|
||||
print(
|
||||
f" opps.skin_sub_addendum_b: {len(skin_codes)} rows "
|
||||
f"({skin_codes['year'].nunique()} years, "
|
||||
f"{skin_codes['hcpcs'].nunique()} unique codes)"
|
||||
)
|
||||
con.execute("DROP TABLE IF EXISTS opps.skin_sub_addendum_b")
|
||||
con.execute(
|
||||
"CREATE TABLE opps.skin_sub_addendum_b AS "
|
||||
"SELECT * FROM opps.addendum_b "
|
||||
"WHERE regexp_matches(CAST(hcpcs AS VARCHAR), '^Q4\\d{2,3}$|^C527[1-8]$')"
|
||||
)
|
||||
n, years, codes = con.execute(
|
||||
"SELECT count(*), count(DISTINCT year), count(DISTINCT hcpcs) "
|
||||
"FROM opps.skin_sub_addendum_b"
|
||||
).fetchone()
|
||||
print(
|
||||
f" opps.skin_sub_addendum_b: {n} rows ({years} years, {codes} unique codes)"
|
||||
)
|
||||
|
||||
# SI distribution for skin subs
|
||||
if "status_indicator" in skin_codes.columns:
|
||||
print(" Status indicators for skin subs by year:")
|
||||
si_year = skin_codes.groupby(["year", "status_indicator"]).size()
|
||||
for (yr, si), ct in si_year.items():
|
||||
print(f" CY{yr} SI={si}: {ct} codes")
|
||||
# SI distribution for skin subs (column absent in some year formats)
|
||||
has_si = con.execute(
|
||||
"SELECT 1 FROM information_schema.columns "
|
||||
"WHERE table_schema = 'opps' AND table_name = 'skin_sub_addendum_b' "
|
||||
"AND column_name = 'status_indicator'"
|
||||
).fetchone()
|
||||
si_rows = (
|
||||
con.execute(
|
||||
"SELECT year, status_indicator, count(*) "
|
||||
"FROM opps.skin_sub_addendum_b "
|
||||
"GROUP BY year, status_indicator ORDER BY year, status_indicator"
|
||||
).fetchall()
|
||||
if has_si
|
||||
else []
|
||||
)
|
||||
if si_rows:
|
||||
print(" Status indicators for skin subs by year:")
|
||||
for yr, si, ct in si_rows:
|
||||
print(f" CY{yr} SI={si}: {ct} codes")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
@@ -332,20 +351,23 @@ def main() -> None:
|
||||
args = parser.parse_args()
|
||||
|
||||
print("Ingesting CMS OPPS files into DuckDB ...")
|
||||
con = duckdb.connect(str(DUCKDB_PATH))
|
||||
# duckdb_batch preflights the single-writer lock (retry + name the
|
||||
# holding PID) instead of dying on a raw IOException when a notebook
|
||||
# kernel holds a connection (#508).
|
||||
from conf.connect import duckdb_batch
|
||||
|
||||
ingest_all(con, year_filter=args.year or "")
|
||||
with duckdb_batch("aco") as con:
|
||||
ingest_all(con, year_filter=args.year or "")
|
||||
|
||||
# Final inventory
|
||||
print("\n--- OPPS tables ---")
|
||||
for r in con.execute("""
|
||||
SELECT table_name FROM information_schema.tables
|
||||
WHERE table_schema = 'opps' ORDER BY table_name
|
||||
""").fetchall():
|
||||
cnt = con.execute(f"SELECT count(*) FROM opps.{r[0]}").fetchone()[0]
|
||||
print(f" opps.{r[0]:30s}: {cnt:>6} rows")
|
||||
# Final inventory
|
||||
print("\n--- OPPS tables ---")
|
||||
for r in con.execute("""
|
||||
SELECT table_name FROM information_schema.tables
|
||||
WHERE table_schema = 'opps' ORDER BY table_name
|
||||
""").fetchall():
|
||||
cnt = con.execute(f"SELECT count(*) FROM opps.{r[0]}").fetchone()[0] # noqa: S608
|
||||
print(f" opps.{r[0]:30s}: {cnt:>6} rows")
|
||||
|
||||
con.close()
|
||||
print("\nDone.")
|
||||
|
||||
|
||||
|
||||
59
docs/docs/duckdb-concurrency.md
Normal file
59
docs/docs/duckdb-concurrency.md
Normal file
@@ -0,0 +1,59 @@
|
||||
---
|
||||
title: DuckDB concurrency & connection hygiene
|
||||
sidebar_position: 90
|
||||
---
|
||||
|
||||
# DuckDB concurrency & connection hygiene
|
||||
|
||||
`data/aco.duckdb` is a **single-writer** store: one process holding *any*
|
||||
connection — even read-only — blocks every writer. The classic failure is a
|
||||
batch ingest dying with:
|
||||
|
||||
```
|
||||
IOException: Could not set lock on file "data/aco.duckdb"
|
||||
```
|
||||
|
||||
while a marimo notebook kernel quietly holds a connection from a cell that ran
|
||||
hours ago.
|
||||
|
||||
## Writing: use `conf.connect.duckdb_batch`
|
||||
|
||||
Batch ingests should not call `duckdb.connect()` directly. The
|
||||
`duckdb_batch()` context manager preflights the lock — retrying with
|
||||
exponential backoff, and if the file is still held, failing with the holder
|
||||
PIDs/commands (via `lsof`) instead of a raw `IOException`:
|
||||
|
||||
```python
|
||||
from conf.connect import duckdb_batch
|
||||
|
||||
with duckdb_batch("aco") as con: # retries, then names the lock holder
|
||||
con.execute("CREATE SCHEMA IF NOT EXISTS opps")
|
||||
...
|
||||
# connection is always closed here — the lock is released even on error
|
||||
```
|
||||
|
||||
`dev/scripts/ingest_opps.py` and `dev/scripts/ingest_asp.py` use this.
|
||||
|
||||
## Reading in notebooks: keep connections short-lived
|
||||
|
||||
- `conf.connect.duckdb()` opens **read-only by default** — keep it that way in
|
||||
notebooks; never pass `read_only=False` from a notebook.
|
||||
- A read-only handle still blocks writers. Don't keep a module-level
|
||||
connection alive for the life of the kernel: open, query, and close within
|
||||
the cell, or wrap access in a small helper that closes after each query.
|
||||
- If an ingest reports the lock is held by a `marimo` PID, close that notebook
|
||||
tab (or stop the PID) and re-run — the notebook loses nothing; it reconnects
|
||||
on the next cell run.
|
||||
|
||||
## Per-year ingests merge, not wipe
|
||||
|
||||
`ingest_opps.py --year YYYY` used to `DROP TABLE` + `CREATE ... AS` from only
|
||||
that year's files, silently erasing every other year (#509). It now merges:
|
||||
existing rows for other years are preserved and the filtered year is replaced.
|
||||
|
||||
## Roadmap
|
||||
|
||||
The strategic fixes (per-domain DB files, read replica, lake storage with
|
||||
snapshot isolation) are tracked as milestones M2–M5 in issues
|
||||
[#510](https://git.fhirworx.io/homelab/stack/issues/510)–[#514](https://git.fhirworx.io/homelab/stack/issues/514),
|
||||
per `docs/superpowers/specs/2026-07-08-duckdb-concurrency-streaming.md`.
|
||||
@@ -23,6 +23,7 @@ const sidebars = {
|
||||
dirName: "research",
|
||||
},
|
||||
"dag",
|
||||
"duckdb-concurrency",
|
||||
],
|
||||
};
|
||||
|
||||
|
||||
@@ -8,6 +8,8 @@ Usage::
|
||||
from conf.connect import duckdb, trino, bib, nessie, polaris, s3, zotero
|
||||
|
||||
con = duckdb() # DuckDB from cfg.db.aco
|
||||
with duckdb_batch() as con: # RW with lock preflight/retry (ingests)
|
||||
...
|
||||
con = duckdb("bib") # DuckDB for bib database
|
||||
store = bib() # bib.Store
|
||||
con = trino() # Trino with catalog from config
|
||||
@@ -21,6 +23,7 @@ from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
from contextlib import contextmanager
|
||||
from typing import Any
|
||||
|
||||
from conf import ROOT, cfg, path
|
||||
@@ -42,6 +45,75 @@ def duckdb(name: str = "aco", *, read_only: bool = True) -> Any:
|
||||
return _duckdb.connect(db_path, read_only=read_only)
|
||||
|
||||
|
||||
def _lock_holders(db_path: str) -> str:
|
||||
"""Best-effort description of the processes holding *db_path* open."""
|
||||
import subprocess
|
||||
|
||||
try:
|
||||
out = subprocess.run(
|
||||
["lsof", "--", db_path],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=10,
|
||||
).stdout.strip()
|
||||
return (
|
||||
out or "(lsof found no holders — lock may be from another host/container)"
|
||||
)
|
||||
except (OSError, subprocess.TimeoutExpired):
|
||||
return "(lsof unavailable)"
|
||||
|
||||
|
||||
@contextmanager
|
||||
def duckdb_batch(
|
||||
name: str = "aco",
|
||||
*,
|
||||
retries: int = 5,
|
||||
backoff_s: float = 2.0,
|
||||
) -> Any:
|
||||
"""Short-lived read-write DuckDB connection for batch ingests.
|
||||
|
||||
DuckDB is single-writer: even a *read-only* handle (typically a
|
||||
marimo notebook kernel) blocks a writer, surfacing as a raw
|
||||
``IOException: Could not set lock on file``. This helper retries
|
||||
with exponential backoff and, if the lock is still held, raises
|
||||
with the holder PIDs/commands so the failure is actionable
|
||||
("close the notebook tab or stop that PID").
|
||||
|
||||
Always closes the connection on exit — callers should do all their
|
||||
writes inside the ``with`` block instead of holding a module-level
|
||||
connection for the life of the process (see
|
||||
docs/docs/duckdb-concurrency.md).
|
||||
"""
|
||||
import time
|
||||
|
||||
import duckdb as _duckdb
|
||||
|
||||
db_path = str(path(f"db.{name}"))
|
||||
last_err: Exception | None = None
|
||||
for attempt in range(retries):
|
||||
try:
|
||||
con = _duckdb.connect(db_path, read_only=False)
|
||||
break
|
||||
except _duckdb.IOException as e:
|
||||
if "lock" not in str(e).lower():
|
||||
raise
|
||||
last_err = e
|
||||
if attempt < retries - 1:
|
||||
time.sleep(backoff_s * (2**attempt))
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Could not acquire write lock on {db_path} after {retries} "
|
||||
f"attempts.\nHolders:\n{_lock_holders(db_path)}\n"
|
||||
"If a holder is a marimo kernel, close the notebook tab or "
|
||||
"stop that PID, then re-run."
|
||||
) from last_err
|
||||
|
||||
try:
|
||||
yield con
|
||||
finally:
|
||||
con.close()
|
||||
|
||||
|
||||
def trino(
|
||||
*,
|
||||
catalog: str = "",
|
||||
|
||||
@@ -189,3 +189,79 @@ class TestTheme:
|
||||
import altair as alt
|
||||
|
||||
assert alt.theme.active == "fhirworx"
|
||||
|
||||
|
||||
class TestDuckdbBatch:
|
||||
def _patch_path(self, monkeypatch, db_file):
|
||||
monkeypatch.setattr(
|
||||
"conf.connect.path",
|
||||
lambda key: db_file if key == "db.custom" else path(key),
|
||||
)
|
||||
|
||||
def test_yields_writable_connection_and_closes(self, tmp_path, monkeypatch):
|
||||
import duckdb as _duckdb
|
||||
|
||||
db_file = tmp_path / "custom.duckdb"
|
||||
_duckdb.connect(str(db_file)).close()
|
||||
self._patch_path(monkeypatch, db_file)
|
||||
|
||||
with connect.duckdb_batch("custom") as con:
|
||||
con.execute("CREATE TABLE t (x INT)")
|
||||
con.execute("INSERT INTO t VALUES (1)")
|
||||
kept = con
|
||||
# closed on exit
|
||||
with pytest.raises(Exception):
|
||||
kept.execute("SELECT 1")
|
||||
# write persisted
|
||||
con2 = connect.duckdb("custom")
|
||||
assert con2.execute("SELECT count(*) FROM t").fetchone()[0] == 1
|
||||
con2.close()
|
||||
|
||||
def test_lock_held_raises_actionable_error(self, tmp_path, monkeypatch):
|
||||
"""A connection held by another process blocks the writer — after
|
||||
retries the error names the holders instead of a raw IOException
|
||||
(#508). Same-process connections share one DuckDB instance, so the
|
||||
holder must be a subprocess."""
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
import duckdb as _duckdb
|
||||
|
||||
db_file = tmp_path / "custom.duckdb"
|
||||
_duckdb.connect(str(db_file)).close()
|
||||
self._patch_path(monkeypatch, db_file)
|
||||
|
||||
holder = subprocess.Popen(
|
||||
[
|
||||
sys.executable,
|
||||
"-c",
|
||||
"import time, duckdb; "
|
||||
f"con = duckdb.connect({str(db_file)!r}); "
|
||||
"print('locked', flush=True); time.sleep(60)",
|
||||
],
|
||||
stdout=subprocess.PIPE,
|
||||
text=True,
|
||||
)
|
||||
try:
|
||||
assert holder.stdout is not None
|
||||
assert holder.stdout.readline().strip() == "locked"
|
||||
with pytest.raises(RuntimeError, match="(?s)write lock.*Holders"):
|
||||
with connect.duckdb_batch("custom", retries=2, backoff_s=0.05):
|
||||
pass
|
||||
finally:
|
||||
holder.kill()
|
||||
holder.wait()
|
||||
|
||||
def test_closes_on_exception_in_body(self, tmp_path, monkeypatch):
|
||||
import duckdb as _duckdb
|
||||
|
||||
db_file = tmp_path / "custom.duckdb"
|
||||
_duckdb.connect(str(db_file)).close()
|
||||
self._patch_path(monkeypatch, db_file)
|
||||
|
||||
with pytest.raises(ValueError, match="boom"):
|
||||
with connect.duckdb_batch("custom") as con:
|
||||
raise ValueError("boom")
|
||||
# lock released — a fresh writer can open immediately
|
||||
with connect.duckdb_batch("custom", retries=1) as con:
|
||||
con.execute("SELECT 1")
|
||||
|
||||
Reference in New Issue
Block a user