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

feat: run dataflows through durable workers

Browse Source
This commit is contained in:
zemion
2026-07-30 02:33:02 +02:00
parent d61e9d8942
commit 5af02933ef
16 changed files with 2306 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/govoplan_dataflow/backend/db/__init__.py
+12 -2
View File
@@ -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",
]
src/govoplan_dataflow/backend/db/models.py
+140 -1
View File
@@ -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__ = (
src/govoplan_dataflow/backend/executor.py
+2
View File
@@ -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)
src/govoplan_dataflow/backend/manifest.py
+21
View File
@@ -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,
src/govoplan_dataflow/backend/migrations/versions/f6c2a9d4e7b1_v0114_production_runs.py
+258
View File
@@ -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,
)
src/govoplan_dataflow/backend/router.py
+106 -1
View File
@@ -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,
src/govoplan_dataflow/backend/run_worker.py
+631
View File
@@ -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",
]
src/govoplan_dataflow/backend/schemas.py
+83 -4
View File
@@ -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
src/govoplan_dataflow/backend/service.py
+557 -35
View File
@@ -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",
src/govoplan_dataflow/backend/triggers.py
+28 -4
View File
@@ -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