GovOPlaN/govoplan-dataflow
2
0
Fork 0
Code Issues 2 Pull Requests Actions Packages Projects Releases Wiki Activity

Fence and reconcile Dataflow runs

Browse Source
This commit is contained in:
zemion
2026-08-03 06:09:53 +02:00
parent e1c092ece7
commit 3fa7a29f48
13 changed files with 1314 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_run_worker.py
+186 -5
View File
@@ -3,7 +3,7 @@ from __future__ import annotations
from datetime import timedelta
import unittest
from sqlalchemy import create_engine
from sqlalchemy import create_engine, select
from sqlalchemy.orm import Session, sessionmaker
from govoplan_core.auth import ApiPrincipal
@@ -12,7 +12,26 @@ from govoplan_core.core.access import (
PrincipalRef,
)
from govoplan_core.core.automation import AutomationPrincipalResolution
from govoplan_core.core.dataflows import DataflowRunRequest
from govoplan_core.core.dataflows import (
DataflowPublicationTarget,
DataflowRunRequest,
)
from govoplan_core.core.datasources import (
CAPABILITY_DATASOURCE_PUBLICATION,
DatasourceDescriptor,
DatasourceMaterialization,
DatasourcePublicationResult,
)
from govoplan_core.core.recovery import (
RecoveryCheckpoint,
RecoveryOperation,
RecoveryStatus,
)
from govoplan_core.core.runtime_coordination import (
DistributedLease,
RuntimeIdentity,
bind_process_runtime_identity,
)
from govoplan_core.db.base import Base, utcnow
from govoplan_dataflow.backend.db.models import (
DataflowPipeline,
@@ -24,6 +43,7 @@ from govoplan_dataflow.backend.run_worker import (
dispatch_pending_runs,
purge_expired_runs,
)
from govoplan_dataflow.backend.recovery import begin_dataflow_run_recovery
from govoplan_dataflow.backend.schemas import (
GraphEdge,
GraphNode,
@@ -51,6 +71,17 @@ def _principal() -> ApiPrincipal:
)
def _runtime_identity() -> RuntimeIdentity:
return RuntimeIdentity(
installation_id="dataflow-worker-tests",
node_id="worker-node",
incarnation="worker-incarnation",
role="worker",
software_version="test",
composition_hash="b" * 64,
)
def _graph() -> PipelineGraph:
return PipelineGraph(
nodes=[
@@ -92,16 +123,54 @@ class _AutomationProvider:
)
class _Registry:
class _PublicationProvider:
def __init__(self) -> None:
self.requests = []
def publish_rows(self, _session, _principal, *, request):
self.requests.append(request)
descriptor = DatasourceDescriptor(
ref="datasource:worker-output",
source_name="worker_output",
name="Worker output",
kind="custom",
mode="static",
shape="tabular",
fingerprint="worker-output-fingerprint",
)
return DatasourcePublicationResult(
ref="publication:worker-output",
status="published",
datasource=descriptor,
materialization=DatasourceMaterialization(
ref="materialization:worker-output",
datasource_ref=descriptor.ref,
revision=1,
state="published",
fingerprint=descriptor.fingerprint,
),
)
class _Registry:
def __init__(
self,
publication_provider: _PublicationProvider | None = None,
) -> None:
self.provider = _AutomationProvider()
self.publication_provider = publication_provider
def has_capability(self, name: str) -> bool:
return name == CAPABILITY_AUTH_AUTOMATION_PRINCIPAL_PROVIDER
return name == CAPABILITY_AUTH_AUTOMATION_PRINCIPAL_PROVIDER or (
name == CAPABILITY_DATASOURCE_PUBLICATION
and self.publication_provider is not None
)
def capability(self, name: str):
if not self.has_capability(name):
raise KeyError(name)
if name == CAPABILITY_DATASOURCE_PUBLICATION:
return self.publication_provider
return self.provider
@@ -111,6 +180,9 @@ class DataflowRunWorkerTests(unittest.TestCase):
Base.metadata.create_all(
self.engine,
tables=[
DistributedLease.__table__,
RecoveryOperation.__table__,
RecoveryCheckpoint.__table__,
DataflowPipeline.__table__,
DataflowPipelineRevision.__table__,
DataflowRun.__table__,
@@ -119,6 +191,7 @@ class DataflowRunWorkerTests(unittest.TestCase):
)
self.Session = sessionmaker(bind=self.engine)
self.session: Session = self.Session()
bind_process_runtime_identity(_runtime_identity())
self.pipeline = create_pipeline(
self.session,
tenant_id="tenant-1",
@@ -133,6 +206,7 @@ class DataflowRunWorkerTests(unittest.TestCase):
self.session.commit()
def tearDown(self) -> None:
bind_process_runtime_identity(None)
self.session.close()
Base.metadata.drop_all(
self.engine,
@@ -141,11 +215,19 @@ class DataflowRunWorkerTests(unittest.TestCase):
DataflowRun.__table__,
DataflowPipelineRevision.__table__,
DataflowPipeline.__table__,
RecoveryCheckpoint.__table__,
RecoveryOperation.__table__,
DistributedLease.__table__,
],
)
self.engine.dispose()
def _queue(self, key: str) -> DataflowRun:
def _queue(
self,
key: str,
*,
publication: bool = False,
) -> DataflowRun:
run, replayed = start_pipeline_run(
self.session,
tenant_id="tenant-1",
@@ -156,6 +238,14 @@ class DataflowRunWorkerTests(unittest.TestCase):
pipeline_ref=f"pipeline:{self.pipeline.id}",
revision=1,
idempotency_key=key,
publication=(
DataflowPublicationTarget(
name="Worker output",
source_name="worker_output",
)
if publication
else None
),
),
defer_execution=True,
)
@@ -238,6 +328,97 @@ class DataflowRunWorkerTests(unittest.TestCase):
self.assertEqual(2, run.output_row_count)
self.assertIsNotNone(run.purged_at)
def test_stale_prepublication_attempt_is_safely_retried(self) -> None:
provider = _PublicationProvider()
registry = _Registry(provider)
run = self._queue("stale-before-publication", publication=True)
run.status = "running"
run.attempts = 1
run.worker_id = "lost-worker"
run.lease_expires_at = utcnow() + timedelta(minutes=5)
self.session.commit()
begin_dataflow_run_recovery(
self.session,
run=run,
lease_ttl_seconds=120,
)
self.session.commit()
lease = self.session.scalar(
select(DistributedLease).where(
DistributedLease.resource_key == f"dataflow:run:{run.id}"
)
)
assert lease is not None
lease.expires_at = utcnow() - timedelta(minutes=1)
run.lease_expires_at = utcnow() - timedelta(minutes=1)
self.session.commit()
result = dispatch_pending_runs(self.session, registry=registry)
self.session.refresh(run)
self.assertEqual(1, result["recovered"])
self.assertEqual("succeeded", run.status)
self.assertEqual(2, run.attempts)
self.assertEqual(1, len(provider.requests))
statuses = set(
self.session.scalars(
select(RecoveryOperation.status).where(
RecoveryOperation.resource_id == run.id
)
)
)
self.assertEqual(
{RecoveryStatus.RECOVERED.value, RecoveryStatus.SUCCEEDED.value},
statuses,
)
def test_stale_publication_attempt_is_not_retried_blindly(self) -> None:
provider = _PublicationProvider()
registry = _Registry(provider)
run = self._queue("stale-after-publication", publication=True)
run.status = "running"
run.attempts = 1
run.worker_id = "lost-worker"
run.lease_expires_at = utcnow() + timedelta(minutes=5)
run.output_row_count = 1
self.session.commit()
recovery = begin_dataflow_run_recovery(
self.session,
run=run,
lease_ttl_seconds=120,
)
recovery.prepare_publication(
self.session,
run=run,
rows=({"id": 1},),
)
lease = self.session.scalar(
select(DistributedLease).where(
DistributedLease.resource_key == f"dataflow:run:{run.id}"
)
)
assert lease is not None
lease.expires_at = utcnow() - timedelta(minutes=1)
run.lease_expires_at = utcnow() - timedelta(minutes=1)
self.session.commit()
result = dispatch_pending_runs(self.session, registry=registry)
self.session.refresh(run)
self.assertEqual(1, result["recovered"])
self.assertEqual(1, result["outcome_unknown"])
self.assertEqual(0, result["claimed"])
self.assertEqual("outcome_unknown", run.status)
self.assertEqual(1, run.attempts)
self.assertEqual([], provider.requests)
operation = self.session.scalar(
select(RecoveryOperation).where(
RecoveryOperation.resource_id == run.id
)
)
assert operation is not None
self.assertEqual(RecoveryStatus.OUTCOME_UNKNOWN.value, operation.status)
if __name__ == "__main__":
unittest.main()
tests/test_service.py
+160
View File
@@ -17,6 +17,18 @@ from govoplan_core.core.datasources import (
DatasourceMaterialization,
DatasourcePublicationResult,
)
from govoplan_core.core.recovery import (
RecoveryCheckpoint,
RecoveryMode,
RecoveryOperation,
RecoveryStatus,
verify_recovery_evidence_chain,
)
from govoplan_core.core.runtime_coordination import (
DistributedLease,
RuntimeIdentity,
bind_process_runtime_identity,
)
from govoplan_core.db.base import Base
from govoplan_dataflow.backend.backends.duckdb import DuckDbExecutionBackend
from govoplan_dataflow.backend.db.models import (
@@ -99,6 +111,17 @@ def principal(tenant_id: str = "tenant-1") -> ApiPrincipal:
)
def runtime_identity() -> RuntimeIdentity:
return RuntimeIdentity(
installation_id="dataflow-service-tests",
node_id="service-node",
incarnation="service-incarnation",
role="api",
software_version="test",
composition_hash="a" * 64,
)
class FakePublicationProvider:
def __init__(self) -> None:
self.requests = []
@@ -128,6 +151,25 @@ class FakePublicationProvider:
)
class FailingPublicationProvider(FakePublicationProvider):
def publish_rows(self, _session, _principal, *, request):
self.requests.append(request)
raise RuntimeError("connection closed after dispatch")
class TamperingPublicationProvider(FakePublicationProvider):
def publish_rows(self, session, principal, *, request):
checkpoint = session.scalar(
select(RecoveryCheckpoint)
.order_by(RecoveryCheckpoint.sequence)
.limit(1)
)
assert checkpoint is not None
checkpoint.summary = "tampered"
session.commit()
return super().publish_rows(session, principal, request=request)
class FakeRegistry:
def __init__(self, publication_provider: FakePublicationProvider) -> None:
self.publication_provider = publication_provider
@@ -159,6 +201,9 @@ class DataflowServiceTests(unittest.TestCase):
Base.metadata.create_all(
self.engine,
tables=[
DistributedLease.__table__,
RecoveryOperation.__table__,
RecoveryCheckpoint.__table__,
DataflowPipeline.__table__,
DataflowPipelineRevision.__table__,
DataflowRun.__table__,
@@ -166,8 +211,10 @@ class DataflowServiceTests(unittest.TestCase):
)
self.Session = sessionmaker(bind=self.engine)
self.session: Session = self.Session()
bind_process_runtime_identity(runtime_identity())
def tearDown(self) -> None:
bind_process_runtime_identity(None)
self.session.close()
Base.metadata.drop_all(
self.engine,
@@ -175,6 +222,9 @@ class DataflowServiceTests(unittest.TestCase):
DataflowRun.__table__,
DataflowPipelineRevision.__table__,
DataflowPipeline.__table__,
RecoveryCheckpoint.__table__,
RecoveryOperation.__table__,
DistributedLease.__table__,
],
)
self.engine.dispose()
@@ -435,6 +485,26 @@ class DataflowServiceTests(unittest.TestCase):
list(publication_provider.requests[0].rows),
)
self.assertEqual(1, len(publication_provider.requests))
operation = self.session.scalar(
select(RecoveryOperation).where(
RecoveryOperation.resource_id == first.id
)
)
assert operation is not None
self.assertEqual(RecoveryMode.FORWARD_RECOVERY.value, operation.mode)
self.assertEqual(RecoveryStatus.SUCCEEDED.value, operation.status)
self.assertTrue(
verify_recovery_evidence_chain(self.session, operation.id)
)
checkpoint_kinds = list(
self.session.scalars(
select(RecoveryCheckpoint.kind)
.where(RecoveryCheckpoint.operation_id == operation.id)
.order_by(RecoveryCheckpoint.sequence)
)
)
self.assertIn("output-publication-dispatch", checkpoint_kinds)
self.assertIn("verified-success", checkpoint_kinds)
def test_run_idempotency_key_rejects_changed_parameters(self) -> None:
pipeline = self._create()
@@ -490,6 +560,96 @@ class DataflowServiceTests(unittest.TestCase):
self.assertEqual("failed", run.status)
self.assertIn("Datasources publication capability", run.error)
operation = self.session.scalar(
select(RecoveryOperation).where(
RecoveryOperation.resource_id == run.id
)
)
assert operation is not None
self.assertEqual(RecoveryStatus.RECOVERED.value, operation.status)
def test_provider_failure_after_dispatch_requires_reconciliation(self) -> None:
pipeline = self._create()
provider = FailingPublicationProvider()
request = DataflowRunRequest(
pipeline_ref=f"pipeline:{pipeline.id}",
revision=1,
idempotency_key="uncertain-publication",
publication=DataflowPublicationTarget(
name="Uncertain output",
source_name="uncertain_output",
),
)
run, replayed = start_pipeline_run(
self.session,
tenant_id="tenant-1",
actor_id="user-1",
principal=principal(),
registry=FakeRegistry(provider),
request=request,
)
replay, second_replayed = start_pipeline_run(
self.session,
tenant_id="tenant-1",
actor_id="user-1",
principal=principal(),
registry=FakeRegistry(provider),
request=request,
)
self.assertFalse(replayed)
self.assertTrue(second_replayed)
self.assertEqual(run.id, replay.id)
self.assertEqual("outcome_unknown", run.status)
self.assertEqual(1, len(provider.requests))
operation = self.session.scalar(
select(RecoveryOperation).where(
RecoveryOperation.resource_id == run.id
)
)
assert operation is not None
self.assertEqual(RecoveryStatus.OUTCOME_UNKNOWN.value, operation.status)
def test_tampered_recovery_chain_prevents_verified_publication(self) -> None:
pipeline = self._create()
provider = TamperingPublicationProvider()
with self.assertRaises(DataflowConflictError):
start_pipeline_run(
self.session,
tenant_id="tenant-1",
actor_id="user-1",
principal=principal(),
registry=FakeRegistry(provider),
request=DataflowRunRequest(
pipeline_ref=f"pipeline:{pipeline.id}",
revision=1,
idempotency_key="tampered-publication",
publication=DataflowPublicationTarget(
name="Tampered output",
source_name="tampered_output",
),
),
)
run = self.session.scalar(
select(DataflowRun).where(
DataflowRun.idempotency_key == "tampered-publication"
)
)
assert run is not None
operation = self.session.scalar(
select(RecoveryOperation).where(
RecoveryOperation.resource_id == run.id
)
)
assert operation is not None
self.assertEqual("outcome_unknown", run.status)
self.assertIsNone(run.output_publication_ref)
self.assertFalse(
verify_recovery_evidence_chain(self.session, operation.id)
)
if __name__ == "__main__":