Fence and reconcile Dataflow runs

src/govoplan_dataflow/backend/service.py

@@ -40,6 +40,7 @@ from govoplan_dataflow.backend.db.models import (
     DataflowPipelineDeployment,
     DataflowPipelineRevision,
     DataflowRun,
+    new_uuid,
 )
 from govoplan_dataflow.backend.executor import (
     EXECUTOR_VERSION,
@@ -79,6 +80,12 @@ from govoplan_dataflow.backend.schemas import (
     PipelineValidationResponse,
     PreviewColumn,
 )
+from govoplan_dataflow.backend.recovery import (
+    DataflowRecoveryError,
+    DataflowRunRecovery,
+    begin_dataflow_run_recovery,
+    dataflow_run_recovery_state,
+)
 from govoplan_dataflow.backend.sql_compiler import (
     SqlCompilationError,
     compile_sql,
@@ -1071,14 +1078,27 @@ def start_pipeline_run(
         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}"
         ),
     }
+    recovery: DataflowRunRecovery | None = None
+    if not defer_execution:
+        try:
+            recovery = begin_dataflow_run_recovery(
+                session,
+                run=run,
+                lease_ttl_seconds=int(
+                    float(run.resource_budget.get("max_wall_seconds") or 30.0)
+                )
+                + 60,
+            )
+        except DataflowRecoveryError as exc:
+            raise DataflowConflictError(str(exc)) from exc
+    session.add(run)
+    session.flush()
     if defer_execution:
         session.flush()
         return run, False
@@ -1090,8 +1110,27 @@ def start_pipeline_run(
         request=request,
         principal=principal,
         registry=registry,
+        recovery=recovery,
     )
-    session.flush()
+    if recovery is not None:
+        try:
+            recovery.finish(session, run=run)
+        except DataflowRecoveryError as exc:
+            session.rollback()
+            if recovery.publication_started:
+                persisted = session.get(DataflowRun, run.id)
+                if persisted is not None:
+                    persisted.status = "outcome_unknown"
+                    persisted.finished_at = utcnow()
+                    persisted.progress_phase = "outcome_unknown"
+                    persisted.error = (
+                        "Output publication completed without verifiable "
+                        "terminal recovery evidence; reconcile the sink."
+                    )
+                    session.commit()
+            raise DataflowConflictError(str(exc)) from exc
+    else:
+        session.flush()
     return run, False
 
 
