diff --git a/tools/checks/check-datasource-composition.py b/tools/checks/check-datasource-composition.py index 15dc778..be9f010 100644 --- a/tools/checks/check-datasource-composition.py +++ b/tools/checks/check-datasource-composition.py @@ -25,6 +25,12 @@ from govoplan_core.core.datasources import ( datasource_publication, ) from govoplan_core.core.modules import ModuleContext +from govoplan_core.core.recovery import RecoveryCheckpoint, RecoveryOperation +from govoplan_core.core.runtime_coordination import ( + DistributedLease, + RuntimeIdentity, + bind_process_runtime_identity, +) from govoplan_core.core.tabular_sources import ( TabularSnapshotInput, tabular_snapshot_writer, @@ -73,6 +79,9 @@ def main() -> int: Base.metadata.create_all( engine, tables=[ + DistributedLease.__table__, + RecoveryOperation.__table__, + RecoveryCheckpoint.__table__, ConnectorTabularSource.__table__, DatasourceRecord.__table__, DatasourcePayloadRecord.__table__, @@ -86,153 +95,168 @@ def main() -> int: ], ) session_factory = sessionmaker(bind=engine) - with session_factory() as session: - principal = _principal() - writer = tabular_snapshot_writer(registry) - lifecycle = datasource_lifecycle(registry) - catalogue = datasource_catalogue(registry) - publisher = datasource_publication(registry) - runner = dataflow_run_lifecycle(registry) - if ( - writer is None - or lifecycle is None - or catalogue is None - or publisher is None - or runner is None - ): - raise RuntimeError("Datasource composition capabilities are incomplete.") + bind_process_runtime_identity(_runtime_identity()) + try: + with session_factory() as session: + principal = _principal() + writer = tabular_snapshot_writer(registry) + lifecycle = datasource_lifecycle(registry) + catalogue = datasource_catalogue(registry) + publisher = datasource_publication(registry) + runner = dataflow_run_lifecycle(registry) + if ( + writer is None + or lifecycle is None + or catalogue is None + or publisher is None + or runner is None + ): + raise RuntimeError( + "Datasource composition capabilities are incomplete." + ) - origin = writer.create_snapshot( - session, - principal, - snapshot=TabularSnapshotInput( - name="Monthly cases", - source_name="connector_monthly_cases", - rows=( - {"id": 1, "amount": 5}, - {"id": 2, "amount": 15}, + origin = writer.create_snapshot( + session, + principal, + snapshot=TabularSnapshotInput( + name="Monthly cases", + source_name="connector_monthly_cases", + rows=( + {"id": 1, "amount": 5}, + {"id": 2, "amount": 15}, + ), ), - ), - ) - datasource = lifecycle.register_origin( - session, - principal, - origin_ref=origin.ref, - name="Monthly cases cache", - source_name="monthly_cases", - mode="cached", - ) - result = preview_pipeline( - session, - tenant_id="tenant-1", - actor_id="account-1", - payload=PipelinePreviewRequest( - graph=_graph( - datasource_ref=datasource.ref, - fingerprint=datasource.fingerprint, + ) + datasource = lifecycle.register_origin( + session, + principal, + origin_ref=origin.ref, + name="Monthly cases cache", + source_name="monthly_cases", + mode="cached", + ) + result = preview_pipeline( + session, + tenant_id="tenant-1", + actor_id="account-1", + payload=PipelinePreviewRequest( + graph=_graph( + datasource_ref=datasource.ref, + fingerprint=datasource.fingerprint, + ), + row_limit=100, ), - row_limit=100, - ), - principal=principal, - registry=registry, - ) - expected_rows = [ - {"id": 1, "amount": 5}, - {"id": 2, "amount": 15}, - ] - if result.status != "succeeded": - raise RuntimeError(f"Dataflow preview failed: {result.diagnostics}") - if result.rows != expected_rows: - raise RuntimeError(f"Unexpected Dataflow rows: {result.rows!r}") - if result.source_fingerprints[0]["source_ref"] != datasource.ref: - raise RuntimeError("Dataflow lineage did not retain the datasource reference.") - pipeline = create_pipeline( - session, - tenant_id="tenant-1", - actor_id="account-1", - payload=PipelineCreateRequest( - name="Monthly case output", - status="active", - graph=_graph( - datasource_ref=datasource.ref, - fingerprint=datasource.fingerprint, + principal=principal, + registry=registry, + ) + expected_rows = [ + {"id": 1, "amount": 5}, + {"id": 2, "amount": 15}, + ] + if result.status != "succeeded": + raise RuntimeError( + f"Dataflow preview failed: {result.diagnostics}" + ) + if result.rows != expected_rows: + raise RuntimeError(f"Unexpected Dataflow rows: {result.rows!r}") + if result.source_fingerprints[0]["source_ref"] != datasource.ref: + raise RuntimeError( + "Dataflow lineage did not retain the datasource reference." + ) + pipeline = create_pipeline( + session, + tenant_id="tenant-1", + actor_id="account-1", + payload=PipelineCreateRequest( + name="Monthly case output", + status="active", + graph=_graph( + datasource_ref=datasource.ref, + fingerprint=datasource.fingerprint, + ), + editor_mode="graph", ), - editor_mode="graph", - ), - ) - run_request = DataflowRunRequest( - pipeline_ref=f"pipeline:{pipeline.id}", - revision=1, - idempotency_key="composition-run-1", - publication=DataflowPublicationTarget( - name="Monthly case result", - source_name="monthly_case_result", - freeze=True, - frozen_label="Composition evidence", - ), - ) - published = runner.start_run( - session, - principal, - request=run_request, - ) - replayed = runner.start_run( - session, - principal, - request=run_request, - ) - if published.status != "queued": - raise RuntimeError( - f"Dataflow run was not queued: {published.status}" ) - worker = SqlDataflowRunWorker( - registry=_AutomationRegistry(registry, principal) - ) - worker_result = worker.dispatch_pending( - session, - worker_id="composition-worker", - ) - if worker_result["succeeded"] != 1: - raise RuntimeError( - f"Dataflow worker failed: {worker_result!r}" + run_request = DataflowRunRequest( + pipeline_ref=f"pipeline:{pipeline.id}", + revision=1, + idempotency_key="composition-run-1", + publication=DataflowPublicationTarget( + name="Monthly case result", + source_name="monthly_case_result", + freeze=True, + frozen_label="Composition evidence", + ), ) - completed = runner.get_run( - session, - principal, - run_ref=published.ref, - ) - if completed is None: - raise RuntimeError("Dataflow run evidence disappeared.") - published = completed - if published.status != "succeeded": - raise RuntimeError(f"Dataflow publication failed: {published.error}") - if replayed.ref != published.ref or not replayed.replayed: - raise RuntimeError("Dataflow run idempotency did not replay the prior run.") - if ( - not published.output_datasource_ref - or not published.output_materialization_ref - ): - raise RuntimeError("Dataflow publication did not retain output references.") - output = catalogue.read_datasource( - session, - principal, - request=DatasourceReadRequest( - datasource_ref=published.output_datasource_ref, - ), - ) - if list(output.rows) != expected_rows: - raise RuntimeError( - f"Unexpected published Dataflow rows: {list(output.rows)!r}" + published = runner.start_run( + session, + principal, + request=run_request, ) - if ( - output.materialization is None - or output.materialization.ref != published.output_materialization_ref - or output.materialization.frozen_at is None - ): - raise RuntimeError( - "Published Datasource materialization is not pinned and frozen." + replayed = runner.start_run( + session, + principal, + request=run_request, ) - engine.dispose() + if published.status != "queued": + raise RuntimeError( + f"Dataflow run was not queued: {published.status}" + ) + worker = SqlDataflowRunWorker( + registry=_AutomationRegistry(registry, principal) + ) + worker_result = worker.dispatch_pending( + session, + worker_id="composition-worker", + ) + if worker_result["succeeded"] != 1: + raise RuntimeError(f"Dataflow worker failed: {worker_result!r}") + completed = runner.get_run( + session, + principal, + run_ref=published.ref, + ) + if completed is None: + raise RuntimeError("Dataflow run evidence disappeared.") + published = completed + if published.status != "succeeded": + raise RuntimeError( + f"Dataflow publication failed: {published.error}" + ) + if replayed.ref != published.ref or not replayed.replayed: + raise RuntimeError( + "Dataflow run idempotency did not replay the prior run." + ) + if ( + not published.output_datasource_ref + or not published.output_materialization_ref + ): + raise RuntimeError( + "Dataflow publication did not retain output references." + ) + output = catalogue.read_datasource( + session, + principal, + request=DatasourceReadRequest( + datasource_ref=published.output_datasource_ref, + ), + ) + if list(output.rows) != expected_rows: + raise RuntimeError( + f"Unexpected published Dataflow rows: {list(output.rows)!r}" + ) + if ( + output.materialization is None + or output.materialization.ref + != published.output_materialization_ref + or output.materialization.frozen_at is None + ): + raise RuntimeError( + "Published Datasource materialization is not pinned and frozen." + ) + finally: + bind_process_runtime_identity(None) + engine.dispose() print( "Connector -> Datasources -> pinned Dataflow publication composition passed." ) @@ -252,6 +276,17 @@ class _AutomationProvider: ) +def _runtime_identity() -> RuntimeIdentity: + return RuntimeIdentity( + installation_id="datasource-composition-check", + node_id="composition-worker", + incarnation="composition-worker-incarnation", + role="worker", + software_version="test", + composition_hash="c" * 64, + ) + + class _AutomationRegistry: def __init__(self, registry, principal: ApiPrincipal) -> None: self.registry = registry