From 5af02933ef27567c8a1ed9229e94a6dd1ed34f37 Mon Sep 17 00:00:00 2001 From: Albrecht Degering Date: Thu, 30 Jul 2026 02:33:02 +0200 Subject: [PATCH] feat: run dataflows through durable workers --- README.md | 28 +- src/govoplan_dataflow/backend/db/__init__.py | 14 +- src/govoplan_dataflow/backend/db/models.py | 141 +++- src/govoplan_dataflow/backend/executor.py | 2 + src/govoplan_dataflow/backend/manifest.py | 21 + .../f6c2a9d4e7b1_v0114_production_runs.py | 258 +++++++ src/govoplan_dataflow/backend/router.py | 107 ++- src/govoplan_dataflow/backend/run_worker.py | 631 ++++++++++++++++++ src/govoplan_dataflow/backend/schemas.py | 87 ++- src/govoplan_dataflow/backend/service.py | 592 +++++++++++++++- src/govoplan_dataflow/backend/triggers.py | 32 +- tests/test_manifest.py | 1 + tests/test_migrations.py | 3 +- tests/test_run_worker.py | 243 +++++++ webui/src/api/dataflow.ts | 75 ++- webui/src/features/dataflow/DataflowPage.tsx | 138 +++- 16 files changed, 2306 insertions(+), 67 deletions(-) create mode 100644 src/govoplan_dataflow/backend/migrations/versions/f6c2a9d4e7b1_v0114_production_runs.py create mode 100644 src/govoplan_dataflow/backend/run_worker.py create mode 100644 tests/test_run_worker.py diff --git a/README.md b/README.md index 3a1e1b9..370c510 100644 --- a/README.md +++ b/README.md @@ -91,11 +91,20 @@ manual-refresh preview modes. Saved revisions can also be started through the versioned `dataflow.runLifecycle` capability or the Run dialog. The first runner is -synchronous and bounded. It records an idempotent terminal run, pins the -pipeline revision, and can atomically publish a complete result through -`datasources.publication`. Publication is rejected when a source or result was -truncated; larger asynchronous and artifact-backed execution belongs to a -subsequent runner slice. +durable and worker-backed: API requests only enqueue idempotent work, while +the `dataflow.runWorker` capability claims leased runs on the dedicated Celery +queue. Access rechecks the persisted least-privilege authorization grant before +each attempt. Runs expose progress, cancellation, bounded exponential retries, +tenant queue/concurrency quotas, notifications, operational metrics, and +evidence retention. Expired retention payloads are purged while hashes, row +counts, outcomes, and publication references remain as audit evidence. + +Development runs may use the bounded reference backend. Staging and production +runs require the short-lived isolated DuckDB process. A revision must be +promoted from development to staging and then from staging to production before +it can run in those environments. Complete results can be published atomically +through `datasources.publication`; publication is rejected when a source or +result was truncated. ## Governed Definitions And Automation @@ -108,10 +117,11 @@ execution, reuse, inheritance, or automation permissions. Complete active flows support explicit user/API starts, administrative backfills, one-time schedules, interval schedules, and exact-match platform -events. Trigger deliveries are durable and idempotent. They pin the pipeline -revision and a least-privilege scope grant, then ask Access to rebuild the -owner's current principal before every run. Revoked memberships or reduced -permissions block the delivery before source access or output publication. +events. Trigger deliveries are durable and idempotent. They enqueue the same +worker-backed run contract, pin the pipeline revision and a least-privilege +scope grant, then ask Access to rebuild the owner's current principal before +both delivery and execution. Revoked memberships or reduced permissions block +the run before source access or output publication. Confidential and restricted events are not accepted through the direct ingress; those require Core's transactional event bridge. diff --git a/src/govoplan_dataflow/backend/db/__init__.py b/src/govoplan_dataflow/backend/db/__init__.py index 3fa4bf0..703ef19 100644 --- a/src/govoplan_dataflow/backend/db/__init__.py +++ b/src/govoplan_dataflow/backend/db/__init__.py @@ -1,5 +1,15 @@ from __future__ import annotations -from govoplan_dataflow.backend.db.models import DataflowPipeline, DataflowPipelineRevision, DataflowRun +from govoplan_dataflow.backend.db.models import ( + DataflowPipeline, + DataflowPipelineDeployment, + DataflowPipelineRevision, + DataflowRun, +) -__all__ = ["DataflowPipeline", "DataflowPipelineRevision", "DataflowRun"] +__all__ = [ + "DataflowPipeline", + "DataflowPipelineDeployment", + "DataflowPipelineRevision", + "DataflowRun", +] diff --git a/src/govoplan_dataflow/backend/db/models.py b/src/govoplan_dataflow/backend/db/models.py index 5037ec3..75a41e9 100644 --- a/src/govoplan_dataflow/backend/db/models.py +++ b/src/govoplan_dataflow/backend/db/models.py @@ -111,6 +111,11 @@ class DataflowPipeline(Base, TimestampMixin): cascade="all, delete-orphan", order_by="DataflowRun.created_at", ) + deployments: Mapped[list["DataflowPipelineDeployment"]] = relationship( + back_populates="pipeline", + cascade="all, delete-orphan", + order_by="DataflowPipelineDeployment.environment", + ) triggers: Mapped[list["DataflowTrigger"]] = relationship( back_populates="pipeline", cascade="all, delete-orphan", @@ -159,6 +164,12 @@ class DataflowRun(Base, TimestampMixin): ), Index("ix_dataflow_runs_tenant_status", "tenant_id", "status"), Index("ix_dataflow_runs_pipeline_created", "pipeline_id", "created_at"), + Index( + "ix_dataflow_runs_worker_queue", + "status", + "available_at", + "created_at", + ), ) id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_uuid) @@ -175,6 +186,18 @@ class DataflowRun(Base, TimestampMixin): ) run_type: Mapped[str] = mapped_column(String(30), default="preview", nullable=False, index=True) status: Mapped[str] = mapped_column(String(30), nullable=False, index=True) + execution_backend: Mapped[str] = mapped_column( + String(30), + default="auto", + nullable=False, + index=True, + ) + environment: Mapped[str] = mapped_column( + String(30), + default="development", + nullable=False, + index=True, + ) executor_version: Mapped[str] = mapped_column(String(40), nullable=False) definition_hash: Mapped[str] = mapped_column(String(64), nullable=False) idempotency_key: Mapped[str | None] = mapped_column( @@ -232,7 +255,73 @@ class DataflowRun(Base, TimestampMixin): String(500), nullable=True, ) - started_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False) + attempts: Mapped[int] = mapped_column(Integer, default=0, nullable=False) + max_attempts: Mapped[int] = mapped_column(Integer, default=3, nullable=False) + queued_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), + nullable=True, + ) + available_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), + nullable=True, + index=True, + ) + claimed_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), + nullable=True, + ) + lease_expires_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), + nullable=True, + index=True, + ) + heartbeat_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), + nullable=True, + ) + worker_id: Mapped[str | None] = mapped_column( + String(255), + nullable=True, + index=True, + ) + cancellation_requested_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), + nullable=True, + ) + progress_percent: Mapped[int] = mapped_column( + Integer, + default=0, + nullable=False, + ) + progress_phase: Mapped[str] = mapped_column( + String(80), + default="queued", + nullable=False, + ) + retention_until: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), + nullable=True, + index=True, + ) + purged_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), + nullable=True, + ) + authorization_: Mapped[dict[str, Any]] = mapped_column( + "authorization", + JSON, + default=dict, + nullable=False, + ) + resource_budget: Mapped[dict[str, Any]] = mapped_column( + JSON, + default=dict, + nullable=False, + ) + started_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), + nullable=True, + ) finished_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) error: Mapped[str | None] = mapped_column(Text) created_by: Mapped[str | None] = mapped_column(String(255), nullable=True, index=True) @@ -240,6 +329,56 @@ class DataflowRun(Base, TimestampMixin): pipeline: Mapped[DataflowPipeline] = relationship(back_populates="runs") +class DataflowPipelineDeployment(Base, TimestampMixin): + __tablename__ = "dataflow_pipeline_deployments" + __table_args__ = ( + UniqueConstraint( + "tenant_id", + "pipeline_id", + "environment", + name="uq_dataflow_pipeline_deployment_environment", + ), + Index( + "ix_dataflow_deployments_tenant_environment", + "tenant_id", + "environment", + ), + ) + + id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_uuid) + tenant_id: Mapped[str] = mapped_column(String(36), nullable=False, index=True) + pipeline_id: Mapped[str] = mapped_column( + ForeignKey("dataflow_pipelines.id", ondelete="CASCADE"), + nullable=False, + index=True, + ) + pipeline_revision_id: Mapped[str] = mapped_column( + ForeignKey("dataflow_pipeline_revisions.id", ondelete="RESTRICT"), + nullable=False, + index=True, + ) + environment: Mapped[str] = mapped_column(String(30), nullable=False, index=True) + source_environment: Mapped[str] = mapped_column(String(30), nullable=False) + status: Mapped[str] = mapped_column( + String(30), + default="active", + nullable=False, + index=True, + ) + provenance: Mapped[dict[str, Any]] = mapped_column( + JSON, + default=dict, + nullable=False, + ) + promoted_by: Mapped[str | None] = mapped_column( + String(255), + nullable=True, + index=True, + ) + + pipeline: Mapped[DataflowPipeline] = relationship(back_populates="deployments") + + class DataflowTrigger(Base, TimestampMixin): __tablename__ = "dataflow_triggers" __table_args__ = ( diff --git a/src/govoplan_dataflow/backend/executor.py b/src/govoplan_dataflow/backend/executor.py index ca5951f..4c29ca2 100644 --- a/src/govoplan_dataflow/backend/executor.py +++ b/src/govoplan_dataflow/backend/executor.py @@ -48,6 +48,7 @@ class PipelineExecutionError(RuntimeError): input_row_count: int = 0, diagnostics: tuple[DataflowDiagnostic, ...] = (), node_preview: NodePreviewResult | None = None, + retryable: bool = False, ) -> None: super().__init__(message) self.node_id = node_id @@ -56,6 +57,7 @@ class PipelineExecutionError(RuntimeError): self.input_row_count = input_row_count self.diagnostics = diagnostics self.node_preview = node_preview + self.retryable = retryable @dataclass(frozen=True) diff --git a/src/govoplan_dataflow/backend/manifest.py b/src/govoplan_dataflow/backend/manifest.py index 1b0be84..3dc33fb 100644 --- a/src/govoplan_dataflow/backend/manifest.py +++ b/src/govoplan_dataflow/backend/manifest.py @@ -12,6 +12,7 @@ from govoplan_core.core.access import ( ) from govoplan_core.core.dataflows import ( CAPABILITY_DATAFLOW_RUN_LIFECYCLE, + CAPABILITY_DATAFLOW_RUN_WORKER, CAPABILITY_DATAFLOW_TRIGGER_DISPATCHER, ) from govoplan_core.core.modules import ( @@ -35,6 +36,9 @@ from govoplan_core.core.datasources import ( from govoplan_core.core.policy import ( CAPABILITY_POLICY_DEFINITION_GOVERNANCE, ) +from govoplan_core.core.notifications import ( + CAPABILITY_NOTIFICATIONS_DISPATCH, +) from govoplan_core.core.references import CAPABILITY_ACCESS_REFERENCE_OPTIONS from govoplan_core.core.search import SearchSourceProviderRegistration from govoplan_core.core.views import ViewSurface @@ -190,6 +194,12 @@ def _run_provider(context: ModuleContext): return SqlDataflowRunLifecycleProvider(registry=context.registry) +def _run_worker(context: ModuleContext): + from govoplan_dataflow.backend.run_worker import SqlDataflowRunWorker + + return SqlDataflowRunWorker(registry=context.registry) + + def _trigger_provider(context: ModuleContext): from govoplan_dataflow.backend.triggers import ( SqlDataflowTriggerDispatcher, @@ -261,12 +271,14 @@ manifest = ModuleManifest( CAPABILITY_DATASOURCE_LIFECYCLE, CAPABILITY_DATASOURCE_PUBLICATION, CAPABILITY_AUTH_AUTOMATION_PRINCIPAL_PROVIDER, + CAPABILITY_NOTIFICATIONS_DISPATCH, CAPABILITY_POLICY_DEFINITION_GOVERNANCE, ), provides_interfaces=( ModuleInterfaceProvider(name="dataflow.pipeline_catalog", version=MODULE_VERSION), ModuleInterfaceProvider(name="dataflow.pipeline_preview", version=MODULE_VERSION), ModuleInterfaceProvider(name="dataflow.run_lifecycle", version=MODULE_VERSION), + ModuleInterfaceProvider(name="dataflow.run_worker", version=MODULE_VERSION), ModuleInterfaceProvider(name="dataflow.dataset_output", version=MODULE_VERSION), ModuleInterfaceProvider( name="dataflow.trigger_dispatcher", @@ -310,6 +322,12 @@ manifest = ModuleManifest( version_max_exclusive="1.0.0", optional=True, ), + ModuleInterfaceRequirement( + name=CAPABILITY_NOTIFICATIONS_DISPATCH, + version_min="0.1.0", + version_max_exclusive="1.0.0", + optional=True, + ), ModuleInterfaceRequirement( name="search.source", version_min="1.0.0", @@ -361,6 +379,7 @@ manifest = ModuleManifest( route_factory=_dataflow_router, capability_factories={ CAPABILITY_DATAFLOW_RUN_LIFECYCLE: _run_provider, + CAPABILITY_DATAFLOW_RUN_WORKER: _run_worker, CAPABILITY_DATAFLOW_TRIGGER_DISPATCHER: _trigger_provider, }, search_sources=( @@ -379,6 +398,7 @@ manifest = ModuleManifest( dataflow_models.DataflowTriggerDelivery, dataflow_models.DataflowTrigger, dataflow_models.DataflowRun, + dataflow_models.DataflowPipelineDeployment, dataflow_models.DataflowPipelineRevision, dataflow_models.DataflowPipeline, label="Dataflow", @@ -392,6 +412,7 @@ manifest = ModuleManifest( persistent_table_uninstall_guard( dataflow_models.DataflowPipeline, dataflow_models.DataflowPipelineRevision, + dataflow_models.DataflowPipelineDeployment, dataflow_models.DataflowRun, dataflow_models.DataflowTrigger, dataflow_models.DataflowTriggerDelivery, diff --git a/src/govoplan_dataflow/backend/migrations/versions/f6c2a9d4e7b1_v0114_production_runs.py b/src/govoplan_dataflow/backend/migrations/versions/f6c2a9d4e7b1_v0114_production_runs.py new file mode 100644 index 0000000..c3f561a --- /dev/null +++ b/src/govoplan_dataflow/backend/migrations/versions/f6c2a9d4e7b1_v0114_production_runs.py @@ -0,0 +1,258 @@ +"""v0.1.14 production Dataflow runs and deployments + +Revision ID: f6c2a9d4e7b1 +Revises: e2f4a8c9d1b7 +Create Date: 2026-07-30 00:00:00.000000 +""" +from __future__ import annotations + +from alembic import op +import sqlalchemy as sa + + +revision = "f6c2a9d4e7b1" +down_revision = "e2f4a8c9d1b7" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + with op.batch_alter_table("dataflow_runs") as batch_op: + batch_op.alter_column( + "started_at", + existing_type=sa.DateTime(timezone=True), + nullable=True, + ) + batch_op.add_column( + sa.Column( + "execution_backend", + sa.String(length=30), + nullable=False, + server_default="auto", + ) + ) + batch_op.add_column( + sa.Column( + "environment", + sa.String(length=30), + nullable=False, + server_default="development", + ) + ) + batch_op.add_column( + sa.Column( + "attempts", + sa.Integer(), + nullable=False, + server_default="0", + ) + ) + batch_op.add_column( + sa.Column( + "max_attempts", + sa.Integer(), + nullable=False, + server_default="3", + ) + ) + batch_op.add_column( + sa.Column("queued_at", sa.DateTime(timezone=True), nullable=True) + ) + batch_op.add_column( + sa.Column("available_at", sa.DateTime(timezone=True), nullable=True) + ) + batch_op.add_column( + sa.Column("claimed_at", sa.DateTime(timezone=True), nullable=True) + ) + batch_op.add_column( + sa.Column( + "lease_expires_at", + sa.DateTime(timezone=True), + nullable=True, + ) + ) + batch_op.add_column( + sa.Column("heartbeat_at", sa.DateTime(timezone=True), nullable=True) + ) + batch_op.add_column( + sa.Column("worker_id", sa.String(length=255), nullable=True) + ) + batch_op.add_column( + sa.Column( + "cancellation_requested_at", + sa.DateTime(timezone=True), + nullable=True, + ) + ) + batch_op.add_column( + sa.Column( + "progress_percent", + sa.Integer(), + nullable=False, + server_default="0", + ) + ) + batch_op.add_column( + sa.Column( + "progress_phase", + sa.String(length=80), + nullable=False, + server_default="queued", + ) + ) + batch_op.add_column( + sa.Column( + "retention_until", + sa.DateTime(timezone=True), + nullable=True, + ) + ) + batch_op.add_column( + sa.Column("purged_at", sa.DateTime(timezone=True), nullable=True) + ) + batch_op.add_column( + sa.Column( + "authorization", + sa.JSON(), + nullable=False, + server_default=sa.text("'{}'"), + ) + ) + batch_op.add_column( + sa.Column( + "resource_budget", + sa.JSON(), + nullable=False, + server_default=sa.text("'{}'"), + ) + ) + op.execute( + sa.text( + "UPDATE dataflow_runs " + "SET queued_at = created_at, available_at = created_at " + "WHERE status IN ('queued', 'retrying')" + ) + ) + for column in ( + "execution_backend", + "environment", + "available_at", + "lease_expires_at", + "worker_id", + "retention_until", + ): + op.create_index( + op.f(f"ix_dataflow_runs_{column}"), + "dataflow_runs", + [column], + unique=False, + ) + op.create_index( + "ix_dataflow_runs_worker_queue", + "dataflow_runs", + ["status", "available_at", "created_at"], + unique=False, + ) + + op.create_table( + "dataflow_pipeline_deployments", + sa.Column("id", sa.String(length=36), nullable=False), + sa.Column("tenant_id", sa.String(length=36), nullable=False), + sa.Column("pipeline_id", sa.String(length=36), nullable=False), + sa.Column( + "pipeline_revision_id", + sa.String(length=36), + nullable=False, + ), + sa.Column("environment", sa.String(length=30), nullable=False), + sa.Column("source_environment", sa.String(length=30), nullable=False), + sa.Column("status", sa.String(length=30), nullable=False), + sa.Column("provenance", sa.JSON(), nullable=False), + sa.Column("promoted_by", sa.String(length=255), nullable=True), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + sa.ForeignKeyConstraint( + ["pipeline_id"], + ["dataflow_pipelines.id"], + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["pipeline_revision_id"], + ["dataflow_pipeline_revisions.id"], + ondelete="RESTRICT", + ), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "tenant_id", + "pipeline_id", + "environment", + name="uq_dataflow_pipeline_deployment_environment", + ), + ) + for column in ( + "tenant_id", + "pipeline_id", + "pipeline_revision_id", + "environment", + "status", + "promoted_by", + ): + op.create_index( + op.f(f"ix_dataflow_pipeline_deployments_{column}"), + "dataflow_pipeline_deployments", + [column], + unique=False, + ) + op.create_index( + "ix_dataflow_deployments_tenant_environment", + "dataflow_pipeline_deployments", + ["tenant_id", "environment"], + unique=False, + ) + + +def downgrade() -> None: + op.drop_index( + "ix_dataflow_deployments_tenant_environment", + table_name="dataflow_pipeline_deployments", + ) + op.drop_table("dataflow_pipeline_deployments") + op.drop_index( + "ix_dataflow_runs_worker_queue", + table_name="dataflow_runs", + ) + for column in ( + "retention_until", + "worker_id", + "lease_expires_at", + "available_at", + "environment", + "execution_backend", + ): + op.drop_index( + op.f(f"ix_dataflow_runs_{column}"), + table_name="dataflow_runs", + ) + with op.batch_alter_table("dataflow_runs") as batch_op: + batch_op.drop_column("resource_budget") + batch_op.drop_column("authorization") + batch_op.drop_column("purged_at") + batch_op.drop_column("retention_until") + batch_op.drop_column("progress_phase") + batch_op.drop_column("progress_percent") + batch_op.drop_column("cancellation_requested_at") + batch_op.drop_column("worker_id") + batch_op.drop_column("heartbeat_at") + batch_op.drop_column("lease_expires_at") + batch_op.drop_column("claimed_at") + batch_op.drop_column("available_at") + batch_op.drop_column("queued_at") + batch_op.drop_column("max_attempts") + batch_op.drop_column("attempts") + batch_op.drop_column("environment") + batch_op.drop_column("execution_backend") + batch_op.alter_column( + "started_at", + existing_type=sa.DateTime(timezone=True), + nullable=False, + ) diff --git a/src/govoplan_dataflow/backend/router.py b/src/govoplan_dataflow/backend/router.py index 0eeb6d2..2d61163 100644 --- a/src/govoplan_dataflow/backend/router.py +++ b/src/govoplan_dataflow/backend/router.py @@ -56,13 +56,17 @@ from govoplan_dataflow.backend.schemas import ( PipelineCreateRequest, PipelineDeriveRequest, PipelineDeleteResponse, + PipelineDeploymentListResponse, + PipelineDeploymentResponse, PipelineDraftRequest, PipelineListResponse, PipelinePreviewRequest, PipelinePreviewResponse, PipelineRunCreateRequest, PipelineRunListResponse, + PipelineRunMetricsResponse, PipelineRunResponse, + PipelinePromotionRequest, PipelineResponse, PipelineSqlResponse, PipelineUpdateRequest, @@ -92,16 +96,20 @@ from govoplan_dataflow.backend.service import ( delete_pipeline, get_pipeline, get_pipeline_run, + list_pipeline_deployments, list_pipeline_runs, list_pipelines, pipeline_response, + pipeline_deployment_response, pipeline_run_response, preview_pipeline, + promote_pipeline, render_graph_sql, start_pipeline_run, update_pipeline, validate_draft, ) +from govoplan_dataflow.backend.run_worker import run_metrics from govoplan_dataflow.backend.triggers import ( create_trigger, delete_trigger, @@ -986,7 +994,7 @@ def api_list_pipeline_runs( @router.post( "/pipelines/{pipeline_id}/runs", response_model=PipelineRunResponse, - status_code=status.HTTP_201_CREATED, + status_code=status.HTTP_202_ACCEPTED, ) def api_start_pipeline_run( pipeline_id: str, @@ -1028,6 +1036,10 @@ def api_start_pipeline_run( revision=payload.revision or pipeline.current_revision, idempotency_key=payload.idempotency_key, row_limit=payload.row_limit, + execution_backend=payload.execution_backend, + environment=payload.environment, + max_attempts=payload.max_attempts, + retention_days=payload.retention_days, publication=publication, invocation=AutomationInvocation( kind=( @@ -1040,6 +1052,7 @@ def api_start_pipeline_run( requested_by=_actor_id(principal), ), ), + defer_execution=True, ) except DataflowError as exc: raise _http_error(exc) from exc @@ -1069,6 +1082,17 @@ def api_start_pipeline_run( return response +@router.get("/runs/metrics", response_model=PipelineRunMetricsResponse) +def api_dataflow_run_metrics( + session: Session = Depends(get_session), + principal: ApiPrincipal = Depends(get_api_principal), +) -> PipelineRunMetricsResponse: + _require_any_scope(principal, READ_SCOPE, RUN_SCOPE, ADMIN_SCOPE) + return PipelineRunMetricsResponse.model_validate( + run_metrics(session, tenant_id=principal.tenant_id) + ) + + @router.get("/runs/{run_id}", response_model=PipelineRunResponse) def api_get_pipeline_run( run_id: str, @@ -1117,6 +1141,87 @@ def api_cancel_pipeline_run( return response +@router.get( + "/pipelines/{pipeline_id}/deployments", + response_model=PipelineDeploymentListResponse, +) +def api_list_pipeline_deployments( + pipeline_id: str, + session: Session = Depends(get_session), + principal: ApiPrincipal = Depends(get_api_principal), +) -> PipelineDeploymentListResponse: + _require_any_scope(principal, READ_SCOPE, RUN_SCOPE, ADMIN_SCOPE) + try: + pipeline = get_pipeline( + session, + tenant_id=principal.tenant_id, + pipeline_id=pipeline_id, + ) + require_definition_action( + pipeline, + principal=principal, + registry=get_registry(), + action="view", + ) + deployments = list_pipeline_deployments( + session, + tenant_id=principal.tenant_id, + pipeline_id=pipeline_id, + ) + return PipelineDeploymentListResponse( + deployments=[ + pipeline_deployment_response(session, item) + for item in deployments + ] + ) + except PermissionError as exc: + raise _governance_http_error(exc) from exc + except DataflowError as exc: + raise _http_error(exc) from exc + + +@router.post( + "/pipelines/{pipeline_id}/promotions", + response_model=PipelineDeploymentResponse, +) +def api_promote_pipeline( + pipeline_id: str, + payload: PipelinePromotionRequest, + session: Session = Depends(get_session), + principal: ApiPrincipal = Depends(get_api_principal), +) -> PipelineDeploymentResponse: + _require_any_scope(principal, WRITE_SCOPE, ADMIN_SCOPE) + try: + deployment = promote_pipeline( + session, + tenant_id=principal.tenant_id, + pipeline_id=pipeline_id, + actor_id=_actor_id(principal), + principal=principal, + registry=get_registry(), + payload=payload, + ) + except DataflowError as exc: + raise _http_error(exc) from exc + audit_event( + session, + tenant_id=principal.tenant_id, + user_id=getattr(principal.user, "id", None), + api_key_id=principal.api_key_id, + action="dataflow.pipeline.promoted", + object_type="dataflow_pipeline", + object_id=pipeline_id, + details={ + "revision_id": deployment.pipeline_revision_id, + "environment": deployment.environment, + "source_environment": deployment.source_environment, + }, + ) + response = pipeline_deployment_response(session, deployment) + session.commit() + return response + + @router.post("/validate", response_model=PipelineValidationResponse) def api_validate_pipeline( payload: PipelineDraftRequest, diff --git a/src/govoplan_dataflow/backend/run_worker.py b/src/govoplan_dataflow/backend/run_worker.py new file mode 100644 index 0000000..64ac86e --- /dev/null +++ b/src/govoplan_dataflow/backend/run_worker.py @@ -0,0 +1,631 @@ +from __future__ import annotations + +from collections.abc import Mapping +from datetime import datetime, timedelta, timezone +import logging +import socket + +from sqlalchemy import func, or_, select +from sqlalchemy.orm import Session + +from govoplan_core.auth import ApiPrincipal +from govoplan_core.core.automation import ( + AutomationPrincipalRequest, + automation_principal_provider, +) +from govoplan_core.core.notifications import ( + NotificationDispatchRequest, + notification_dispatch_provider, +) +from govoplan_core.db.base import utcnow +from govoplan_dataflow.backend.db.models import ( + DataflowPipeline, + DataflowPipelineRevision, + DataflowRun, + DataflowTrigger, +) +from govoplan_dataflow.backend.governance import require_definition_action +from govoplan_dataflow.backend.service import ( + _execute_pipeline_run, + _require_deployed_revision, + pipeline_run_request, +) + + +MAX_ACTIVE_RUNS_PER_TENANT = 4 +DEFAULT_LEASE_SECONDS = 120 +logger = logging.getLogger(__name__) + + +class DataflowWorkerError(RuntimeError): + pass + + +def dispatch_pending_runs( + session: Session, + *, + registry: object | None, + now: datetime | None = None, + limit: int = 10, + worker_id: str | None = None, +) -> dict[str, object]: + current = _as_utc(now or utcnow()) + resolved_worker_id = ( + str(worker_id).strip() + if worker_id and str(worker_id).strip() + else socket.gethostname() + ) + recovered = _recover_expired_leases(session, now=current) + session.commit() + summary: dict[str, object] = { + "claimed": 0, + "succeeded": 0, + "retrying": 0, + "failed": 0, + "cancelled": 0, + "recovered": recovered, + "runs": [], + } + for _ in range(max(1, min(int(limit), 100))): + run_id = _claim_next_run( + session, + now=current, + worker_id=resolved_worker_id, + ) + if run_id is None: + session.rollback() + break + session.commit() + summary["claimed"] = int(summary["claimed"]) + 1 + outcome = _execute_claimed_run( + session, + run_id=run_id, + registry=registry, + now=current, + ) + session.commit() + summary[outcome] = int(summary[outcome]) + 1 + runs = list(summary["runs"]) + runs.append({"ref": f"dataflow-run:{run_id}", "status": outcome}) + summary["runs"] = runs + return summary + + +def purge_expired_runs( + session: Session, + *, + now: datetime | None = None, + limit: int = 500, +) -> dict[str, object]: + current = _as_utc(now or utcnow()) + runs = list( + session.scalars( + select(DataflowRun) + .where( + DataflowRun.status.in_( + ("succeeded", "failed", "cancelled") + ), + DataflowRun.retention_until.is_not(None), + DataflowRun.retention_until <= current, + DataflowRun.purged_at.is_(None), + ) + .order_by(DataflowRun.retention_until, DataflowRun.id) + .limit(max(1, min(int(limit), 5_000))) + .with_for_update(skip_locked=True) + ) + ) + for run in runs: + run.request_ = {} + run.source_fingerprints = [] + run.diagnostics = [] + authorization = dict(run.authorization_) + authorization.pop("submitted_principal", None) + authorization["personal_data_purged"] = True + run.authorization_ = authorization + run.purged_at = current + session.flush() + return { + "purged": len(runs), + "run_refs": [f"dataflow-run:{run.id}" for run in runs], + } + + +def run_metrics( + session: Session, + *, + tenant_id: str, + now: datetime | None = None, +) -> dict[str, object]: + current = _as_utc(now or utcnow()) + rows = session.execute( + select(DataflowRun.status, func.count()) + .where(DataflowRun.tenant_id == tenant_id) + .group_by(DataflowRun.status) + ).all() + statuses = {str(status): int(count) for status, count in rows} + recent = list( + session.execute( + select( + DataflowRun.status, + DataflowRun.started_at, + DataflowRun.finished_at, + ).where( + DataflowRun.tenant_id == tenant_id, + DataflowRun.finished_at.is_not(None), + DataflowRun.finished_at >= current - timedelta(hours=24), + ) + ) + ) + durations = [ + ( + _as_utc(finished_at) - _as_utc(started_at) + ).total_seconds() + for _status, started_at, finished_at in recent + if finished_at is not None and started_at is not None + ] + oldest_queued_at = session.scalar( + select(func.min(DataflowRun.queued_at)).where( + DataflowRun.tenant_id == tenant_id, + DataflowRun.status.in_(("queued", "retrying")), + ) + ) + return { + "statuses": statuses, + "active": sum( + statuses.get(status, 0) + for status in ("queued", "retrying", "running") + ), + "completed_last_24_hours": len(recent), + "failed_last_24_hours": sum( + 1 for status, _started_at, _finished_at in recent + if status == "failed" + ), + "average_duration_seconds": ( + round(sum(durations) / len(durations), 3) + if durations + else None + ), + "oldest_queued_at": oldest_queued_at, + } + + +def _recover_expired_leases( + session: Session, + *, + now: datetime, +) -> int: + runs = list( + session.scalars( + select(DataflowRun) + .where( + DataflowRun.status == "running", + DataflowRun.lease_expires_at.is_not(None), + DataflowRun.lease_expires_at < now, + ) + .with_for_update(skip_locked=True) + ) + ) + for run in runs: + run.worker_id = None + run.claimed_at = None + run.lease_expires_at = None + run.heartbeat_at = None + if run.cancellation_requested_at is not None: + _cancel_run(run, now=now) + elif run.attempts < run.max_attempts: + run.status = "retrying" + run.available_at = now + run.progress_phase = "retrying_after_worker_loss" + run.error = "The previous worker lease expired." + else: + _fail_run( + run, + now=now, + message="The worker lease expired after the final attempt.", + ) + return len(runs) + + +def _claim_next_run( + session: Session, + *, + now: datetime, + worker_id: str, +) -> str | None: + candidates = list( + session.scalars( + select(DataflowRun) + .where( + DataflowRun.status.in_(("queued", "retrying")), + or_( + DataflowRun.available_at.is_(None), + DataflowRun.available_at <= now, + ), + ) + .order_by(DataflowRun.available_at, DataflowRun.created_at) + .limit(40) + .with_for_update(skip_locked=True) + ) + ) + for run in candidates: + active = int( + session.scalar( + select(func.count()) + .select_from(DataflowRun) + .where( + DataflowRun.tenant_id == run.tenant_id, + DataflowRun.status == "running", + ) + ) + or 0 + ) + if active >= MAX_ACTIVE_RUNS_PER_TENANT: + continue + run.status = "running" + run.attempts += 1 + run.worker_id = worker_id[:255] + run.claimed_at = now + run.heartbeat_at = now + run.lease_expires_at = now + timedelta( + seconds=_lease_seconds(run) + ) + run.started_at = run.started_at or now + run.finished_at = None + run.progress_percent = max(run.progress_percent, 5) + run.progress_phase = "authorizing" + session.flush() + return run.id + return None + + +def _execute_claimed_run( + session: Session, + *, + run_id: str, + registry: object | None, + now: datetime, +) -> str: + run = session.get(DataflowRun, run_id) + if run is None: + return "failed" + if run.cancellation_requested_at is not None: + _cancel_run(run, now=now) + return "cancelled" + try: + principal, provenance = _resolve_principal( + session, + run=run, + registry=registry, + ) + run.authorization_ = { + **dict(run.authorization_), + "last_resolution": provenance, + "resolved_at": now.isoformat(), + } + pipeline = session.get(DataflowPipeline, run.pipeline_id) + revision = session.get( + DataflowPipelineRevision, + run.pipeline_revision_id, + ) + if pipeline is None or revision is None: + raise DataflowWorkerError( + "The pipeline or pinned revision no longer exists." + ) + action = ( + "run" + if run.invocation_kind in {"manual", "api", "backfill"} + else "automate" + ) + require_definition_action( + pipeline, + principal=principal, + registry=registry, + action=action, + ) + _require_deployed_revision( + session, + tenant_id=run.tenant_id, + pipeline=pipeline, + revision=revision, + environment=run.environment, + ) + _notify_run( + session, + registry=registry, + run=run, + event_kind="dataflow.run.started", + subject=f"Dataflow run started: {pipeline.name}", + ) + retryable = _execute_pipeline_run( + session, + run=run, + pipeline=pipeline, + revision=revision, + request=pipeline_run_request(run), + principal=principal, + registry=registry, + ) + if ( + run.status == "failed" + and retryable + and run.attempts < run.max_attempts + ): + _retry_run(run, now=now) + _finish_claim(run) + _update_trigger_status(session, run) + _notify_terminal_run( + session, + registry=registry, + run=run, + pipeline_name=pipeline.name, + ) + return run.status + except (DataflowWorkerError, PermissionError, ValueError) as exc: + _fail_run(run, now=now, message=str(exc)) + _finish_claim(run) + _update_trigger_status(session, run) + _notify_terminal_run( + session, + registry=registry, + run=run, + pipeline_name="Dataflow", + ) + return "failed" + except Exception as exc: + logger.exception( + "Unexpected Dataflow worker failure for run %s", + run_id, + ) + session.rollback() + run = session.get(DataflowRun, run_id) + if run is None: + return "failed" + _fail_run( + run, + now=now, + message=( + "Unexpected worker failure " + f"({type(exc).__name__})." + ), + ) + _finish_claim(run) + return "failed" + + +def _resolve_principal( + session: Session, + *, + run: DataflowRun, + registry: object | None, +) -> tuple[ApiPrincipal, Mapping[str, object]]: + provider = automation_principal_provider(registry) + if provider is None: + raise DataflowWorkerError( + "Queued Dataflow execution requires the automation-principal " + "capability." + ) + authorization = dict(run.authorization_) + subject_kind = str( + authorization.get("subject_kind") or "delegated_user" + ) + common = { + "tenant_id": run.tenant_id, + "authorization_ref": str( + authorization.get("authorization_ref") + or f"dataflow-run:{run.id}" + ), + "grant_scopes": tuple( + str(scope) + for scope in authorization.get("grant_scopes") or () + ), + "context": { + "run_ref": f"dataflow-run:{run.id}", + "pipeline_ref": f"pipeline:{run.pipeline_id}", + "environment": run.environment, + "attempt": run.attempts, + }, + } + if subject_kind == "service_account": + service_account_id = str( + authorization.get("service_account_id") or "" + ).strip() + request = AutomationPrincipalRequest.service_account( + service_account_id=service_account_id, + **common, + ) + else: + request = AutomationPrincipalRequest.delegated_user( + account_id=str(authorization.get("account_id") or ""), + membership_id=str( + authorization.get("membership_id") or "" + ), + **common, + ) + resolution = provider.resolve_automation_principal( + session, + request=request, + ) + if not resolution.allowed or not isinstance( + resolution.principal, + ApiPrincipal, + ): + raise DataflowWorkerError( + resolution.reason or "Dataflow run authorization was denied." + ) + return resolution.principal, resolution.provenance + + +def _retry_run(run: DataflowRun, *, now: datetime) -> None: + delay_seconds = min(15 * (2 ** max(0, run.attempts - 1)), 900) + run.status = "retrying" + run.available_at = now + timedelta(seconds=delay_seconds) + run.finished_at = None + run.progress_percent = 0 + run.progress_phase = "retrying" + + +def _finish_claim(run: DataflowRun) -> None: + run.worker_id = None + run.claimed_at = None + run.lease_expires_at = None + run.heartbeat_at = None + + +def _cancel_run(run: DataflowRun, *, now: datetime) -> None: + run.status = "cancelled" + run.finished_at = now + run.error = "Cancelled by request." + run.progress_phase = "cancelled" + _finish_claim(run) + + +def _fail_run( + run: DataflowRun, + *, + now: datetime, + message: str, +) -> None: + run.status = "failed" + run.finished_at = now + run.error = message + run.progress_phase = "failed" + + +def _update_trigger_status(session: Session, run: DataflowRun) -> None: + if not run.trigger_id: + return + trigger = session.get(DataflowTrigger, run.trigger_id) + if trigger is None: + return + trigger.last_status = run.status + trigger.last_error = run.error + + +def _notify_terminal_run( + session: Session, + *, + registry: object | None, + run: DataflowRun, + pipeline_name: str, +) -> None: + if run.status == "retrying": + event_kind = "dataflow.run.retrying" + subject = f"Dataflow run will retry: {pipeline_name}" + elif run.status == "succeeded": + event_kind = "dataflow.run.completed" + subject = f"Dataflow run completed: {pipeline_name}" + elif run.status == "cancelled": + event_kind = "dataflow.run.cancelled" + subject = f"Dataflow run cancelled: {pipeline_name}" + else: + event_kind = "dataflow.run.failed" + subject = f"Dataflow run failed: {pipeline_name}" + _notify_run( + session, + registry=registry, + run=run, + event_kind=event_kind, + subject=subject, + ) + + +def _notify_run( + session: Session, + *, + registry: object | None, + run: DataflowRun, + event_kind: str, + subject: str, +) -> None: + provider = notification_dispatch_provider(registry) + authorization = dict(run.authorization_) + recipient_id = str(authorization.get("account_id") or "").strip() + if provider is None or not recipient_id: + return + try: + provider.enqueue_notification( + session, + NotificationDispatchRequest( + tenant_id=run.tenant_id, + source_module="dataflow", + source_resource_type="dataflow_run", + source_resource_id=run.id, + event_kind=event_kind, + recipient_type="account", + recipient_id=recipient_id, + subject=subject, + body_text=run.error, + action_url="/dataflow", + payload={ + "run_ref": f"dataflow-run:{run.id}", + "pipeline_ref": f"pipeline:{run.pipeline_id}", + "status": run.status, + "environment": run.environment, + "attempt": run.attempts, + }, + ), + ) + except Exception: + # Notification delivery is optional and cannot change run outcome. + logger.warning( + "Could not enqueue Dataflow notification for run %s", + run.id, + exc_info=True, + ) + return + + +def _lease_seconds(run: DataflowRun) -> int: + budget = dict(run.resource_budget) + wall_seconds = float(budget.get("max_wall_seconds") or 30.0) + return max(DEFAULT_LEASE_SECONDS, int(wall_seconds) + 60) + + +def _as_utc(value: datetime | None) -> datetime: + if value is None: + return datetime.now(timezone.utc) + if value.tzinfo is None: + return value.replace(tzinfo=timezone.utc) + return value.astimezone(timezone.utc) + + +class SqlDataflowRunWorker: + def __init__(self, *, registry: object | None = None) -> None: + self._registry = registry + + def dispatch_pending( + self, + session: object, + *, + now: datetime | None = None, + limit: int = 10, + worker_id: str | None = None, + ) -> Mapping[str, object]: + if not isinstance(session, Session): + raise TypeError("Dataflow run dispatch requires a Session.") + return dispatch_pending_runs( + session, + registry=self._registry, + now=now, + limit=limit, + worker_id=worker_id, + ) + + def purge_expired( + self, + session: object, + *, + now: datetime | None = None, + limit: int = 500, + ) -> Mapping[str, object]: + if not isinstance(session, Session): + raise TypeError("Dataflow retention requires a Session.") + return purge_expired_runs(session, now=now, limit=limit) + + +__all__ = [ + "SqlDataflowRunWorker", + "dispatch_pending_runs", + "purge_expired_runs", + "run_metrics", +] diff --git a/src/govoplan_dataflow/backend/schemas.py b/src/govoplan_dataflow/backend/schemas.py index b90075f..dccacd6 100644 --- a/src/govoplan_dataflow/backend/schemas.py +++ b/src/govoplan_dataflow/backend/schemas.py @@ -22,6 +22,16 @@ TriggerDeliveryStatus = Literal[ "blocked", "skipped", ] +DataflowEnvironment = Literal["development", "staging", "production"] +DataflowExecutionBackend = Literal["auto", "reference", "duckdb"] +DataflowRunStatus = Literal[ + "queued", + "retrying", + "running", + "succeeded", + "failed", + "cancelled", +] class GraphPosition(BaseModel): @@ -267,7 +277,11 @@ class PipelineRunPublicationRequest(BaseModel): class PipelineRunCreateRequest(BaseModel): revision: int | None = Field(default=None, ge=1) idempotency_key: str = Field(min_length=1, max_length=255) - row_limit: int = Field(default=500, ge=1, le=500) + row_limit: int = Field(default=500, ge=1, le=10_000) + execution_backend: DataflowExecutionBackend = "auto" + environment: DataflowEnvironment = "development" + max_attempts: int = Field(default=3, ge=1, le=5) + retention_days: int = Field(default=30, ge=1, le=365) publication: PipelineRunPublicationRequest | None = None invocation_kind: Literal["manual", "api", "backfill"] = "manual" @@ -277,8 +291,10 @@ class PipelineRunResponse(BaseModel): pipeline_id: str revision: int run_type: str - status: Literal["running", "succeeded", "failed", "cancelled"] + status: DataflowRunStatus idempotency_key: str | None + execution_backend: str + environment: DataflowEnvironment definition_hash: str executor_version: str source_fingerprints: list[dict[str, Any]] @@ -294,8 +310,18 @@ class PipelineRunResponse(BaseModel): delivery_ref: str | None correlation_id: str | None causation_id: str | None + attempts: int + max_attempts: int + available_at: datetime | None + claimed_at: datetime | None + lease_expires_at: datetime | None + cancellation_requested_at: datetime | None + progress_percent: int + progress_phase: str + retention_until: datetime | None + purged_at: datetime | None error: str | None - started_at: datetime + started_at: datetime | None finished_at: datetime | None created_by: str | None created_at: datetime @@ -306,6 +332,51 @@ class PipelineRunListResponse(BaseModel): runs: list[PipelineRunResponse] +class PipelineRunMetricsResponse(BaseModel): + statuses: dict[str, int] + active: int + completed_last_24_hours: int + failed_last_24_hours: int + average_duration_seconds: float | None + oldest_queued_at: datetime | None + + +class PipelinePromotionRequest(BaseModel): + revision: int | None = Field(default=None, ge=1) + source_environment: DataflowEnvironment = "development" + target_environment: Literal["staging", "production"] + + @model_validator(mode="after") + def validate_transition(self) -> "PipelinePromotionRequest": + allowed = { + ("development", "staging"), + ("staging", "production"), + } + if (self.source_environment, self.target_environment) not in allowed: + raise ValueError( + "Promotions must move from development to staging or " + "from staging to production." + ) + return self + + +class PipelineDeploymentResponse(BaseModel): + id: str + pipeline_id: str + revision: int + environment: DataflowEnvironment + source_environment: DataflowEnvironment + status: str + provenance: dict[str, Any] + promoted_by: str | None + created_at: datetime + updated_at: datetime + + +class PipelineDeploymentListResponse(BaseModel): + deployments: list[PipelineDeploymentResponse] + + JsonScalar = str | int | float | bool | None @@ -353,7 +424,11 @@ class DataflowTriggerCreateRequest(BaseModel): schedule: DataflowTriggerSchedule | None = None event: DataflowTriggerEvent | None = None publication: PipelineRunPublicationRequest | None = None - row_limit: int = Field(default=500, ge=1, le=500) + row_limit: int = Field(default=500, ge=1, le=10_000) + environment: DataflowEnvironment = "development" + execution_backend: DataflowExecutionBackend = "auto" + max_attempts: int = Field(default=3, ge=1, le=5) + retention_days: int = Field(default=30, ge=1, le=365) catch_up_policy: Literal["skip", "coalesce"] = "coalesce" max_concurrent_runs: int = Field(default=1, ge=1, le=10) @@ -392,6 +467,10 @@ class DataflowTriggerResponse(BaseModel): event: DataflowTriggerEvent | None publication: PipelineRunPublicationRequest | None row_limit: int + environment: DataflowEnvironment + execution_backend: DataflowExecutionBackend + max_attempts: int + retention_days: int catch_up_policy: Literal["skip", "coalesce"] max_concurrent_runs: int next_fire_at: datetime | None diff --git a/src/govoplan_dataflow/backend/service.py b/src/govoplan_dataflow/backend/service.py index d1d68a0..9b82d4b 100644 --- a/src/govoplan_dataflow/backend/service.py +++ b/src/govoplan_dataflow/backend/service.py @@ -4,13 +4,15 @@ from collections.abc import Mapping import hashlib import json from dataclasses import dataclass +from datetime import datetime, timedelta -from sqlalchemy import or_, select +from sqlalchemy import func, or_, select from sqlalchemy.orm import Session from govoplan_core.auth import ApiPrincipal from govoplan_core.core.automation import AutomationInvocation from govoplan_core.core.dataflows import ( + DataflowPublicationTarget, DataflowRunConflictError, DataflowRunDescriptor, DataflowRunError, @@ -21,6 +23,7 @@ from govoplan_core.core.datasources import ( DatasourceError, DatasourcePublicationRequest, DatasourceReadRequest, + DatasourceUnavailableError, datasource_catalogue, datasource_publication, ) @@ -34,6 +37,7 @@ from govoplan_dataflow.backend.backends import ( from govoplan_dataflow.backend.batches import TypedBatch from govoplan_dataflow.backend.db.models import ( DataflowPipeline, + DataflowPipelineDeployment, DataflowPipelineRevision, DataflowRun, ) @@ -65,6 +69,8 @@ from govoplan_dataflow.backend.schemas import ( PipelineGraph, PipelinePreviewRequest, PipelinePreviewResponse, + PipelineDeploymentResponse, + PipelinePromotionRequest, PipelineResponse, PipelineRevisionResponse, PipelineRunResponse, @@ -98,6 +104,13 @@ class DataflowValidationError(DataflowError): self.diagnostics = diagnostics +MAX_PENDING_RUNS_PER_TENANT = 100 +MAX_PRODUCTION_ROWS = 10_000 +RUN_SCOPE = "dataflow:pipeline:run" +DATASOURCE_READ_SCOPE = "datasources:catalogue:read" +DATASOURCE_WRITE_SCOPE = "datasources:source:write" + + @dataclass(frozen=True) class NormalizedDefinition: graph: PipelineGraph @@ -694,6 +707,7 @@ def _execute_pipeline_preview( backend: str, row_limit: int, preview_node_id: str | None, + budget: ExecutionBudget | None = None, ) -> tuple[PipelineExecutionResult, str]: source_resolver = _preview_source_resolver( session=session, @@ -713,13 +727,14 @@ def _execute_pipeline_preview( sources = _typed_backend_sources( graph, source_resolver=source_resolver, + source_limit=max(MAX_SOURCE_ROWS, row_limit), ) try: result = execute_typed_graph( graph, backend=backend, sources=sources, - budget=ExecutionBudget(max_output_rows=row_limit), + budget=budget or ExecutionBudget(max_output_rows=row_limit), preview_node_id=preview_node_id, ) except BackendExecutionError as exc: @@ -727,6 +742,7 @@ def _execute_pipeline_preview( str(exc), node_id=exc.node_id, diagnostics=tuple(exc.diagnostics), + retryable=exc.code == "backend.capacity", ) from exc columns = [ PreviewColumn( @@ -802,12 +818,13 @@ def _typed_backend_sources( graph: PipelineGraph, *, source_resolver, + source_limit: int = MAX_SOURCE_ROWS, ) -> dict[str, BackendSource]: sources: dict[str, BackendSource] = {} for node in graph.nodes: if node.type != "source.reference": continue - resolved = source_resolver(node, MAX_SOURCE_ROWS) + resolved = source_resolver(node, source_limit) sources[node.id] = BackendSource( node_id=node.id, batch=TypedBatch.from_rows(resolved.rows), @@ -863,6 +880,137 @@ def list_pipeline_runs( return list(session.scalars(statement)) +def list_pipeline_deployments( + session: Session, + *, + tenant_id: str, + pipeline_id: str, +) -> list[DataflowPipelineDeployment]: + get_pipeline( + session, + tenant_id=tenant_id, + pipeline_id=pipeline_id, + ) + return list( + session.scalars( + select(DataflowPipelineDeployment) + .where( + DataflowPipelineDeployment.tenant_id == tenant_id, + DataflowPipelineDeployment.pipeline_id == pipeline_id, + DataflowPipelineDeployment.status == "active", + ) + .order_by(DataflowPipelineDeployment.environment) + ) + ) + + +def promote_pipeline( + session: Session, + *, + tenant_id: str, + pipeline_id: str, + actor_id: str | None, + principal: ApiPrincipal, + registry: object | None, + payload: PipelinePromotionRequest, +) -> DataflowPipelineDeployment: + pipeline = get_pipeline( + session, + tenant_id=tenant_id, + pipeline_id=pipeline_id, + ) + try: + require_definition_action( + pipeline, + principal=principal, + registry=registry, + action="edit", + ) + except PermissionError as exc: + raise DataflowConflictError(str(exc)) from exc + revision = get_pipeline_revision( + session, + pipeline=pipeline, + revision=payload.revision, + ) + if payload.source_environment == "staging": + source = session.scalar( + select(DataflowPipelineDeployment).where( + DataflowPipelineDeployment.tenant_id == tenant_id, + DataflowPipelineDeployment.pipeline_id == pipeline.id, + DataflowPipelineDeployment.environment == "staging", + DataflowPipelineDeployment.status == "active", + ) + ) + if source is None or source.pipeline_revision_id != revision.id: + raise DataflowConflictError( + "Only the revision currently promoted to staging can be " + "promoted to production." + ) + deployment = session.scalar( + select(DataflowPipelineDeployment).where( + DataflowPipelineDeployment.tenant_id == tenant_id, + DataflowPipelineDeployment.pipeline_id == pipeline.id, + DataflowPipelineDeployment.environment + == payload.target_environment, + ) + ) + previous_revision_id = ( + deployment.pipeline_revision_id if deployment is not None else None + ) + if deployment is None: + deployment = DataflowPipelineDeployment( + tenant_id=tenant_id, + pipeline_id=pipeline.id, + pipeline_revision_id=revision.id, + environment=payload.target_environment, + source_environment=payload.source_environment, + status="active", + promoted_by=actor_id, + ) + session.add(deployment) + deployment.pipeline_revision_id = revision.id + deployment.source_environment = payload.source_environment + deployment.status = "active" + deployment.promoted_by = actor_id + deployment.provenance = { + "pipeline_ref": f"pipeline:{pipeline.id}", + "revision": revision.revision, + "definition_hash": revision.content_hash, + "source_environment": payload.source_environment, + "target_environment": payload.target_environment, + "previous_revision_id": previous_revision_id, + "promoted_by": actor_id, + "promoted_at": utcnow().isoformat(), + } + session.flush() + return deployment + + +def pipeline_deployment_response( + session: Session, + deployment: DataflowPipelineDeployment, +) -> PipelineDeploymentResponse: + revision = session.get( + DataflowPipelineRevision, + deployment.pipeline_revision_id, + ) + if revision is None: + raise DataflowNotFoundError("Dataflow pipeline revision not found") + return PipelineDeploymentResponse( + id=deployment.id, + pipeline_id=deployment.pipeline_id, + revision=revision.revision, + environment=deployment.environment, # type: ignore[arg-type] + source_environment=deployment.source_environment, # type: ignore[arg-type] + status=deployment.status, + provenance=dict(deployment.provenance), + promoted_by=deployment.promoted_by, + created_at=deployment.created_at, + updated_at=deployment.updated_at, + ) + + def get_pipeline_run( session: Session, *, @@ -891,6 +1039,7 @@ def start_pipeline_run( principal: ApiPrincipal, registry: object | None, request: DataflowRunRequest, + defer_execution: bool = False, ) -> tuple[DataflowRun, bool]: pipeline, revision = _run_definition( session, @@ -909,6 +1058,8 @@ def start_pipeline_run( ) if existing is not None: return existing, True + if defer_execution: + _require_run_queue_capacity(session, tenant_id=tenant_id) run = _new_pipeline_run( tenant_id=tenant_id, actor_id=actor_id, @@ -917,9 +1068,20 @@ def start_pipeline_run( request=request, idempotency_key=idempotency_key, request_hash=request_hash, + principal=principal, + defer_execution=defer_execution, ) session.add(run) session.flush() + run.authorization_ = { + **dict(run.authorization_), + "authorization_ref": ( + request.invocation.trigger_ref or f"dataflow-run:{run.id}" + ), + } + if defer_execution: + session.flush() + return run, False _execute_pipeline_run( session, run=run, @@ -968,6 +1130,13 @@ def _run_definition( pipeline=pipeline, revision=request.revision, ) + _require_deployed_revision( + session, + tenant_id=tenant_id, + pipeline=pipeline, + revision=revision, + environment=request.environment, + ) return pipeline, revision @@ -977,13 +1146,77 @@ def _validated_run_identity(request: DataflowRunRequest) -> tuple[str, str]: raise DataflowConflictError( "A Dataflow run idempotency key of at most 255 characters is required." ) - if request.row_limit < 1 or request.row_limit > 500: + if request.row_limit < 1 or request.row_limit > MAX_PRODUCTION_ROWS: raise DataflowConflictError( - "The bounded Dataflow runner supports between 1 and 500 output rows." + "The bounded Dataflow runner supports between 1 and 10,000 " + "output rows." + ) + if request.execution_backend not in {"auto", "reference", "duckdb"}: + raise DataflowConflictError("Unknown Dataflow execution backend.") + if request.environment not in {"development", "staging", "production"}: + raise DataflowConflictError("Unknown Dataflow environment.") + if request.max_attempts < 1 or request.max_attempts > 5: + raise DataflowConflictError( + "Dataflow runs support between one and five attempts." + ) + if request.retention_days < 1 or request.retention_days > 365: + raise DataflowConflictError( + "Dataflow run retention must be between one and 365 days." ) return idempotency_key, _run_request_hash(request) +def _require_run_queue_capacity( + session: Session, + *, + tenant_id: str, +) -> None: + pending = int( + session.scalar( + select(func.count()) + .select_from(DataflowRun) + .where( + DataflowRun.tenant_id == tenant_id, + DataflowRun.status.in_(("queued", "retrying", "running")), + ) + ) + or 0 + ) + if pending >= MAX_PENDING_RUNS_PER_TENANT: + raise DataflowConflictError( + "The tenant Dataflow queue is full; wait for an active run to " + "finish before submitting more work." + ) + + +def _require_deployed_revision( + session: Session, + *, + tenant_id: str, + pipeline: DataflowPipeline, + revision: DataflowPipelineRevision, + environment: str, +) -> None: + if environment == "development": + return + deployment = session.scalar( + select(DataflowPipelineDeployment).where( + DataflowPipelineDeployment.tenant_id == tenant_id, + DataflowPipelineDeployment.pipeline_id == pipeline.id, + DataflowPipelineDeployment.environment == environment, + DataflowPipelineDeployment.status == "active", + ) + ) + if ( + deployment is None + or deployment.pipeline_revision_id != revision.id + ): + raise DataflowConflictError( + f"Pipeline revision {revision.revision} is not promoted to " + f"{environment}." + ) + + def _existing_pipeline_run( session: Session, *, @@ -1017,13 +1250,19 @@ def _new_pipeline_run( request: DataflowRunRequest, idempotency_key: str, request_hash: str, + principal: ApiPrincipal, + defer_execution: bool, ) -> DataflowRun: + now = utcnow() + budget = _run_resource_budget(request) return DataflowRun( tenant_id=tenant_id, pipeline_id=pipeline.id, pipeline_revision_id=revision.id, run_type="published" if request.publication else "run", - status="running", + status="queued" if defer_execution else "running", + execution_backend=request.execution_backend, + environment=request.environment, executor_version=EXECUTOR_VERSION, definition_hash=revision.content_hash, idempotency_key=idempotency_key, @@ -1045,11 +1284,73 @@ def _new_pipeline_run( diagnostics=[], input_row_count=0, output_row_count=0, - started_at=utcnow(), + attempts=0 if defer_execution else 1, + max_attempts=request.max_attempts, + queued_at=now if defer_execution else None, + available_at=now if defer_execution else None, + progress_percent=0 if defer_execution else 10, + progress_phase="queued" if defer_execution else "executing", + retention_until=now + timedelta(days=request.retention_days), + authorization_=_run_authorization_payload( + principal, + request=request, + revision=revision, + ), + resource_budget=budget, + started_at=None if defer_execution else now, created_by=actor_id, ) +def _run_authorization_payload( + principal: ApiPrincipal, + *, + request: DataflowRunRequest, + revision: DataflowPipelineRevision, +) -> dict[str, object]: + principal_ref = principal.to_platform_principal() + graph = PipelineGraph.model_validate(revision.graph) + scopes = {RUN_SCOPE} + if any( + node.type == "source.reference" + or bool(node.config.get("source_ref")) + for node in graph.nodes + ): + scopes.add(DATASOURCE_READ_SCOPE) + if request.publication is not None: + scopes.add(DATASOURCE_WRITE_SCOPE) + return { + "contract_version": "1", + "subject_kind": ( + "service_account" + if principal_ref.service_account_id + else "delegated_user" + ), + "account_id": principal_ref.account_id, + "membership_id": principal_ref.membership_id, + "service_account_id": principal_ref.service_account_id, + "grant_scopes": sorted(scopes), + "submitted_principal": principal_ref.to_dict(), + "authorization_ref": request.invocation.trigger_ref, + "resolution": "rechecked_by_worker", + } + + +def _run_resource_budget( + request: DataflowRunRequest, +) -> dict[str, object]: + production = request.environment in {"staging", "production"} + return { + "max_output_rows": request.row_limit, + "max_batch_bytes": 64 * 1024 * 1024 if production else 8 * 1024 * 1024, + "max_wall_seconds": 30.0 if production else 10.0, + "max_memory_bytes": ( + 512 * 1024 * 1024 if production else 256 * 1024 * 1024 + ), + "max_concurrency": 1, + } + + def _execute_pipeline_run( session: Session, *, @@ -1059,18 +1360,40 @@ def _execute_pipeline_run( request: DataflowRunRequest, principal: ApiPrincipal, registry: object | None, -) -> None: +) -> bool: try: - result = execute_preview( + if run.cancellation_requested_at is not None: + _mark_pipeline_run_cancelled(run) + return False + execution_backend = run.execution_backend + if run.environment in {"staging", "production"}: + if execution_backend == "reference": + raise PipelineExecutionError( + "Staging and production runs require the isolated " + "DuckDB execution backend." + ) + execution_backend = "duckdb" + run.progress_percent = 20 + run.progress_phase = "reading_sources" + result, executor_version = _execute_pipeline_preview( PipelineGraph.model_validate(revision.graph), + session=session, + principal=principal, + registry=registry, + backend=execution_backend, row_limit=request.row_limit, - source_resolver=_datasource_source_resolver( - session=session, - principal=principal, - registry=registry, - ), + preview_node_id=None, + budget=_execution_budget(run), ) + run.executor_version = executor_version _apply_pipeline_result(run, result) + run.progress_percent = 80 + run.progress_phase = "publishing" if request.publication else "finalizing" + session.flush() + session.expire(run, ["cancellation_requested_at"]) + if run.cancellation_requested_at is not None: + _mark_pipeline_run_cancelled(run) + return False if request.publication: _ensure_publishable(result) _publish_pipeline_result( @@ -1086,8 +1409,27 @@ def _execute_pipeline_run( run.status = "succeeded" run.finished_at = utcnow() run.error = None + run.progress_percent = 100 + run.progress_phase = "completed" + return False except (DatasourceError, PipelineExecutionError) as exc: _mark_pipeline_run_failed(run, exc) + return isinstance(exc, DatasourceUnavailableError) or bool( + getattr(exc, "retryable", False) + ) + + +def _execution_budget(run: DataflowRun) -> ExecutionBudget: + value = dict(run.resource_budget) + return ExecutionBudget( + max_output_rows=int(value.get("max_output_rows") or 500), + max_batch_bytes=int(value.get("max_batch_bytes") or 8 * 1024 * 1024), + max_wall_seconds=float(value.get("max_wall_seconds") or 10.0), + max_memory_bytes=int( + value.get("max_memory_bytes") or 256 * 1024 * 1024 + ), + max_concurrency=int(value.get("max_concurrency") or 1), + ) def _apply_pipeline_result( @@ -1176,6 +1518,7 @@ def _mark_pipeline_run_failed( run.status = "failed" run.finished_at = utcnow() run.error = str(exc) + run.progress_phase = "failed" diagnostics = list(getattr(exc, "diagnostics", ())) diagnostics.append( DataflowDiagnostic( @@ -1194,6 +1537,13 @@ def _mark_pipeline_run_failed( run.input_row_count = int(getattr(exc, "input_row_count", 0)) +def _mark_pipeline_run_cancelled(run: DataflowRun) -> None: + run.status = "cancelled" + run.finished_at = utcnow() + run.error = "Cancelled by request." + run.progress_phase = "cancelled" + + def cancel_pipeline_run( session: Session, *, @@ -1205,13 +1555,15 @@ def cancel_pipeline_run( tenant_id=tenant_id, run_ref=run_ref, ) - if run.status not in {"queued", "running"}: + if run.status not in {"queued", "retrying", "running"}: raise DataflowConflictError( f"Dataflow run is already {run.status} and cannot be cancelled." ) - run.status = "cancelled" - run.finished_at = utcnow() - run.error = "Cancelled by request." + run.cancellation_requested_at = utcnow() + if run.status in {"queued", "retrying"}: + _mark_pipeline_run_cancelled(run) + else: + run.progress_phase = "cancellation_requested" session.flush() return run @@ -1232,6 +1584,8 @@ def pipeline_run_response( run_type=run.run_type, status=run.status, # type: ignore[arg-type] idempotency_key=run.idempotency_key, + execution_backend=run.execution_backend, + environment=run.environment, # type: ignore[arg-type] definition_hash=run.definition_hash, executor_version=run.executor_version, source_fingerprints=list(run.source_fingerprints), @@ -1253,6 +1607,16 @@ def pipeline_run_response( ), correlation_id=run.correlation_id, causation_id=run.causation_id, + attempts=run.attempts, + max_attempts=run.max_attempts, + available_at=run.available_at, + claimed_at=run.claimed_at, + lease_expires_at=run.lease_expires_at, + cancellation_requested_at=run.cancellation_requested_at, + progress_percent=run.progress_percent, + progress_phase=run.progress_phase, + retention_until=run.retention_until, + purged_at=run.purged_at, error=run.error, started_at=run.started_at, finished_at=run.finished_at, @@ -1298,6 +1662,12 @@ def pipeline_run_descriptor( replayed=replayed, metadata={ "run_type": run.run_type, + "execution_backend": run.execution_backend, + "environment": run.environment, + "attempts": run.attempts, + "max_attempts": run.max_attempts, + "progress_percent": run.progress_percent, + "progress_phase": run.progress_phase, "source_fingerprints": list(run.source_fingerprints), "diagnostics": list(run.diagnostics), }, @@ -1323,6 +1693,7 @@ class SqlDataflowRunLifecycleProvider: principal=api_principal, registry=self._registry, request=request, + defer_execution=True, ) return pipeline_run_descriptor(db, run, replayed=replayed) @@ -1421,30 +1792,66 @@ def _datasource_source_resolver( "catalogue capability.", node_id=node.id, ) + rows: list[dict[str, object]] = [] + resolved = None + offset = 0 + expected_fingerprint = _clean_optional( + node.config.get("expected_fingerprint") + ) try: - resolved = provider.read_datasource( - session, - principal, - request=DatasourceReadRequest( - datasource_ref=str(node.config["source_ref"]), - consistency=str( - node.config.get("consistency") or "current" - ), # type: ignore[arg-type] - limit=limit, - expected_fingerprint=_clean_optional( - node.config.get("expected_fingerprint") + while offset < limit: + page = provider.read_datasource( + session, + principal, + request=DatasourceReadRequest( + datasource_ref=str(node.config["source_ref"]), + consistency=str( + node.config.get("consistency") or "current" + ), # type: ignore[arg-type] + limit=min(500, limit - offset), + offset=offset, + expected_fingerprint=expected_fingerprint, ), - ), - ) + ) + if resolved is None: + resolved = page + expected_fingerprint = page.datasource.fingerprint + elif ( + page.datasource.fingerprint + != resolved.datasource.fingerprint + ): + raise PipelineExecutionError( + "Datasource changed while the run was reading it.", + node_id=node.id, + retryable=True, + ) + rows.extend(dict(row) for row in page.rows) + offset += len(page.rows) + if not page.truncated or not page.rows: + resolved = page + break + resolved = page + except DatasourceUnavailableError as exc: + raise PipelineExecutionError( + str(exc), + node_id=node.id, + retryable=True, + ) from exc except DatasourceError as exc: raise PipelineExecutionError(str(exc), node_id=node.id) from exc + if resolved is None: + raise PipelineExecutionError( + "Datasource returned no result.", + node_id=node.id, + retryable=True, + ) return ResolvedSource( - rows=tuple(dict(row) for row in resolved.rows), + rows=tuple(rows), source_ref=resolved.datasource.ref, provider=resolved.datasource.provider or "datasources", fingerprint=resolved.datasource.fingerprint, total_rows=resolved.total_rows, - truncated=resolved.truncated, + truncated=len(rows) < resolved.total_rows, ) return resolve_source @@ -1482,6 +1889,10 @@ def _run_request_payload(request: DataflowRunRequest) -> dict[str, object]: "pipeline_ref": request.pipeline_ref, "revision": request.revision, "row_limit": request.row_limit, + "execution_backend": request.execution_backend, + "environment": request.environment, + "max_attempts": request.max_attempts, + "retention_days": request.retention_days, "publication": ( { "target_datasource_ref": target.target_datasource_ref, @@ -1521,6 +1932,113 @@ def _invocation_payload( } +def pipeline_run_request(run: DataflowRun) -> DataflowRunRequest: + value = dict(run.request_) + publication_value = value.get("publication") + publication = ( + DataflowPublicationTarget( + target_datasource_ref=_mapping_optional_text( + publication_value, + "target_datasource_ref", + ), + name=_mapping_optional_text(publication_value, "name"), + source_name=_mapping_optional_text( + publication_value, + "source_name", + ), + description=_mapping_optional_text( + publication_value, + "description", + ), + freeze=bool(publication_value.get("freeze", False)), + frozen_label=_mapping_optional_text( + publication_value, + "frozen_label", + ), + set_current=bool(publication_value.get("set_current", True)), + metadata=( + dict(publication_value.get("metadata") or {}) + if isinstance(publication_value.get("metadata"), Mapping) + else {} + ), + ) + if isinstance(publication_value, Mapping) + else None + ) + invocation_value = value.get("invocation") + invocation_mapping = ( + invocation_value if isinstance(invocation_value, Mapping) else {} + ) + metadata = invocation_mapping.get("metadata") + return DataflowRunRequest( + pipeline_ref=str(value.get("pipeline_ref") or f"pipeline:{run.pipeline_id}"), + revision=int(value.get("revision") or 1), + idempotency_key=str(value.get("idempotency_key") or run.idempotency_key or run.id), + row_limit=int(value.get("row_limit") or 500), + execution_backend=str( + value.get("execution_backend") or run.execution_backend + ), + environment=str(value.get("environment") or run.environment), + max_attempts=int(value.get("max_attempts") or run.max_attempts), + retention_days=int(value.get("retention_days") or 30), + publication=publication, + invocation=AutomationInvocation( + kind=str( + invocation_mapping.get("kind") or run.invocation_kind + ), # type: ignore[arg-type] + trigger_ref=_mapping_optional_text( + invocation_mapping, + "trigger_ref", + ), + delivery_ref=_mapping_optional_text( + invocation_mapping, + "delivery_ref", + ), + event_id=_mapping_optional_text(invocation_mapping, "event_id"), + event_type=_mapping_optional_text( + invocation_mapping, + "event_type", + ), + correlation_id=_mapping_optional_text( + invocation_mapping, + "correlation_id", + ), + causation_id=_mapping_optional_text( + invocation_mapping, + "causation_id", + ), + scheduled_for=_optional_datetime( + invocation_mapping.get("scheduled_for") + ), + requested_by=_mapping_optional_text( + invocation_mapping, + "requested_by", + ), + metadata=dict(metadata) if isinstance(metadata, Mapping) else {}, + ), + ) + + +def _mapping_optional_text( + value: object, + key: str, +) -> str | None: + if not isinstance(value, Mapping): + return None + return _clean_optional(value.get(key)) + + +def _optional_datetime(value: object) -> datetime | None: + if isinstance(value, datetime): + return value + if isinstance(value, str) and value.strip(): + try: + return datetime.fromisoformat(value) + except ValueError: + return None + return None + + def _run_request_hash(request: DataflowRunRequest) -> str: encoded = json.dumps( _run_request_payload(request), @@ -1531,10 +2049,10 @@ def _run_request_hash(request: DataflowRunRequest) -> str: return hashlib.sha256(encoded.encode("utf-8")).hexdigest() -def _clean_optional(value: str | None) -> str | None: +def _clean_optional(value: object | None) -> str | None: if value is None: return None - cleaned = value.strip() + cleaned = str(value).strip() return cleaned or None @@ -1600,12 +2118,16 @@ __all__ = [ "get_pipeline", "get_pipeline_revision", "get_pipeline_run", + "list_pipeline_deployments", "list_pipeline_runs", "list_pipelines", "normalize_definition", "pipeline_response", + "pipeline_deployment_response", "pipeline_run_descriptor", + "pipeline_run_request", "pipeline_run_response", + "promote_pipeline", "preview_pipeline", "render_graph_sql", "start_pipeline_run", diff --git a/src/govoplan_dataflow/backend/triggers.py b/src/govoplan_dataflow/backend/triggers.py index 112421b..e935f11 100644 --- a/src/govoplan_dataflow/backend/triggers.py +++ b/src/govoplan_dataflow/backend/triggers.py @@ -429,6 +429,14 @@ def trigger_response( else None ), row_limit=trigger.row_limit, + environment=str( + config.get("environment") or "development" + ), # type: ignore[arg-type] + execution_backend=str( + config.get("execution_backend") or "auto" + ), # type: ignore[arg-type] + max_attempts=int(config.get("max_attempts") or 3), + retention_days=int(config.get("retention_days") or 30), catch_up_policy=trigger.catch_up_policy, # type: ignore[arg-type] max_concurrent_runs=trigger.max_concurrent_runs, next_fire_at=trigger.next_fire_at, @@ -528,6 +536,10 @@ def _trigger_config( "event": ( payload.event.model_dump(mode="json") if payload.event else None ), + "environment": payload.environment, + "execution_backend": payload.execution_backend, + "max_attempts": payload.max_attempts, + "retention_days": payload.retention_days, } @@ -800,6 +812,18 @@ def _dispatch_delivery( ), idempotency_key=f"trigger-delivery:{delivery.id}", row_limit=trigger.row_limit, + execution_backend=str( + trigger.config_.get("execution_backend") or "auto" + ), + environment=str( + trigger.config_.get("environment") or "development" + ), + max_attempts=int( + trigger.config_.get("max_attempts") or 3 + ), + retention_days=int( + trigger.config_.get("retention_days") or 30 + ), publication=publication, invocation=AutomationInvocation( kind=delivery.invocation_kind, # type: ignore[arg-type] @@ -822,6 +846,7 @@ def _dispatch_delivery( }, ), ), + defer_execution=True, ) except (DataflowConflictError, PermissionError, ValueError) as exc: trigger.last_status = "blocked" @@ -836,12 +861,11 @@ def _dispatch_delivery( delivery.run_id = run.id trigger.last_status = run.status trigger.last_error = run.error - status = "succeeded" if run.status == "succeeded" else "failed" return _finish_delivery( delivery, - status=status, + status="succeeded", now=now, - error=run.error, + error=None, ) @@ -865,7 +889,7 @@ def _running_count(session: Session, trigger_id: str) -> int: .select_from(DataflowRun) .where( DataflowRun.trigger_id == trigger_id, - DataflowRun.status.in_(("queued", "running")), + DataflowRun.status.in_(("queued", "retrying", "running")), ) ) or 0 diff --git a/tests/test_manifest.py b/tests/test_manifest.py index 162f4d7..22b5604 100644 --- a/tests/test_manifest.py +++ b/tests/test_manifest.py @@ -36,6 +36,7 @@ class DataflowManifestTests(unittest.TestCase): "dataflow.pipeline_catalog", "dataflow.pipeline_preview", "dataflow.run_lifecycle", + "dataflow.run_worker", "dataflow.dataset_output", "dataflow.trigger_dispatcher", }, diff --git a/tests/test_migrations.py b/tests/test_migrations.py index 01d21f8..f4508ec 100644 --- a/tests/test_migrations.py +++ b/tests/test_migrations.py @@ -24,13 +24,14 @@ class DataflowMigrationTests(unittest.TestCase): try: with engine.connect() as connection: self.assertIn( - "e2f4a8c9d1b7", + "f6c2a9d4e7b1", set(MigrationContext.configure(connection).get_current_heads()), ) self.assertEqual( { "dataflow_pipelines", "dataflow_pipeline_revisions", + "dataflow_pipeline_deployments", "dataflow_runs", "dataflow_triggers", "dataflow_trigger_deliveries", diff --git a/tests/test_run_worker.py b/tests/test_run_worker.py new file mode 100644 index 0000000..cc7327a --- /dev/null +++ b/tests/test_run_worker.py @@ -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() diff --git a/webui/src/api/dataflow.ts b/webui/src/api/dataflow.ts index b9ff9c2..2367b60 100644 --- a/webui/src/api/dataflow.ts +++ b/webui/src/api/dataflow.ts @@ -214,8 +214,10 @@ export type PipelineRun = { pipeline_id: string; revision: number; run_type: string; - status: "running" | "succeeded" | "failed" | "cancelled"; + status: "queued" | "retrying" | "running" | "succeeded" | "failed" | "cancelled"; idempotency_key?: string | null; + execution_backend: "auto" | "reference" | "duckdb"; + environment: "development" | "staging" | "production"; definition_hash: string; executor_version: string; source_fingerprints: Record[]; @@ -231,14 +233,37 @@ export type PipelineRun = { delivery_ref?: string | null; correlation_id?: string | null; causation_id?: string | null; + attempts: number; + max_attempts: number; + available_at?: string | null; + claimed_at?: string | null; + lease_expires_at?: string | null; + cancellation_requested_at?: string | null; + progress_percent: number; + progress_phase: string; + retention_until?: string | null; + purged_at?: string | null; error?: string | null; - started_at: string; + started_at?: string | null; finished_at?: string | null; created_by?: string | null; created_at: string; replayed: boolean; }; +export type PipelineDeployment = { + id: string; + pipeline_id: string; + revision: number; + environment: "development" | "staging" | "production"; + source_environment: "development" | "staging" | "production"; + status: string; + provenance: Record; + promoted_by?: string | null; + created_at: string; + updated_at: string; +}; + export type DataflowTriggerKind = "once" | "interval" | "event"; export type DataflowTrigger = { id: string; @@ -505,6 +530,10 @@ export function runDataflowPipeline( revision: number; idempotency_key: string; row_limit?: number; + execution_backend?: "reference" | "duckdb" | "auto"; + environment?: "development" | "staging" | "production"; + max_attempts?: number; + retention_days?: number; publication?: { target_datasource_ref?: string | null; name?: string | null; @@ -526,3 +555,45 @@ export function runDataflowPipeline( } ); } + +export function cancelDataflowPipelineRun( + settings: ApiSettings, + runRef: string +): Promise { + const runId = runRef.replace(/^dataflow-run:/, ""); + return apiFetch( + settings, + `/api/v1/dataflow/runs/${encodeURIComponent(runId)}/cancel`, + { method: "POST" } + ); +} + +export async function listDataflowPipelineDeployments( + settings: ApiSettings, + pipelineId: string +): Promise { + const response = await apiFetch<{ deployments: PipelineDeployment[] }>( + settings, + `/api/v1/dataflow/pipelines/${encodeURIComponent(pipelineId)}/deployments` + ); + return response.deployments; +} + +export function promoteDataflowPipeline( + settings: ApiSettings, + pipelineId: string, + payload: { + revision: number; + source_environment: "development" | "staging"; + target_environment: "staging" | "production"; + } +): Promise { + return apiFetch( + settings, + `/api/v1/dataflow/pipelines/${encodeURIComponent(pipelineId)}/promotions`, + { + method: "POST", + body: JSON.stringify(payload) + } + ); +} diff --git a/webui/src/features/dataflow/DataflowPage.tsx b/webui/src/features/dataflow/DataflowPage.tsx index 75afa2a..7414d72 100644 --- a/webui/src/features/dataflow/DataflowPage.tsx +++ b/webui/src/features/dataflow/DataflowPage.tsx @@ -8,6 +8,7 @@ import { } from "react"; import { CheckCircle2, + CircleStop, Clock3, Code2, CopyPlus, @@ -16,6 +17,7 @@ import { Play, Plus, RefreshCw, + Rocket, RotateCcw, Save, Settings2, @@ -46,6 +48,7 @@ import { ReactFlowProvider } from "@xyflow/react"; import { useLocation } from "react-router-dom"; import { compileDataflowSql, + cancelDataflowPipelineRun, createDataflowTrigger, createDataflowPipeline, createDataflowSourceSnapshot, @@ -55,10 +58,12 @@ import { deriveDataflowPipeline, listDataflowNodeTypes, listDataflowPipelineRuns, + listDataflowPipelineDeployments, listDataflowPipelines, listDataflowSources, listDataflowTriggers, previewDataflowPipeline, + promoteDataflowPipeline, runDataflowPipeline, renderDataflowSql, updateDataflowPipeline, @@ -75,6 +80,7 @@ import { type Pipeline, type PipelineGraphNode, type PipelinePreview, + type PipelineDeployment, type PipelineRun, type TabularSource } from "../../api/dataflow"; @@ -925,6 +931,7 @@ export default function DataflowPage({ settings, auth }: { settings: ApiSettings revision: draft.currentRevision } : null} sources={sources} + canPromote={canWrite} onClose={() => setRunOpen(false)} onCompleted={(run) => { if (run.status === "succeeded") { @@ -933,8 +940,10 @@ export default function DataflowPage({ settings, auth }: { settings: ApiSettings ? `Run published ${run.output_row_count} rows to ${run.output_datasource_ref}.` : `Run completed with ${run.output_row_count} output rows.` ); - } else { + } else if (run.status === "failed") { setError(run.error || "Dataflow run failed."); + } else { + setSuccess(`Run ${run.ref} was queued.`); } if (run.output_datasource_ref) { void listDataflowSources(settings).then((catalogue) => { @@ -1586,6 +1595,7 @@ function RunPipelineDialog({ settings, pipeline, sources, + canPromote, onClose, onCompleted }: { @@ -1593,6 +1603,7 @@ function RunPipelineDialog({ settings: ApiSettings; pipeline: { id: string; name: string; revision: number } | null; sources: TabularSource[]; + canPromote: boolean; onClose: () => void; onCompleted: (run: PipelineRun) => void; }) { @@ -1604,6 +1615,8 @@ function RunPipelineDialog({ const [freeze, setFreeze] = useState(false); const [frozenLabel, setFrozenLabel] = useState(""); const [idempotencyKey, setIdempotencyKey] = useState(""); + const [environment, setEnvironment] = useState<"development" | "staging" | "production">("development"); + const [deployments, setDeployments] = useState([]); const [runs, setRuns] = useState([]); const [loadingRuns, setLoadingRuns] = useState(false); const [busy, setBusy] = useState(false); @@ -1622,14 +1635,35 @@ function RunPipelineDialog({ setFreeze(false); setFrozenLabel(""); setIdempotencyKey(crypto.randomUUID()); + setEnvironment("development"); setError(""); setLoadingRuns(true); - void listDataflowPipelineRuns(settings, pipeline.id) - .then(setRuns) + void Promise.all([ + listDataflowPipelineRuns(settings, pipeline.id), + listDataflowPipelineDeployments(settings, pipeline.id) + ]) + .then(([loadedRuns, loadedDeployments]) => { + setRuns(loadedRuns); + setDeployments(loadedDeployments); + }) .catch((loadError) => setError(apiErrorMessage(loadError))) .finally(() => setLoadingRuns(false)); }, [open, pipeline?.id, pipeline?.revision, settings]); + const hasActiveRuns = runs.some((run) => + run.status === "queued" || run.status === "retrying" || run.status === "running" + ); + + useEffect(() => { + if (!open || !pipeline || !hasActiveRuns) return; + const timer = window.setInterval(() => { + void listDataflowPipelineRuns(settings, pipeline.id) + .then(setRuns) + .catch(() => undefined); + }, 2000); + return () => window.clearInterval(timer); + }, [hasActiveRuns, open, pipeline?.id, settings]); + const start = async () => { if (!pipeline) return; setBusy(true); @@ -1639,6 +1673,8 @@ function RunPipelineDialog({ revision: pipeline.revision, idempotency_key: idempotencyKey, row_limit: 500, + execution_backend: "auto", + environment, publication: mode === "publish" ? { target_datasource_ref: targetRef === "new" ? null : targetRef, name: targetRef === "new" ? name.trim() : null, @@ -1655,7 +1691,7 @@ function RunPipelineDialog({ ]); onCompleted(run); setIdempotencyKey(crypto.randomUUID()); - if (run.status !== "succeeded") { + if (run.status === "failed") { setError(run.error || "Dataflow run failed."); } } catch (runError) { @@ -1665,6 +1701,39 @@ function RunPipelineDialog({ } }; + const cancelRun = async (run: PipelineRun) => { + setError(""); + try { + const cancelled = await cancelDataflowPipelineRun(settings, run.ref); + setRuns((current) => current.map((item) => + item.ref === cancelled.ref ? cancelled : item + )); + } catch (cancelError) { + setError(apiErrorMessage(cancelError)); + } + }; + + const promote = async (target: "staging" | "production") => { + if (!pipeline) return; + setBusy(true); + setError(""); + try { + const deployment = await promoteDataflowPipeline(settings, pipeline.id, { + revision: pipeline.revision, + source_environment: target === "staging" ? "development" : "staging", + target_environment: target + }); + setDeployments((current) => [ + ...current.filter((item) => item.environment !== target), + deployment + ]); + } catch (promotionError) { + setError(apiErrorMessage(promotionError)); + } finally { + setBusy(false); + } + }; + const targetComplete = mode === "run" || targetRef !== "new" || Boolean(name.trim() && sourceName.trim()); @@ -1685,7 +1754,7 @@ function RunPipelineDialog({ onClick={() => void start()} disabled={busy || !pipeline || !targetComplete} > - {busy ? "Running..." : "Run revision"} + {busy ? "Queueing..." : "Queue run"} )} @@ -1710,6 +1779,44 @@ function RunPipelineDialog({ Revision {pipeline?.revision ?? "-"} +
+ + + + {canPromote ? ( + <> + + + + ) : null} +
{mode === "publish" ? (
@@ -1774,27 +1881,42 @@ function RunPipelineDialog({ - + + + {runs.map((run) => ( - + + + ))} {!loadingRuns && !runs.length ? ( - + ) : null}
StartedQueued Revision StatusProgress Rows Output
{new Date(run.started_at).toLocaleString()}{new Date(run.started_at ?? run.created_at).toLocaleString()} {run.revision} {run.progress_percent}% ยท {run.progress_phase} {run.output_row_count.toLocaleString()} {run.output_datasource_ref ?? "Run evidence"} + {run.status === "queued" + || run.status === "retrying" + || run.status === "running" ? ( + } + variant="danger" + onClick={() => void cancelRun(run)} + /> + ) : null} +
No runs yet
No runs yet