fix(cli): stack pfs elements --dry-run reads the replica, never the writer (refs #684)

This commit is contained in:
kert
2026-09-09 12:18:00 -04:00
parent 72a2d06e8a
commit 7fcec9c908
2 changed files with 45 additions and 12 deletions

View File

@@ -66,6 +66,16 @@ def _classifier() -> Any:
return closed_vocab_classifier(cfg, HostPool.from_config(cfg))
def _run_elements(
con: Any, store: Any, targets: list[str], classify: Any, *, write: bool
) -> None:
for c in targets:
x = extract_code(store, con, c, classify=classify)
if write:
write_elements(con, c, x.rows, x.reviews)
typer.echo(f"{c}: {len(x.rows)} elements, {len(x.reviews)} for review")
def _codes_for(con: Any, codes: list[str], family: str, all_payable: bool) -> list[str]:
out: list[str] = [c.upper() for c in codes]
if family:
@@ -106,15 +116,21 @@ def elements(
"""Extract typed elements for codes into pfs.code_element (+ review queue)."""
store = _store()
classify = None if no_llm else _classifier()
with _batch() as con:
ensure_tables(con)
targets = _codes_for(con, code, family, all_payable)
for c in targets:
x = extract_code(store, con, c, classify=classify)
if not dry_run:
write_elements(con, c, x.rows, x.reviews)
typer.echo(f"{c}: {len(x.rows)} elements, {len(x.reviews)} for review")
if not dry_run:
if dry_run:
# A preview never contends for the DuckDB single-writer lock a
# notebook may be holding (#508-#514) — read the replica, don't
# open a batch writer.
con = _read()
try:
targets = _codes_for(con, code, family, all_payable)
_run_elements(con, store, targets, classify, write=False)
finally:
con.close()
else:
with _batch() as con:
ensure_tables(con)
targets = _codes_for(con, code, family, all_payable)
_run_elements(con, store, targets, classify, write=True)
_publish()
@@ -155,7 +171,7 @@ def families(write: bool = typer.Option(False, "--write")) -> None:
"UNION SELECT DISTINCT hcpcs FROM pfs.rvu WHERE year = (SELECT max(year) FROM pfs.rvu) AND status_code IN ('A','R','T')"
).fetchall()
]
elements = {c: read_elements(con, c) for c in codes}
elements_by_code = {c: read_elements(con, c) for c in codes}
events = {c: read_events(con, c) for c in codes}
descriptions = {
r[0]: r[1]
@@ -163,7 +179,7 @@ def families(write: bool = typer.Option(False, "--write")) -> None:
"SELECT hcpcs, arg_max(description, year) FROM pfs.rvu WHERE mod IS NULL OR mod = '' GROUP BY hcpcs"
).fetchall()
}
rows = derive_families(elements, events, descriptions)
rows = derive_families(elements_by_code, events, descriptions)
by_key: dict[str, list[str]] = {}
for r in rows:
by_key.setdefault(r.key, []).append(f"{r.code}({r.role})")

View File

@@ -121,10 +121,27 @@ class TestElements:
"extract_code",
lambda s, c, code, *, classify=None: Extraction(code, (), ()),
)
def fail_batch():
raise AssertionError("--dry-run must not open a RW duckdb_batch connection")
monkeypatch.setattr(pfs_cli, "_batch", fail_batch)
read_calls = []
orig_read = pfs_cli._read
def spy_read():
read_calls.append(True)
return orig_read()
monkeypatch.setattr(pfs_cli, "_read", spy_read)
res = runner.invoke(
app, ["pfs", "elements", "--code", "G0556", "--no-llm", "--dry-run"]
)
assert res.exit_code == 0 and con.published == []
assert res.exit_code == 0, res.output
assert con.published == []
assert read_calls == [True]
def test_all_payable_selects_art_status_in_newest_year(self, con, monkeypatch):
seen = []