Exercise recovery in datasource composition check

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,6 +95,8 @@ def main() -> int:
         ],
     )
     session_factory = sessionmaker(bind=engine)
+    bind_process_runtime_identity(_runtime_identity())
+    try:
         with session_factory() as session:
             principal = _principal()
             writer = tabular_snapshot_writer(registry)
@@ -100,7 +111,9 @@ def main() -> int:
                 or publisher is None
                 or runner is None
             ):
-            raise RuntimeError("Datasource composition capabilities are incomplete.")
+                raise RuntimeError(
+                    "Datasource composition capabilities are incomplete."
+                )
 
             origin = writer.create_snapshot(
                 session,
@@ -141,11 +154,15 @@ def main() -> int:
                 {"id": 2, "amount": 15},
             ]
             if result.status != "succeeded":
-            raise RuntimeError(f"Dataflow preview failed: {result.diagnostics}")
+                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.")
+                raise RuntimeError(
+                    "Dataflow lineage did not retain the datasource reference."
+                )
             pipeline = create_pipeline(
                 session,
                 tenant_id="tenant-1",
@@ -193,9 +210,7 @@ def main() -> int:
                 worker_id="composition-worker",
             )
             if worker_result["succeeded"] != 1:
-            raise RuntimeError(
+                raise RuntimeError(f"Dataflow worker failed: {worker_result!r}")
-                f"Dataflow worker failed: {worker_result!r}"
-            )
             completed = runner.get_run(
                 session,
                 principal,
@@ -205,14 +220,20 @@ def main() -> int:
                 raise RuntimeError("Dataflow run evidence disappeared.")
             published = completed
             if published.status != "succeeded":
-            raise RuntimeError(f"Dataflow publication failed: {published.error}")
+                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.")
+                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.")
+                raise RuntimeError(
+                    "Dataflow publication did not retain output references."
+                )
             output = catalogue.read_datasource(
                 session,
                 principal,
@@ -226,12 +247,15 @@ def main() -> int:
                 )
             if (
                 output.materialization is None
-            or output.materialization.ref != published.output_materialization_ref
+                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