feat: run dataflows through durable workers
This commit is contained in:
@@ -0,0 +1,243 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import timedelta
|
||||
import unittest
|
||||
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
|
||||
from govoplan_core.auth import ApiPrincipal
|
||||
from govoplan_core.core.access import (
|
||||
CAPABILITY_AUTH_AUTOMATION_PRINCIPAL_PROVIDER,
|
||||
PrincipalRef,
|
||||
)
|
||||
from govoplan_core.core.automation import AutomationPrincipalResolution
|
||||
from govoplan_core.core.dataflows import DataflowRunRequest
|
||||
from govoplan_core.db.base import Base, utcnow
|
||||
from govoplan_dataflow.backend.db.models import (
|
||||
DataflowPipeline,
|
||||
DataflowPipelineDeployment,
|
||||
DataflowPipelineRevision,
|
||||
DataflowRun,
|
||||
)
|
||||
from govoplan_dataflow.backend.run_worker import (
|
||||
dispatch_pending_runs,
|
||||
purge_expired_runs,
|
||||
)
|
||||
from govoplan_dataflow.backend.schemas import (
|
||||
GraphEdge,
|
||||
GraphNode,
|
||||
GraphPosition,
|
||||
PipelineCreateRequest,
|
||||
PipelineGraph,
|
||||
)
|
||||
from govoplan_dataflow.backend.service import (
|
||||
cancel_pipeline_run,
|
||||
create_pipeline,
|
||||
start_pipeline_run,
|
||||
)
|
||||
|
||||
|
||||
def _principal() -> ApiPrincipal:
|
||||
return ApiPrincipal(
|
||||
principal=PrincipalRef(
|
||||
account_id="account-1",
|
||||
membership_id="membership-1",
|
||||
tenant_id="tenant-1",
|
||||
scopes=frozenset({"dataflow:pipeline:run"}),
|
||||
),
|
||||
account=object(),
|
||||
user=object(),
|
||||
)
|
||||
|
||||
|
||||
def _graph() -> PipelineGraph:
|
||||
return PipelineGraph(
|
||||
nodes=[
|
||||
GraphNode(
|
||||
id="source",
|
||||
type="source.inline",
|
||||
label="Input",
|
||||
position=GraphPosition(x=0, y=0),
|
||||
config={
|
||||
"source_name": "input_rows",
|
||||
"rows": [{"id": 1}, {"id": 2}],
|
||||
},
|
||||
),
|
||||
GraphNode(
|
||||
id="output",
|
||||
type="output",
|
||||
label="Output",
|
||||
position=GraphPosition(x=200, y=0),
|
||||
config={},
|
||||
),
|
||||
],
|
||||
edges=[
|
||||
GraphEdge(
|
||||
id="source-output",
|
||||
source="source",
|
||||
target="output",
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
class _AutomationProvider:
|
||||
def resolve_automation_principal(self, _session, *, request):
|
||||
return AutomationPrincipalResolution(
|
||||
allowed=True,
|
||||
principal=_principal(),
|
||||
granted_scopes=request.grant_scopes,
|
||||
provenance={"status": "rechecked"},
|
||||
)
|
||||
|
||||
|
||||
class _Registry:
|
||||
def __init__(self) -> None:
|
||||
self.provider = _AutomationProvider()
|
||||
|
||||
def has_capability(self, name: str) -> bool:
|
||||
return name == CAPABILITY_AUTH_AUTOMATION_PRINCIPAL_PROVIDER
|
||||
|
||||
def capability(self, name: str):
|
||||
if not self.has_capability(name):
|
||||
raise KeyError(name)
|
||||
return self.provider
|
||||
|
||||
|
||||
class DataflowRunWorkerTests(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self.engine = create_engine("sqlite:///:memory:")
|
||||
Base.metadata.create_all(
|
||||
self.engine,
|
||||
tables=[
|
||||
DataflowPipeline.__table__,
|
||||
DataflowPipelineRevision.__table__,
|
||||
DataflowRun.__table__,
|
||||
DataflowPipelineDeployment.__table__,
|
||||
],
|
||||
)
|
||||
self.Session = sessionmaker(bind=self.engine)
|
||||
self.session: Session = self.Session()
|
||||
self.pipeline = create_pipeline(
|
||||
self.session,
|
||||
tenant_id="tenant-1",
|
||||
actor_id="account-1",
|
||||
payload=PipelineCreateRequest(
|
||||
name="Worker flow",
|
||||
status="active",
|
||||
graph=_graph(),
|
||||
editor_mode="graph",
|
||||
),
|
||||
)
|
||||
self.session.commit()
|
||||
|
||||
def tearDown(self) -> None:
|
||||
self.session.close()
|
||||
Base.metadata.drop_all(
|
||||
self.engine,
|
||||
tables=[
|
||||
DataflowPipelineDeployment.__table__,
|
||||
DataflowRun.__table__,
|
||||
DataflowPipelineRevision.__table__,
|
||||
DataflowPipeline.__table__,
|
||||
],
|
||||
)
|
||||
self.engine.dispose()
|
||||
|
||||
def _queue(self, key: str) -> DataflowRun:
|
||||
run, replayed = start_pipeline_run(
|
||||
self.session,
|
||||
tenant_id="tenant-1",
|
||||
actor_id="account-1",
|
||||
principal=_principal(),
|
||||
registry=_Registry(),
|
||||
request=DataflowRunRequest(
|
||||
pipeline_ref=f"pipeline:{self.pipeline.id}",
|
||||
revision=1,
|
||||
idempotency_key=key,
|
||||
),
|
||||
defer_execution=True,
|
||||
)
|
||||
self.session.commit()
|
||||
self.assertFalse(replayed)
|
||||
return run
|
||||
|
||||
def test_worker_claims_authorizes_and_executes_queued_run(self) -> None:
|
||||
run = self._queue("worker-success")
|
||||
self.assertEqual("queued", run.status)
|
||||
self.assertIsNone(run.started_at)
|
||||
|
||||
result = dispatch_pending_runs(
|
||||
self.session,
|
||||
registry=_Registry(),
|
||||
worker_id="test-worker",
|
||||
)
|
||||
|
||||
self.session.refresh(run)
|
||||
self.assertEqual(1, result["claimed"])
|
||||
self.assertEqual(1, result["succeeded"])
|
||||
self.assertEqual("succeeded", run.status)
|
||||
self.assertEqual(2, run.output_row_count)
|
||||
self.assertEqual(1, run.attempts)
|
||||
self.assertEqual(100, run.progress_percent)
|
||||
self.assertIsNone(run.worker_id)
|
||||
self.assertEqual(
|
||||
"rechecked",
|
||||
run.authorization_["last_resolution"]["status"],
|
||||
)
|
||||
|
||||
def test_cancelled_queued_run_is_never_claimed(self) -> None:
|
||||
run = self._queue("worker-cancel")
|
||||
cancel_pipeline_run(
|
||||
self.session,
|
||||
tenant_id="tenant-1",
|
||||
run_ref=f"dataflow-run:{run.id}",
|
||||
)
|
||||
self.session.commit()
|
||||
|
||||
result = dispatch_pending_runs(
|
||||
self.session,
|
||||
registry=_Registry(),
|
||||
)
|
||||
|
||||
self.assertEqual(0, result["claimed"])
|
||||
self.session.refresh(run)
|
||||
self.assertEqual("cancelled", run.status)
|
||||
|
||||
def test_expired_worker_lease_is_recovered(self) -> None:
|
||||
run = self._queue("worker-recover")
|
||||
run.status = "running"
|
||||
run.attempts = 1
|
||||
run.worker_id = "lost-worker"
|
||||
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)
|
||||
|
||||
def test_retention_purges_payload_but_keeps_run_evidence(self) -> None:
|
||||
run = self._queue("worker-retention")
|
||||
dispatch_pending_runs(self.session, registry=_Registry())
|
||||
run.retention_until = utcnow() - timedelta(seconds=1)
|
||||
self.session.commit()
|
||||
|
||||
result = purge_expired_runs(self.session)
|
||||
|
||||
self.session.refresh(run)
|
||||
self.assertEqual(1, result["purged"])
|
||||
self.assertEqual({}, run.request_)
|
||||
self.assertEqual("succeeded", run.status)
|
||||
self.assertEqual(2, run.output_row_count)
|
||||
self.assertIsNotNone(run.purged_at)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user