Exercise recovery in datasource composition check
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user