From 62d3be2fff5027653ca1aa6a98e0015589299aef Mon Sep 17 00:00:00 2001 From: Georgie Kennedy Date: Mon, 14 Sep 2026 17:18:15 +1000 Subject: [PATCH 1/2] partial load safety --- README.md | 5 + src/pbs_client/cli/main.py | 28 +++- src/pbs_client/db/engine.py | 32 ++++- src/pbs_client/sync/__init__.py | 16 ++- src/pbs_client/sync/orchestrator.py | 30 ++++ .../toolkit/analytics/indications.py | 36 +++-- src/pbs_client/toolkit/core/service.py | 4 +- tests/fixtures/criteria-restriction-path.json | 86 +++++++++++ tests/test_fk_enforcement.py | 4 + tests/test_query.py | 134 +++++------------- tests/test_sync.py | 71 ++++++++++ 11 files changed, 328 insertions(+), 118 deletions(-) create mode 100644 tests/fixtures/criteria-restriction-path.json diff --git a/README.md b/README.md index 05d902c..81f3964 100644 --- a/README.md +++ b/README.md @@ -9,8 +9,13 @@ uv sync --extra dev uv run omop-config configure pbs_client uv run pbs-client init-db uv run pbs-client sync +uv run pbs-client verify ``` +`verify` is a read-only check that each recorded sync has written at least the +API-reported `total_records`; it compares checkpoint writes, not distinct table +row counts, because some resources intentionally upsert duplicate keys. + The package registers `PBSClientConfig` with `oa-configurator` under the `pbs_client` tool name. The configuration wizard creates a `[tools.pbs_client]` section and a named generic database entry (`pbs_db`) that can be shared with diff --git a/src/pbs_client/cli/main.py b/src/pbs_client/cli/main.py index b69b2e2..5f50178 100644 --- a/src/pbs_client/cli/main.py +++ b/src/pbs_client/cli/main.py @@ -12,7 +12,7 @@ from pbs_client.db import init_db, make_session_factory from pbs_client.errors import PBSHTTPError, PBSInvalidResponseError, PBSSyncError, PBSTransportError from pbs_client.http import DEFAULT_PAGE_SIZE, PBSClient -from pbs_client.sync import SyncOrchestrator, mirror_status +from pbs_client.sync import SyncOrchestrator, mirror_status, sync_integrity_issues app = typer.Typer(help="Maintain a local offline mirror of the PBS Public Data API v3.") @@ -203,5 +203,31 @@ def status() -> None: typer.echo(line) +@app.command() +def verify() -> None: + """Check that recorded sync writes are not below API-reported totals.""" + + _, engine, _ = _runtime() + with make_session_factory(engine)() as session: + issues = sync_integrity_issues(session) + if not issues: + typer.echo("PBS mirror verification passed: no recorded resource is below its API total.") + return + + typer.echo( + f"PBS mirror verification failed for {len(issues)} resource(s):", + err=True, + ) + for state in issues: + total_records = state.metadata_json["total_records"] + missing_records = total_records - state.records_written + typer.echo( + f" {state.resource}: {state.records_written:,}/{total_records:,} " + f"records written (missing {missing_records:,}; status={state.status})", + err=True, + ) + raise typer.Exit(code=1) + + def main() -> None: app() diff --git a/src/pbs_client/db/engine.py b/src/pbs_client/db/engine.py index 5a9a484..4e54691 100644 --- a/src/pbs_client/db/engine.py +++ b/src/pbs_client/db/engine.py @@ -5,7 +5,6 @@ from collections.abc import Iterator from contextlib import contextmanager -import sqlalchemy as sa from orm_loader.backends.resolve import resolve_backend from orm_loader.helpers.bulk import engine_with_replica_role from sqlalchemy import Engine, event @@ -14,8 +13,25 @@ from pbs_client.db.model import Base +def _set_sqlite_foreign_keys(dbapi_connection, enabled: bool) -> int: + """Set this connection's FK flag outside any SQLite transaction.""" + + autocommit = dbapi_connection.autocommit + dbapi_connection.autocommit = True + try: + value = "ON" if enabled else "OFF" + dbapi_connection.execute(f"PRAGMA foreign_keys = {value}").close() + cursor = dbapi_connection.execute("PRAGMA foreign_keys") + try: + return cursor.fetchone()[0] + finally: + cursor.close() + finally: + dbapi_connection.autocommit = autocommit + + def _enable_sqlite_foreign_keys(dbapi_connection, _connection_record) -> None: - dbapi_connection.execute("PRAGMA foreign_keys = ON") + _set_sqlite_foreign_keys(dbapi_connection, True) @contextmanager @@ -43,11 +59,11 @@ def fk_checks_disabled_for_refresh(engine: Engine) -> Iterator[None]: raise NotImplementedError(f"Unsupported database backend for PBS refresh: {backend.name}") def disable_on_pool_event(dbapi_connection, *_args) -> None: - dbapi_connection.execute("PRAGMA foreign_keys = OFF") + _set_sqlite_foreign_keys(dbapi_connection, False) def enable_on_checkin(dbapi_connection, *_args) -> None: if dbapi_connection is not None: - dbapi_connection.execute("PRAGMA foreign_keys = ON") + _set_sqlite_foreign_keys(dbapi_connection, True) event.listen(engine, "connect", disable_on_pool_event) event.listen(engine, "checkout", disable_on_pool_event) @@ -59,8 +75,10 @@ def enable_on_checkin(dbapi_connection, *_args) -> None: event.remove(engine, "checkout", disable_on_pool_event) try: with engine.connect() as connection: - connection.execute(sa.text("PRAGMA foreign_keys = ON")) - state = connection.execute(sa.text("PRAGMA foreign_keys")).scalar_one() + state = _set_sqlite_foreign_keys( + connection.connection.driver_connection, + True, + ) if state != 1: raise RuntimeError("Failed to restore SQLite foreign-key enforcement") finally: @@ -80,5 +98,5 @@ def init_db(engine: Engine) -> None: if not event.contains(engine, "connect", _enable_sqlite_foreign_keys): event.listen(engine, "connect", _enable_sqlite_foreign_keys) with engine.connect() as connection: - connection.execute(sa.text("PRAGMA foreign_keys = ON")) + _set_sqlite_foreign_keys(connection.connection.driver_connection, True) Base.metadata.create_all(engine) diff --git a/src/pbs_client/sync/__init__.py b/src/pbs_client/sync/__init__.py index 2215bcd..8e4d7cf 100644 --- a/src/pbs_client/sync/__init__.py +++ b/src/pbs_client/sync/__init__.py @@ -1,5 +1,17 @@ """PBS API to local database synchronization.""" -from pbs_client.sync.orchestrator import SyncOrchestrator, SyncResult, mirror_status, upsert_records +from pbs_client.sync.orchestrator import ( + SyncOrchestrator, + SyncResult, + mirror_status, + sync_integrity_issues, + upsert_records, +) -__all__ = ["SyncOrchestrator", "SyncResult", "mirror_status", "upsert_records"] +__all__ = [ + "SyncOrchestrator", + "SyncResult", + "mirror_status", + "sync_integrity_issues", + "upsert_records", +] diff --git a/src/pbs_client/sync/orchestrator.py b/src/pbs_client/sync/orchestrator.py index 524b07e..01eeea9 100644 --- a/src/pbs_client/sync/orchestrator.py +++ b/src/pbs_client/sync/orchestrator.py @@ -55,6 +55,17 @@ def upsert_records(session: Session, model: type[Base], records: Iterable[dict[s return written +def _underreported_total(state: SyncState) -> int | None: + total_records = state.metadata_json.get("total_records") + if ( + isinstance(total_records, int) + and not isinstance(total_records, bool) + and state.records_written < total_records + ): + return total_records + return None + + class SyncOrchestrator: """Coordinate API pages and local transactions in the required order.""" @@ -143,6 +154,11 @@ def _sync_resource(self, name: str, state: SyncState, *, limit: int) -> SyncResu session.commit() state = current logger.info("Synced %s page %s (%s records)", name, page.page, count) + if (total_records := _underreported_total(state)) is not None: + raise PBSSyncError( + f"{name} wrote {state.records_written:,} of the API-reported " + f"{total_records:,} records" + ) state.complete() self._save_state(state) except Exception as exc: @@ -226,3 +242,17 @@ def mirror_status(session: Session) -> list[dict[str, Any]]: } ) return rows + + +def sync_integrity_issues(session: Session) -> list[SyncState]: + """Find checkpoints that wrote fewer API records than the final page reports. + + Compare the API page count with ``records_written``, not table row counts: + upserts can legitimately collapse repeated records across schedules. + """ + + issues = [] + for state in session.scalars(select(SyncState).order_by(SyncState.resource)): + if _underreported_total(state) is not None: + issues.append(state) + return issues diff --git a/src/pbs_client/toolkit/analytics/indications.py b/src/pbs_client/toolkit/analytics/indications.py index 7e27d78..bd39ff6 100644 --- a/src/pbs_client/toolkit/analytics/indications.py +++ b/src/pbs_client/toolkit/analytics/indications.py @@ -13,10 +13,11 @@ CriteriaParameterRltd, DispensingRule, Item, - ItemPrescribingTxtRltd, + ItemRestrictionRltd, Parameter, PrescribingTxt, Program, + RstrctnPrscrbngTxtRltd, Schedule, ) from pbs_client.toolkit.core import ( @@ -49,7 +50,7 @@ class ParameterText: class CriteriaText: """One item-linked eligibility criterion with its linked parameter detail.""" - item_relationship: ItemPrescribingTxtRltd + item_relationship: ItemRestrictionRltd criteria: Criteria prescribing_text: PrescribingTxt parameters: tuple[ParameterText, ...] @@ -75,29 +76,44 @@ def get_item_criteria_breakdown(session: Session, item: Item) -> ItemCriteriaBre criteria_rows = [] if item.pbs_code is not None: criteria_rows = session.execute( - select(ItemPrescribingTxtRltd, Criteria, PrescribingTxt) + select(ItemRestrictionRltd, Criteria, PrescribingTxt) + .select_from(ItemRestrictionRltd) + .join( + RstrctnPrscrbngTxtRltd, + and_( + RstrctnPrscrbngTxtRltd.schedule_code + == ItemRestrictionRltd.schedule_code, + RstrctnPrscrbngTxtRltd.res_code == ItemRestrictionRltd.res_code, + ), + ) .join( PrescribingTxt, and_( - PrescribingTxt.schedule_code == ItemPrescribingTxtRltd.schedule_code, - PrescribingTxt.prescribing_txt_id == ItemPrescribingTxtRltd.prescribing_txt_id, + PrescribingTxt.schedule_code + == RstrctnPrscrbngTxtRltd.schedule_code, + PrescribingTxt.prescribing_txt_id + == RstrctnPrscrbngTxtRltd.prescribing_text_id, ), ) .join( Criteria, and_( Criteria.schedule_code == PrescribingTxt.schedule_code, - Criteria.criteria_prescribing_txt_id == PrescribingTxt.prescribing_txt_id, + Criteria.criteria_prescribing_txt_id + == PrescribingTxt.prescribing_txt_id, ), ) .where( - ItemPrescribingTxtRltd.schedule_code == item.schedule_code, - ItemPrescribingTxtRltd.pbs_code == item.pbs_code, + ItemRestrictionRltd.schedule_code == item.schedule_code, + ItemRestrictionRltd.pbs_code == item.pbs_code, + ItemRestrictionRltd.restriction_indicator == "Y", PrescribingTxt.prescribing_type == "CRITERIA", ) .order_by( - ItemPrescribingTxtRltd.pt_position, - ItemPrescribingTxtRltd.prescribing_txt_id, + ItemRestrictionRltd.res_position, + ItemRestrictionRltd.res_code, + RstrctnPrscrbngTxtRltd.pt_position, + RstrctnPrscrbngTxtRltd.prescribing_text_id, ) ).all() diff --git a/src/pbs_client/toolkit/core/service.py b/src/pbs_client/toolkit/core/service.py index f669295..0366be8 100644 --- a/src/pbs_client/toolkit/core/service.py +++ b/src/pbs_client/toolkit/core/service.py @@ -247,7 +247,9 @@ def get_item_indication_text(session: Session, item: Item) -> list[IndicationTex Notes and cautions are excluded using the PBS relationship's ``restriction_indicator`` field. A restriction produces a fallback only - when it has no usable linked ``INDICATION`` condition. + when it has no usable linked ``INDICATION`` condition. ``source`` records + provenance, not a confidence level; the fallback is not currently a + calibrated confidence signal. """ if item.pbs_code is None: diff --git a/tests/fixtures/criteria-restriction-path.json b/tests/fixtures/criteria-restriction-path.json new file mode 100644 index 0000000..ab73672 --- /dev/null +++ b/tests/fixtures/criteria-restriction-path.json @@ -0,0 +1,86 @@ +{ + "_capture_note": "Selected columns copied from raw_payload records in the read-only mirror for schedule 4708, item 10003L. No ItemPrescribingTxtRltd criterion row exists in this capture.", + "Schedule": { + "schedule_code": 4708, + "revision_number": 2, + "start_tsp": "2026-08-01T00:00:00.000+10:00", + "effective_date": "2026-08-01", + "effective_month": "AUGUST", + "effective_year": 2026, + "publication_status": "PUBLISHED" + }, + "Program": { + "schedule_code": 4708, + "program_code": "GE", + "program_title": "General Schedule" + }, + "DispensingRule": { + "schedule_code": 4708, + "dispensing_rule_mnem": "s90-cp", + "dispensing_rule_reference": "rp-s90-cp", + "dispensing_rule_title": "Community Pharmacy", + "community_pharmacy_indicator": "true" + }, + "Item": { + "schedule_code": 4708, + "li_item_id": "10003L_13467_29812_29815_29817", + "pbs_code": "10003L", + "drug_name": "Dabrafenib", + "program_code": "GE" + }, + "RestrictionText": { + "schedule_code": 4708, + "res_code": "17806_17949_R" + }, + "ItemRestrictionRltd": { + "schedule_code": 4708, + "pbs_code": "10003L", + "res_code": "17806_17949_R", + "benefit_type_code": "S", + "restriction_indicator": "Y", + "res_position": 1 + }, + "RstrctnPrscrbngTxtRltd": { + "schedule_code": 4708, + "res_code": "17806_17949_R", + "prescribing_text_id": 7738, + "pt_position": 2 + }, + "PrescribingTxt": [ + { + "schedule_code": 4708, + "prescribing_txt_id": 7738, + "prescribing_type": "CRITERIA", + "prescribing_txt": "Clinical criteria: Patient must have previously been issued with an authority prescription for this drug" + }, + { + "schedule_code": 4708, + "prescribing_txt_id": 7737, + "prescribing_type": "PARAMETER", + "prescribing_txt": "Patient must have previously been issued with an authority prescription for this drug" + } + ], + "Criteria": { + "schedule_code": 4708, + "criteria_prescribing_txt_id": 7738, + "criteria_type": "CLINICAL", + "parameter_relationship": "ANY" + }, + "CriteriaParameterRltd": { + "schedule_code": 4708, + "criteria_prescribing_txt_id": 7738, + "parameter_prescribing_txt_id": 7737, + "pt_position": 1 + }, + "Parameter": { + "schedule_code": 4708, + "parameter_prescribing_txt_id": 7737, + "assessment_type": "IMMEDIATE", + "parameter_type": "CLINICAL_PATIENT" + }, + "ItemDispensingRuleRltd": { + "schedule_code": 4708, + "li_item_id": "10003L_13467_29812_29815_29817", + "dispensing_rule_mnem": "s90-cp" + } +} diff --git a/tests/test_fk_enforcement.py b/tests/test_fk_enforcement.py index 0706b6b..06fa6fa 100644 --- a/tests/test_fk_enforcement.py +++ b/tests/test_fk_enforcement.py @@ -25,6 +25,7 @@ def orphan_item_atc(code: str) -> dict[str, object]: def test_sqlite_refresh_disables_fk_for_all_pool_connections_and_restores(tmp_path): engine = create_engine( f"sqlite:///{tmp_path / 'pbs.sqlite'}", + connect_args={"autocommit": False}, future=True, pool_size=2, max_overflow=0, @@ -32,6 +33,9 @@ def test_sqlite_refresh_disables_fk_for_all_pool_connections_and_restores(tmp_pa init_db(engine) sessions = sessionmaker(bind=engine, expire_on_commit=False, future=True) + with sessions() as session: + assert session.execute(text("PRAGMA foreign_keys")).scalar_one() == 1 + with fk_checks_disabled_for_refresh(engine): first, second = sessions(), sessions() try: diff --git a/tests/test_query.py b/tests/test_query.py index 1ca8912..21e4877 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -1,10 +1,11 @@ from __future__ import annotations +import json + +from pbs_client.db import MODEL_BY_NAME from pbs_client.db.model import ( ATC, Copayment, - Criteria, - CriteriaParameterRltd, DispensingRule, Fee, Indication, @@ -13,18 +14,17 @@ ItemAtcRltd, ItemDispensingRuleRltd, ItemOrganisationRltd, - ItemPrescribingTxtRltd, ItemPricingEvent, ItemRestrictionRltd, MarkupBand, Organisation, - Parameter, PrescribingTxt, Program, RestrictionText, RstrctnPrscrbngTxtRltd, Schedule, ) +from pbs_client.sync import upsert_records from pbs_client.toolkit.analytics import get_item_criteria_breakdown, indication_candidates from pbs_client.toolkit.core import ( BenefitTypeCode, @@ -426,109 +426,49 @@ def test_item_manufacturer_uses_item_manufacturer_link_not_wholesaler_link(sessi assert manufacturer.name == "Item manufacturer" -def test_item_criteria_breakdown_preserves_parameter_program_and_rule_context(session_factory): +def test_item_criteria_breakdown_uses_captured_restriction_path( + session_factory, fixture_dir +): + captured = json.loads((fixture_dir / "criteria-restriction-path.json").read_text()) with session_factory() as session: - session.add(Schedule(schedule_code=51, effective_date="2026-09-01", effective_year=2026)) - session.commit() - session.add_all( - [ - Program(schedule_code=51, program_code="CT", program_title="Chemotherapy"), - DispensingRule( - schedule_code=51, - dispensing_rule_mnem="HOSP", - dispensing_rule_title="Public hospital", - community_pharmacy_indicator="N", - ), - PrescribingTxt( - schedule_code=51, - prescribing_txt_id=100, - prescribing_type="CRITERIA", - prescribing_txt="Patient has the required clinical condition.", - ), - PrescribingTxt( - schedule_code=51, - prescribing_txt_id=200, - prescribing_type="PARAMETER", - prescribing_txt="Document the patient's clinical status.", - ), - PrescribingTxt( - schedule_code=51, - prescribing_txt_id=300, - prescribing_type="INDICATION", - prescribing_txt="Indication text is kept separate.", - ), - ] - ) - session.commit() - session.add_all( - [ - Criteria( - schedule_code=51, - criteria_prescribing_txt_id=100, - criteria_type="CLINICAL_PATIENT", - parameter_relationship="AND", - ), - Parameter( - schedule_code=51, - assessment_type="CLINICAL", - parameter_prescribing_txt_id=200, - parameter_type="CLINICAL_PATIENT", - ), - ] - ) - session.commit() - session.add( - Item( - schedule_code=51, - li_item_id="li-criteria", - pbs_code="C1", - program_code="CT", - ) - ) - session.commit() - session.add_all( - [ - ItemPrescribingTxtRltd( - schedule_code=51, - pbs_code="C1", - prescribing_txt_id=100, - pt_position=1, - ), - ItemPrescribingTxtRltd( - schedule_code=51, - pbs_code="C1", - prescribing_txt_id=300, - pt_position=2, - ), - CriteriaParameterRltd( - schedule_code=51, - criteria_prescribing_txt_id=100, - parameter_prescribing_txt_id=200, - pt_position=1, - ), - ItemDispensingRuleRltd( - schedule_code=51, - li_item_id="li-criteria", - dispensing_rule_mnem="HOSP", - ), - ] - ) + for resource in ( + "Schedule", + "Program", + "DispensingRule", + "RestrictionText", + "PrescribingTxt", + "Item", + "Criteria", + "Parameter", + "ItemRestrictionRltd", + "RstrctnPrscrbngTxtRltd", + "CriteriaParameterRltd", + "ItemDispensingRuleRltd", + ): + payload = captured[resource] + records = payload if isinstance(payload, list) else [payload] + upsert_records(session, MODEL_BY_NAME[resource], records) + session.flush() session.commit() - item = session.get(Item, (51, "li-criteria")) + item = session.get(Item, (4708, "10003L_13467_29812_29815_29817")) result = get_item_criteria_breakdown(session, item) links = get_item_dispensing_rule_links(session, item) assert len(result.criteria) == 1 criterion = result.criteria[0] - assert criterion.prescribing_text.prescribing_txt.startswith("Patient has") - assert criterion.criteria.criteria_type == "CLINICAL_PATIENT" + assert criterion.prescribing_text.prescribing_txt.startswith("Clinical criteria:") + assert criterion.item_relationship.res_code == "17806_17949_R" + assert isinstance(criterion.item_relationship, ItemRestrictionRltd) + assert criterion.criteria.criteria_type == "CLINICAL" assert len(criterion.parameters) == 1 - assert criterion.parameters[0].prescribing_text.prescribing_txt.startswith("Document") + assert criterion.parameters[0].prescribing_text.prescribing_txt.startswith( + "Patient must have previously" + ) assert criterion.parameters[0].parameters[0].parameter_type == "CLINICAL_PATIENT" - assert result.program.program_title == "Chemotherapy" - assert [rule.dispensing_rule_mnem for rule in result.dispensing_rules] == ["HOSP"] - assert [link.dispensing_rule_mnem for link in links] == ["HOSP"] + assert result.program.program_title == "General Schedule" + assert [rule.dispensing_rule_mnem for rule in result.dispensing_rules] == ["s90-cp"] + assert [link.dispensing_rule_mnem for link in links] == ["s90-cp"] def test_item_pricing_breakdown_returns_source_inputs_without_calculating_patient_amount( diff --git a/tests/test_sync.py b/tests/test_sync.py index 17ed481..ced700d 100644 --- a/tests/test_sync.py +++ b/tests/test_sync.py @@ -1,6 +1,10 @@ from __future__ import annotations from datetime import UTC, datetime +from importlib import import_module + +import pytest +from typer.testing import CliRunner from pbs_client.cli.main import _report_sync_failure, _status_lines from pbs_client.db import MODEL_BY_NAME, SyncState @@ -58,6 +62,73 @@ def test_sync_is_idempotent_and_upserts(session_factory): assert first[0].status == second[0].status == "complete" +def test_sync_fails_instead_of_completing_below_api_total(session_factory): + page = Page( + "/schedules", + 1, + 1, + [schedule(1, "2026-01-01")], + {"total_records": 2}, + [], + ) + + with pytest.raises(PBSSyncError, match="wrote 1 of the API-reported 2 records"): + SyncOrchestrator(FakeClient([page]), session_factory).run( + resource="Schedule", limit=1 + ) + + with session_factory() as session: + state = session.get(SyncState, "Schedule") + assert state.status == "failed" + assert state.records_written == 1 + assert state.metadata_json["total_records"] == 2 + assert state.completed_at is None + + +@pytest.mark.parametrize( + ("resource", "records_written", "total_records", "exit_code"), + [ + ("MarkupBand", 1748, 1748, 0), + ("ItemDispensingRuleRltd", 56181, 501181, 1), + ], +) +def test_verify_command_checks_written_count_not_table_rows( + session_factory, monkeypatch, resource, records_written, total_records, exit_code +): + endpoint = "/markup-bands" if resource == "MarkupBand" else "/item-dispensing-rule-relationships" + with session_factory() as session: + session.add( + SyncState( + resource=resource, + endpoint=endpoint, + status="complete", + records_written=records_written, + metadata_json={"total_records": total_records}, + ) + ) + session.commit() + + probe = session_factory() + engine = probe.get_bind() + probe.close() + cli_main = import_module("pbs_client.cli.main") + monkeypatch.setattr(cli_main, "_runtime", lambda: (None, engine, "test_db")) + monkeypatch.setattr( + cli_main, + "init_db", + lambda _: pytest.fail("verify must not mutate or initialize the mirror"), + ) + + result = CliRunner().invoke(cli_main.app, ["verify"]) + + assert result.exit_code == exit_code + if exit_code: + assert "ItemDispensingRuleRltd" in result.output + assert "56,181/501,181" in result.output + else: + assert "verification passed" in result.output + + def test_sync_resumes_after_a_committed_page(session_factory): pages = [ Page( From 204470d48b85698927878fac8c55892743a690f7 Mon Sep 17 00:00:00 2001 From: Georgie Kennedy Date: Mon, 14 Sep 2026 17:33:58 +1000 Subject: [PATCH 2/2] progress bar and fix locked DB issue --- README.md | 19 +++----- src/pbs_client/db/engine.py | 52 ++++++++++++---------- src/pbs_client/sync/orchestrator.py | 69 +++++++++++++++++++++-------- tests/test_sync.py | 57 ++++++++++++++++++++++++ 4 files changed, 142 insertions(+), 55 deletions(-) diff --git a/README.md b/README.md index 81f3964..6bc5393 100644 --- a/README.md +++ b/README.md @@ -12,14 +12,11 @@ uv run pbs-client sync uv run pbs-client verify ``` -`verify` is a read-only check that each recorded sync has written at least the -API-reported `total_records`; it compares checkpoint writes, not distinct table -row counts, because some resources intentionally upsert duplicate keys. +`verify` is a read-only check that each recorded sync has written at least the API-reported `total_records`; it compares checkpoint writes, not distinct table row counts, because some resources intentionally upsert duplicate keys. -The package registers `PBSClientConfig` with `oa-configurator` under the -`pbs_client` tool name. The configuration wizard creates a `[tools.pbs_client]` -section and a named generic database entry (`pbs_db`) that can be shared with -downstream packages such as Groundworkers. +In an interactive terminal, `sync` also shows a per-resource record progress bar from the API totals, advancing after each page is committed. + +The package registers `PBSClientConfig` with `oa-configurator` under the `pbs_client` tool name. The configuration wizard creates a `[tools.pbs_client]` section and a named generic database entry (`pbs_db`) that can be shared with downstream packages such as Groundworkers. For a non-interactive local SQLite setup, the relevant stack configuration is: @@ -37,12 +34,8 @@ pbs_db = "pbs_db" subscription_key = "your-subscription-key" ``` -Set `OA_CONFIG_PATH` before invoking the commands when using a config file -outside `~/.config/omop/config.toml`. +Set `OA_CONFIG_PATH` before invoking the commands when using a config file outside `~/.config/omop/config.toml`. The public API is deliberately rate limited to one request per twenty seconds. The client enforces that interval process-wide, including retries and page continuations, and rejects configured intervals below three seconds because the quota is shared across users. The default page size is 5,000 records: this keeps large responses manageable without creating unnecessary calls against the shared quota. Use `--limit 1000` when an endpoint still returns an empty or non-JSON response; if a page-size change is made during a resume, the affected resource safely restarts from page one. Refreshes are upserts and intentionally retain rows no longer returned by a later response, preserving local PBS history. Tests use local fixtures and never call the API. -All configuration — the subscription key, base URL, rate limit, and the -shared mirror database — is read from `oa-configurator`. There is no -environment-variable or CLI-flag fallback; run `uv run omop-config configure -pbs_client` before using the library or CLI. +All configuration — the subscription key, base URL, rate limit, and the shared mirror database — is read from `oa-configurator`. There is no environment-variable or CLI-flag fallback; run `uv run omop-config configure pbs_client` before using the library or CLI. diff --git a/src/pbs_client/db/engine.py b/src/pbs_client/db/engine.py index 4e54691..b6a8d26 100644 --- a/src/pbs_client/db/engine.py +++ b/src/pbs_client/db/engine.py @@ -17,6 +17,7 @@ def _set_sqlite_foreign_keys(dbapi_connection, enabled: bool) -> int: """Set this connection's FK flag outside any SQLite transaction.""" autocommit = dbapi_connection.autocommit + dbapi_connection.rollback() dbapi_connection.autocommit = True try: value = "ON" if enabled else "OFF" @@ -40,9 +41,9 @@ def fk_checks_disabled_for_refresh(engine: Engine) -> Iterator[None]: The orchestrator commits each API page in a fresh session, so a session scoped PRAGMA/replication-role change would only affect one arbitrary pool - connection. SQLite uses pool checkout/checkin listeners; Postgres uses - orm_loader's engine-scoped replica role and disposes the pool afterwards so - no connection carrying the disabled role can be reused. + connection. SQLite disables checks on checkout and restores them after the + refresh; Postgres disposes the pool so disabled replica connections cannot + be reused. """ backend = resolve_backend(engine) @@ -58,31 +59,34 @@ def fk_checks_disabled_for_refresh(engine: Engine) -> Iterator[None]: if backend.name != "sqlite": raise NotImplementedError(f"Unsupported database backend for PBS refresh: {backend.name}") - def disable_on_pool_event(dbapi_connection, *_args) -> None: - _set_sqlite_foreign_keys(dbapi_connection, False) - - def enable_on_checkin(dbapi_connection, *_args) -> None: - if dbapi_connection is not None: - _set_sqlite_foreign_keys(dbapi_connection, True) + def disable_on_checkout(dbapi_connection, *_args) -> None: + cursor = dbapi_connection.execute("PRAGMA foreign_keys") + try: + if cursor.fetchone()[0] != 0: + _set_sqlite_foreign_keys(dbapi_connection, False) + finally: + cursor.close() - event.listen(engine, "connect", disable_on_pool_event) - event.listen(engine, "checkout", disable_on_pool_event) - event.listen(engine, "checkin", enable_on_checkin) + event.listen(engine, "checkout", disable_on_checkout) try: yield finally: - event.remove(engine, "connect", disable_on_pool_event) - event.remove(engine, "checkout", disable_on_pool_event) - try: - with engine.connect() as connection: - state = _set_sqlite_foreign_keys( - connection.connection.driver_connection, - True, - ) - if state != 1: - raise RuntimeError("Failed to restore SQLite foreign-key enforcement") - finally: - event.remove(engine, "checkin", enable_on_checkin) + event.remove(engine, "checkout", disable_on_checkout) + database = engine.url.database + in_memory = ( + database in (None, "", ":memory:", "file::memory:") + or engine.url.query.get("mode") == "memory" + ) + if not in_memory: + # File-backed pool connections may still have FK checks disabled. + engine.dispose() + with engine.connect() as connection: + state = _set_sqlite_foreign_keys( + connection.connection.driver_connection, + True, + ) + if state != 1: + raise RuntimeError("Failed to restore SQLite foreign-key enforcement") def make_session_factory(engine: Engine): diff --git a/src/pbs_client/sync/orchestrator.py b/src/pbs_client/sync/orchestrator.py index 01eeea9..5b89536 100644 --- a/src/pbs_client/sync/orchestrator.py +++ b/src/pbs_client/sync/orchestrator.py @@ -3,10 +3,19 @@ from __future__ import annotations import logging +import sys from collections.abc import Callable, Iterable from dataclasses import dataclass from typing import Any +from rich.progress import ( + BarColumn, + MofNCompleteColumn, + Progress, + TextColumn, + TimeElapsedColumn, + TimeRemainingColumn, +) from sqlalchemy import Engine, func, select from sqlalchemy.orm import Session @@ -136,24 +145,48 @@ def _sync_resource(self, name: str, state: SyncState, *, limit: int) -> SyncResu self._save_state(state) pages = 0 try: - for page in self.client.iter_pages(spec.endpoint, limit=limit, start_page=start_page): - pages += 1 - with self.session_factory() as session: - current = session.get(SyncState, name) - if current is None: - raise PBSSyncError(f"sync state disappeared for {name}") - count = upsert_records(session, model, page.records) - metadata = { - "total_records": page.total_records, - "messages": page.messages, - "links": page.links, - "synced_at": page.metadata.get("synced_at"), - "page_limit": limit, - } - current.checkpoint(page.page, count, metadata) - session.commit() - state = current - logger.info("Synced %s page %s (%s records)", name, page.page, count) + initial_total = state.metadata_json.get("total_records") if can_resume else None + if isinstance(initial_total, bool) or not isinstance(initial_total, int): + initial_total = None + with Progress( + TextColumn("{task.description}"), + BarColumn(), + MofNCompleteColumn(), + TextColumn("records"), + TimeElapsedColumn(), + TimeRemainingColumn(), + disable=not sys.stderr.isatty(), + transient=True, + ) as progress: + task_id = progress.add_task( + name, + total=initial_total, + completed=state.records_written, + ) + for page in self.client.iter_pages( + spec.endpoint, + limit=limit, + start_page=start_page, + ): + pages += 1 + with self.session_factory() as session: + current = session.get(SyncState, name) + if current is None: + raise PBSSyncError(f"sync state disappeared for {name}") + count = upsert_records(session, model, page.records) + metadata = { + "total_records": page.total_records, + "messages": page.messages, + "links": page.links, + "synced_at": page.metadata.get("synced_at"), + "page_limit": limit, + } + current.checkpoint(page.page, count, metadata) + session.commit() + state = current + progress.update(task_id, total=page.total_records) + progress.advance(task_id, len(page.records)) + logger.info("Synced %s page %s (%s records)", name, page.page, count) if (total_records := _underreported_total(state)) is not None: raise PBSSyncError( f"{name} wrote {state.records_written:,} of the API-reported " diff --git a/tests/test_sync.py b/tests/test_sync.py index ced700d..508d012 100644 --- a/tests/test_sync.py +++ b/tests/test_sync.py @@ -85,6 +85,63 @@ def test_sync_fails_instead_of_completing_below_api_total(session_factory): assert state.completed_at is None +def test_sync_progress_tracks_page_totals(session_factory, monkeypatch): + observed = [] + + class RecordingProgress: + def __init__(self, *_columns, **_options): + self.total = None + self.completed = 0 + observed.append(self) + + def __enter__(self): + return self + + def __exit__(self, *_exc): + return False + + def add_task(self, _description, *, total=None, completed=0): + self.total = total + self.completed = completed + return 1 + + def update(self, _task_id, *, total=None): + if total is not None: + self.total = total + + def advance(self, _task_id, amount): + self.completed += amount + + sync_module = import_module("pbs_client.sync.orchestrator") + monkeypatch.setattr(sync_module, "Progress", RecordingProgress) + pages = [ + Page( + "/schedules", + 1, + 2, + [schedule(1, "2026-01-01"), schedule(2, "2026-01-01")], + {"total_records": 3}, + [], + ), + Page( + "/schedules", + 2, + 2, + [schedule(3, "2026-01-01")], + {"total_records": 3}, + [], + ), + ] + + SyncOrchestrator(FakeClient(pages), session_factory).run( + resource="Schedule", limit=2 + ) + + assert len(observed) == 1 + assert observed[0].total == 3 + assert observed[0].completed == 3 + + @pytest.mark.parametrize( ("resource", "records_written", "total_records", "exit_code"), [