@@ -1256,6 +1295,7 @@ def _new_pipeline_run(
     now = utcnow()
     budget = _run_resource_budget(request)
     return DataflowRun(
+        id=new_uuid(),
         tenant_id=tenant_id,
         pipeline_id=pipeline.id,
         pipeline_revision_id=revision.id,
@@ -1360,6 +1400,7 @@ def _execute_pipeline_run(
     request: DataflowRunRequest,
     principal: ApiPrincipal,
     registry: object | None,
+    recovery: DataflowRunRecovery | None = None,
 ) -> bool:
     try:
         if run.cancellation_requested_at is not None:
@@ -1405,6 +1446,7 @@ def _execute_pipeline_run(
                 result=result,
                 principal=principal,
                 registry=registry,
+                recovery=recovery,
             )
         run.status = "succeeded"
         run.finished_at = utcnow()
@@ -1414,6 +1456,10 @@ def _execute_pipeline_run(
         return False
     except (DatasourceError, PipelineExecutionError) as exc:
         _mark_pipeline_run_failed(run, exc)
+        if recovery is not None and recovery.publication_started:
+            run.status = "outcome_unknown"
+            run.progress_phase = "outcome_unknown"
+            return False
         return isinstance(exc, DatasourceUnavailableError) or bool(
             getattr(exc, "retryable", False)
         )
@@ -1469,6 +1515,7 @@ def _publish_pipeline_result(
     result: PipelineExecutionResult,
     principal: ApiPrincipal,
     registry: object | None,
+    recovery: DataflowRunRecovery | None = None,
 ) -> None:
     publisher = datasource_publication(registry)
     if publisher is None:
@@ -1479,33 +1526,51 @@ def _publish_pipeline_result(
     target = request.publication
     if target is None:
         return
-    publication = publisher.publish_rows(
+    if recovery is None:
+        raise PipelineExecutionError(
+            "Publishing Dataflow output requires a durable recovery operation."
+        )
+    recovery.prepare_publication(
         session,
-        principal,
-        request=DatasourcePublicationRequest(
-            producer_module="dataflow",
-            producer_run_ref=f"dataflow-run:{run.id}",
-            idempotency_key=f"{pipeline.id}:{request.idempotency_key.strip()}",
-            rows=tuple(dict(row) for row in result.rows),
-            target_datasource_ref=target.target_datasource_ref,
-            name=target.name or f"{pipeline.name} output",
-            source_name=target.source_name,
-            description=target.description,
-            freeze=target.freeze,
-            frozen_label=target.frozen_label,
-            set_current=target.set_current,
-            provenance={
-                "pipeline_ref": f"pipeline:{pipeline.id}",
-                "pipeline_revision": revision.revision,
-                "definition_hash": revision.content_hash,
-                "source_fingerprints": result.source_fingerprints,
-            },
-            metadata={
-                **dict(target.metadata),
-                "dataflow_run_ref": f"dataflow-run:{run.id}",
-            },
-        ),
+        run=run,
+        rows=tuple(dict(row) for row in result.rows),
     )
+    try:
+        publication = publisher.publish_rows(
+            session,
+            principal,
+            request=DatasourcePublicationRequest(
+                producer_module="dataflow",
+                producer_run_ref=f"dataflow-run:{run.id}",
+                idempotency_key=(
+                    f"{pipeline.id}:{request.idempotency_key.strip()}"
+                ),
+                rows=tuple(dict(row) for row in result.rows),
+                target_datasource_ref=target.target_datasource_ref,
+                name=target.name or f"{pipeline.name} output",
+                source_name=target.source_name,
+                description=target.description,
+                freeze=target.freeze,
+                frozen_label=target.frozen_label,
+                set_current=target.set_current,
+                provenance={
+                    "pipeline_ref": f"pipeline:{pipeline.id}",
+                    "pipeline_revision": revision.revision,
+                    "definition_hash": revision.content_hash,
+                    "source_fingerprints": result.source_fingerprints,
+                },
+                metadata={
+                    **dict(target.metadata),
+                    "dataflow_run_ref": f"dataflow-run:{run.id}",
+                },
+            ),
+        )
+    except DatasourceError:
+        raise
+    except Exception as exc:
+        raise PipelineExecutionError(
+            "The output provider failed after publication dispatch began."
+        ) from exc
     run.output_publication_ref = publication.ref
     run.output_datasource_ref = publication.datasource.ref
     run.output_materialization_ref = publication.materialization.ref
@@ -1573,10 +1638,16 @@ def pipeline_run_response(
     run: DataflowRun,
     *,
     replayed: bool = False,
+    recovery_state: Mapping[str, object] | None = None,
 ) -> PipelineRunResponse:
     revision = session.get(DataflowPipelineRevision, run.pipeline_revision_id)
     if revision is None:
         raise DataflowNotFoundError("Dataflow pipeline revision not found")
+    recovery = (
+        dict(recovery_state) or None
+        if recovery_state is not None
+        else dataflow_run_recovery_state(session, run_id=run.id)
+    )
     return PipelineRunResponse(
         ref=f"dataflow-run:{run.id}",
         pipeline_id=run.pipeline_id,
@@ -1622,6 +1693,27 @@ def pipeline_run_response(
         finished_at=run.finished_at,
         created_by=run.created_by,
         created_at=run.created_at,
+        recovery_operation_id=(
+            str(recovery["operation_id"]) if recovery is not None else None
+        ),
+        recovery_operation_type=(
+            str(recovery["operation_type"]) if recovery is not None else None
+        ),
+        recovery_mode=(
+            str(recovery["mode"]) if recovery is not None else None
+        ),
+        recovery_status=(
+            str(recovery["status"]) if recovery is not None else None
+        ),
+        recovery_requires_attention=(
+            bool(recovery["requires_attention"])
+            or run.status == "outcome_unknown"
+            if recovery is not None
+            else run.status == "outcome_unknown"
+        ),
+        recovery_explanation=(
+            str(recovery["explanation"]) if recovery is not None else None
+        ),
         replayed=replayed,
     )
 
@@ -1635,6 +1727,7 @@ def pipeline_run_descriptor(
     revision = session.get(DataflowPipelineRevision, run.pipeline_revision_id)
     if revision is None:
         raise DataflowNotFoundError("Dataflow pipeline revision not found")
+    recovery = dataflow_run_recovery_state(session, run_id=run.id)
     return DataflowRunDescriptor(
         ref=f"dataflow-run:{run.id}",
         pipeline_ref=f"pipeline:{run.pipeline_id}",
@@ -1670,6 +1763,7 @@ def pipeline_run_descriptor(
             "progress_phase": run.progress_phase,
             "source_fingerprints": list(run.source_fingerprints),
             "diagnostics": list(run.diagnostics),
+            "recovery": dict(recovery) if recovery is not None else None,
         },
     )