Files
govoplan-dataflow/src/govoplan_dataflow/backend/service.py
T

2231 lines
72 KiB
Python

from __future__ import annotations
from collections.abc import Mapping
import hashlib
import json
from dataclasses import dataclass
from datetime import datetime, timedelta
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,
DataflowRunNotFoundError,
DataflowRunRequest,
)
from govoplan_core.core.datasources import (
DatasourceError,
DatasourcePublicationRequest,
DatasourceReadRequest,
DatasourceUnavailableError,
datasource_catalogue,
datasource_publication,
)
from govoplan_core.db.base import utcnow
from govoplan_dataflow.backend.backends import (
BackendExecutionError,
BackendSource,
ExecutionBudget,
execute_typed_graph,
)
from govoplan_dataflow.backend.batches import TypedBatch
from govoplan_dataflow.backend.db.models import (
DataflowPipeline,
DataflowPipelineDeployment,
DataflowPipelineRevision,
DataflowRun,
new_uuid,
)
from govoplan_dataflow.backend.executor import (
EXECUTOR_VERSION,
MAX_SOURCE_ROWS,
PipelineExecutionError,
PipelineExecutionResult,
ResolvedSource,
execute_preview,
)
from govoplan_dataflow.backend.governance import (
definition_governance_payload,
require_definition_action,
)
from govoplan_dataflow.backend.graph import (
canonical_graph_payload,
definition_hash,
preserve_compatible_graph_layout,
validate_graph,
)
from govoplan_dataflow.backend.schemas import (
DataflowDiagnostic,
GraphNode,
NodePreviewResult,
PipelineCreateRequest,
PipelineDeriveRequest,
PipelineDraftRequest,
PipelineGraph,
PipelinePreviewRequest,
PipelinePreviewResponse,
PipelineDeploymentResponse,
PipelinePromotionRequest,
PipelineResponse,
PipelineRevisionResponse,
PipelineRunResponse,
PipelineSqlResponse,
PipelineUpdateRequest,
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,
render_sql,
)
class DataflowError(DataflowRunError):
pass
class DataflowNotFoundError(DataflowError, DataflowRunNotFoundError):
pass
class DataflowConflictError(DataflowError, DataflowRunConflictError):
pass
class DataflowValidationError(DataflowError):
def __init__(self, diagnostics: list[DataflowDiagnostic]) -> None:
super().__init__(diagnostics[0].message if diagnostics else "Pipeline validation failed")
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
sql_text: str | None
diagnostics: list[DataflowDiagnostic]
def list_pipelines(session: Session, *, tenant_id: str) -> list[DataflowPipeline]:
return list(
session.scalars(
select(DataflowPipeline)
.where(
or_(
DataflowPipeline.tenant_id == tenant_id,
DataflowPipeline.tenant_id.is_(None),
),
DataflowPipeline.deleted_at.is_(None),
)
.order_by(DataflowPipeline.updated_at.desc(), DataflowPipeline.name)
)
)
def get_pipeline(
session: Session,
*,
tenant_id: str,
pipeline_id: str,
) -> DataflowPipeline:
pipeline = session.scalar(
select(DataflowPipeline).where(
DataflowPipeline.id == pipeline_id,
or_(
DataflowPipeline.tenant_id == tenant_id,
DataflowPipeline.tenant_id.is_(None),
),
DataflowPipeline.deleted_at.is_(None),
)
)
if pipeline is None:
raise DataflowNotFoundError("Dataflow pipeline not found")
return pipeline
def get_pipeline_revision(
session: Session,
*,
pipeline: DataflowPipeline,
revision: int | None = None,
) -> DataflowPipelineRevision:
revision_number = revision or pipeline.current_revision
item = session.scalar(
select(DataflowPipelineRevision).where(
DataflowPipelineRevision.pipeline_id == pipeline.id,
DataflowPipelineRevision.tenant_id == pipeline.tenant_id,
DataflowPipelineRevision.revision == revision_number,
)
)
if item is None:
raise DataflowNotFoundError("Dataflow pipeline revision not found")
return item
def create_pipeline(
session: Session,
*,
tenant_id: str,
actor_id: str | None,
payload: PipelineCreateRequest,
) -> DataflowPipeline:
definition = normalize_definition(
graph=payload.graph,
sql_text=payload.sql_text,
editor_mode=payload.editor_mode,
)
content_hash = definition_hash(definition.graph, definition.sql_text)
stored_tenant_id = None if payload.scope_type == "system" else tenant_id
scope_id = (
None
if payload.scope_type == "system"
else tenant_id
if payload.scope_type == "tenant"
else payload.scope_id
)
pipeline = DataflowPipeline(
tenant_id=stored_tenant_id,
scope_type=payload.scope_type,
scope_id=scope_id,
definition_kind=payload.definition_kind,
inherit_to_lower_scopes=payload.inherit_to_lower_scopes,
allow_run=payload.allow_run,
allow_reuse=payload.allow_reuse,
allow_automation=payload.allow_automation,
name=payload.name.strip(),
description=_clean_optional(payload.description),
status=payload.status,
current_revision=1,
created_by=actor_id,
updated_by=actor_id,
metadata_={},
)
revision = DataflowPipelineRevision(
tenant_id=stored_tenant_id,
revision=1,
schema_version=definition.graph.schema_version,
graph=canonical_graph_payload(definition.graph),
sql_text=definition.sql_text,
editor_mode=payload.editor_mode,
content_hash=content_hash,
created_by=actor_id,
)
pipeline.revisions.append(revision)
session.add(pipeline)
session.flush()
return pipeline
def update_pipeline(
session: Session,
*,
tenant_id: str,
pipeline_id: str,
actor_id: str | None,
payload: PipelineUpdateRequest,
) -> DataflowPipeline:
pipeline = get_pipeline(session, tenant_id=tenant_id, pipeline_id=pipeline_id)
if payload.expected_revision != pipeline.current_revision:
raise DataflowConflictError(
f"Pipeline changed on the server; expected revision {payload.expected_revision}, "
f"current revision is {pipeline.current_revision}"
)
if (
payload.scope_type != pipeline.scope_type
or payload.scope_id != pipeline.scope_id
and not (
pipeline.scope_type == "tenant"
and payload.scope_id in {None, pipeline.scope_id}
)
):
raise DataflowConflictError(
"Definition scope is immutable; derive a scoped copy instead."
)
if payload.definition_kind != pipeline.definition_kind:
raise DataflowConflictError(
"Definition kind is immutable; derive a flow or template instead."
)
definition = normalize_definition(
graph=payload.graph,
sql_text=payload.sql_text,
editor_mode=payload.editor_mode,
)
content_hash = definition_hash(definition.graph, definition.sql_text)
current = get_pipeline_revision(session, pipeline=pipeline)
pipeline.name = payload.name.strip()
pipeline.description = _clean_optional(payload.description)
pipeline.status = payload.status
ancestor_limits = _ancestor_governance_limits(
pipeline.derivation_provenance
)
pipeline.inherit_to_lower_scopes = (
payload.inherit_to_lower_scopes
and ancestor_limits["inherit_to_lower_scopes"]
)
pipeline.allow_run = payload.allow_run and ancestor_limits["allow_run"]
pipeline.allow_reuse = (
payload.allow_reuse and ancestor_limits["allow_reuse"]
)
pipeline.allow_automation = (
payload.allow_automation and ancestor_limits["allow_automation"]
)
pipeline.updated_by = actor_id
if current.content_hash != content_hash or current.editor_mode != payload.editor_mode:
pipeline.current_revision += 1
pipeline.revisions.append(
DataflowPipelineRevision(
tenant_id=pipeline.tenant_id,
revision=pipeline.current_revision,
schema_version=definition.graph.schema_version,
graph=canonical_graph_payload(definition.graph),
sql_text=definition.sql_text,
editor_mode=payload.editor_mode,
content_hash=content_hash,
created_by=actor_id,
)
)
session.flush()
return pipeline
def derive_pipeline(
session: Session,
*,
tenant_id: str,
actor_id: str | None,
principal: ApiPrincipal,
registry: object | None,
source_pipeline_id: str,
payload: PipelineDeriveRequest,
) -> DataflowPipeline:
source = get_pipeline(
session,
tenant_id=tenant_id,
pipeline_id=source_pipeline_id,
)
reuse_decision = require_definition_action(
source,
principal=principal,
registry=registry,
action="derive",
)
source_revision = get_pipeline_revision(
session,
pipeline=source,
revision=payload.source_revision,
)
stored_tenant_id = None if payload.scope_type == "system" else tenant_id
scope_id = (
None
if payload.scope_type == "system"
else tenant_id
if payload.scope_type == "tenant"
else payload.scope_id
)
source_limits = _effective_governance_limits(
source,
decision_details=reuse_decision.details,
)
effective_limits = {
"inherit_to_lower_scopes": (
source_limits["inherit_to_lower_scopes"]
and payload.inherit_to_lower_scopes
),
"allow_run": source_limits["allow_run"] and payload.allow_run,
"allow_reuse": source_limits["allow_reuse"] and payload.allow_reuse,
"allow_automation": (
source_limits["allow_automation"]
and payload.allow_automation
),
}
provenance = {
"source_ref": f"pipeline:{source.id}",
"source_scope": {
"scope_type": source.scope_type,
"scope_id": source.scope_id,
},
"source_definition_kind": source.definition_kind,
"source_revision": source_revision.revision,
"source_hash": source_revision.content_hash,
"source_effective_limits": effective_limits,
"policy_decision": reuse_decision.to_dict(),
"derived_by": actor_id,
"derived_at": utcnow().isoformat(),
}
pipeline = DataflowPipeline(
tenant_id=stored_tenant_id,
scope_type=payload.scope_type,
scope_id=scope_id,
definition_kind=payload.definition_kind,
inherit_to_lower_scopes=effective_limits[
"inherit_to_lower_scopes"
],
allow_run=effective_limits["allow_run"],
allow_reuse=effective_limits["allow_reuse"],
allow_automation=effective_limits["allow_automation"],
derived_from_pipeline_id=source.id,
derived_from_revision=source_revision.revision,
derived_from_hash=source_revision.content_hash,
derivation_provenance=provenance,
name=payload.name.strip(),
description=_clean_optional(payload.description),
status="draft",
current_revision=1,
created_by=actor_id,
updated_by=actor_id,
metadata_={},
)
pipeline.revisions.append(
DataflowPipelineRevision(
tenant_id=stored_tenant_id,
revision=1,
schema_version=source_revision.schema_version,
graph=dict(source_revision.graph),
sql_text=source_revision.sql_text,
editor_mode=source_revision.editor_mode,
content_hash=source_revision.content_hash,
created_by=actor_id,
)
)
session.add(pipeline)
session.flush()
return pipeline
def delete_pipeline(
session: Session,
*,
tenant_id: str,
pipeline_id: str,
actor_id: str | None,
) -> DataflowPipeline:
pipeline = get_pipeline(session, tenant_id=tenant_id, pipeline_id=pipeline_id)
pipeline.deleted_at = utcnow()
pipeline.updated_by = actor_id
session.flush()
return pipeline
def pipeline_response(
session: Session,
pipeline: DataflowPipeline,
*,
principal: ApiPrincipal,
registry: object | None,
) -> PipelineResponse:
revision = get_pipeline_revision(session, pipeline=pipeline)
return PipelineResponse(
id=pipeline.id,
tenant_id=pipeline.tenant_id,
name=pipeline.name,
description=pipeline.description,
status=pipeline.status,
current_revision=pipeline.current_revision,
created_by=pipeline.created_by,
updated_by=pipeline.updated_by,
created_at=pipeline.created_at,
updated_at=pipeline.updated_at,
revision=PipelineRevisionResponse.model_validate(revision),
governance=definition_governance_payload(
pipeline,
principal=principal,
registry=registry,
),
)
def validate_draft(payload: PipelineDraftRequest) -> PipelineValidationResponse:
if payload.sql_text and payload.sql_text.strip():
try:
graph, sql_text, diagnostics = compile_sql(
payload.sql_text,
source_nodes=_source_nodes(payload.graph, payload.source_nodes),
)
except SqlCompilationError as exc:
return PipelineValidationResponse(
valid=False,
graph=payload.graph,
sql_text=payload.sql_text,
diagnostics=exc.diagnostics,
)
return PipelineValidationResponse(
valid=True,
graph=graph,
sql_text=sql_text,
diagnostics=diagnostics,
)
if payload.graph is None:
diagnostic = DataflowDiagnostic(
severity="error",
code="definition.required",
message="Provide a graph or SQL query.",
)
return PipelineValidationResponse(
valid=False,
graph=None,
sql_text=None,
diagnostics=[diagnostic],
)
diagnostics = validate_graph(payload.graph)
sql_text: str | None = None
if not any(item.severity == "error" for item in diagnostics):
try:
sql_text, render_diagnostics = render_sql(payload.graph)
diagnostics.extend(render_diagnostics)
except SqlCompilationError as exc:
diagnostics.extend(
DataflowDiagnostic(
severity="warning",
code=item.code,
message=item.message,
node_id=item.node_id,
field=item.field,
)
for item in exc.diagnostics
)
return PipelineValidationResponse(
valid=not any(item.severity == "error" for item in diagnostics),
graph=payload.graph,
sql_text=sql_text,
diagnostics=diagnostics,
)
def compile_sql_draft(payload: PipelineDraftRequest) -> PipelineSqlResponse:
if not payload.sql_text:
diagnostic = DataflowDiagnostic(
severity="error",
code="sql.empty",
message="Enter a SELECT query.",
field="sql_text",
)
return PipelineSqlResponse(valid=False, graph=payload.graph, sql_text="", diagnostics=[diagnostic])
try:
graph, sql_text, diagnostics = compile_sql(
payload.sql_text,
source_nodes=_source_nodes(payload.graph, payload.source_nodes),
)
if payload.graph is not None:
graph = preserve_compatible_graph_layout(payload.graph, graph)
except SqlCompilationError as exc:
return PipelineSqlResponse(
valid=False,
graph=payload.graph,
sql_text=payload.sql_text,
diagnostics=exc.diagnostics,
)
return PipelineSqlResponse(valid=True, graph=graph, sql_text=sql_text, diagnostics=diagnostics)
def render_graph_sql(payload: PipelineDraftRequest) -> PipelineSqlResponse:
if payload.graph is None:
diagnostic = DataflowDiagnostic(
severity="error",
code="graph.required",
message="Provide a graph to render.",
)
return PipelineSqlResponse(valid=False, graph=None, sql_text=None, diagnostics=[diagnostic])
try:
sql_text, diagnostics = render_sql(payload.graph)
except SqlCompilationError as exc:
return PipelineSqlResponse(
valid=False,
graph=payload.graph,
sql_text=None,
diagnostics=exc.diagnostics,
)
return PipelineSqlResponse(valid=True, graph=payload.graph, sql_text=sql_text, diagnostics=diagnostics)
def preview_pipeline(
session: Session,
*,
tenant_id: str,
actor_id: str | None,
payload: PipelinePreviewRequest,
principal: ApiPrincipal | None = None,
registry: object | None = None,
) -> PipelinePreviewResponse:
pipeline: DataflowPipeline | None = None
revision: DataflowPipelineRevision | None = None
if payload.pipeline_id:
pipeline = get_pipeline(session, tenant_id=tenant_id, pipeline_id=payload.pipeline_id)
if principal is None:
raise DataflowConflictError(
"Saved pipeline previews require a tenant API principal."
)
action = "edit" if pipeline.status == "draft" else "run"
try:
require_definition_action(
pipeline,
principal=principal,
registry=registry,
action=action,
)
except PermissionError as exc:
raise DataflowConflictError(str(exc)) from exc
revision = get_pipeline_revision(session, pipeline=pipeline, revision=payload.revision)
graph = PipelineGraph.model_validate(revision.graph)
sql_text = revision.sql_text
else:
draft = PipelineDraftRequest(
graph=payload.graph,
sql_text=payload.sql_text,
source_nodes=payload.source_nodes,
)
validated = validate_draft(draft)
if not validated.valid or validated.graph is None:
return PipelinePreviewResponse(
run_id=None,
pipeline_id=None,
revision=None,
status="failed",
columns=[],
rows=[],
total_rows=0,
truncated=False,
diagnostics=validated.diagnostics,
node_diagnostics=[],
node_preview=None,
source_fingerprints=[],
input_row_count=0,
definition_hash="",
executor_version=EXECUTOR_VERSION,
)
graph = validated.graph
sql_text = validated.sql_text
graph_hash = definition_hash(graph, sql_text)
started_at = utcnow()
run: DataflowRun | None = None
try:
result, executor_version = _execute_pipeline_preview(
graph,
session=session,
principal=principal,
registry=registry,
backend=payload.execution_backend,
row_limit=payload.row_limit,
preview_node_id=payload.preview_node_id,
)
status = "succeeded"
error = None
diagnostics = result.diagnostics
columns = result.columns
rows = result.rows
total_rows = result.total_rows
truncated = result.truncated
node_diagnostics = result.node_diagnostics
node_preview = result.node_preview
source_fingerprints = result.source_fingerprints
input_row_count = result.input_row_count
except PipelineExecutionError as exc:
executor_version = (
payload.execution_backend
if payload.execution_backend != "reference"
else EXECUTOR_VERSION
)
status = "failed"
error = str(exc)
diagnostics = [
*exc.diagnostics,
DataflowDiagnostic(
severity="error",
code="preview.execution",
message=str(exc),
node_id=exc.node_id,
)
]
columns = []
rows = []
total_rows = 0
truncated = False
node_diagnostics = list(exc.node_diagnostics)
node_preview = exc.node_preview
source_fingerprints = list(exc.source_fingerprints)
input_row_count = exc.input_row_count
if pipeline is not None and revision is not None:
run = DataflowRun(
tenant_id=tenant_id,
pipeline_id=pipeline.id,
pipeline_revision_id=revision.id,
run_type="preview",
status=status,
executor_version=executor_version,
definition_hash=graph_hash,
source_fingerprints=source_fingerprints,
result_schema=[item.model_dump(mode="json") for item in columns],
diagnostics=[item.model_dump(mode="json") for item in diagnostics],
input_row_count=input_row_count,
output_row_count=total_rows,
started_at=started_at,
finished_at=utcnow(),
error=error,
created_by=actor_id,
)
session.add(run)
session.flush()
return PipelinePreviewResponse(
run_id=run.id if run else None,
pipeline_id=pipeline.id if pipeline else None,
revision=revision.revision if revision else None,
status=status,
columns=columns,
rows=rows,
total_rows=total_rows,
truncated=truncated,
diagnostics=diagnostics,
node_diagnostics=node_diagnostics,
node_preview=node_preview,
source_fingerprints=source_fingerprints,
input_row_count=input_row_count,
definition_hash=graph_hash,
executor_version=executor_version,
)
def _execute_pipeline_preview(
graph: PipelineGraph,
*,
session: Session,
principal: ApiPrincipal | None,
registry: object | None,
backend: str,
row_limit: int,
preview_node_id: str | None,
budget: ExecutionBudget | None = None,
) -> tuple[PipelineExecutionResult, str]:
source_resolver = _preview_source_resolver(
session=session,
principal=principal,
registry=registry,
)
if backend == "reference":
return (
execute_preview(
graph,
row_limit=row_limit,
source_resolver=source_resolver,
preview_node_id=preview_node_id,
),
EXECUTOR_VERSION,
)
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=budget or ExecutionBudget(max_output_rows=row_limit),
preview_node_id=preview_node_id,
)
except BackendExecutionError as exc:
raise PipelineExecutionError(
str(exc),
node_id=exc.node_id,
diagnostics=tuple(exc.diagnostics),
retryable=exc.code == "backend.capacity",
) from exc
columns = [
PreviewColumn(
name=field.name,
type=field.type,
nullable=field.nullable,
)
for field in result.batch.schema.fields
]
source_fingerprints = list(
result.metadata.get(
"source_fingerprints",
result.contract.lineage.source_fingerprints,
)
)
node_preview = _typed_node_preview(
result,
preview_node_id=preview_node_id,
columns=columns,
)
return (
PipelineExecutionResult(
rows=result.rows,
total_rows=result.contract.row_count,
truncated=result.contract.truncated,
columns=columns,
diagnostics=list(result.contract.diagnostics),
node_diagnostics=list(result.node_diagnostics),
node_preview=node_preview,
source_fingerprints=source_fingerprints,
input_row_count=int(
result.metadata.get(
"input_row_count",
sum(
int(
item.get(
"row_count",
item.get("total_rows", 0),
)
)
for item in source_fingerprints
),
)
),
),
result.contract.backend_version,
)
def _preview_source_resolver(
*,
session: Session,
principal: ApiPrincipal | None,
registry: object | None,
):
if principal is not None:
return _datasource_source_resolver(
session=session,
principal=principal,
registry=registry,
)
def unavailable(node: GraphNode, _limit: int) -> ResolvedSource:
raise PipelineExecutionError(
"Datasource-backed preview requires a tenant API principal.",
node_id=node.id,
)
return unavailable
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, source_limit)
sources[node.id] = BackendSource(
node_id=node.id,
batch=TypedBatch.from_rows(resolved.rows),
source_ref=resolved.source_ref,
provider=resolved.provider,
fingerprint=resolved.fingerprint,
total_rows=resolved.total_rows,
truncated=resolved.truncated,
source_name=str(node.config.get("source_name") or ""),
)
return sources
def _typed_node_preview(
result,
*,
preview_node_id: str | None,
columns: list[PreviewColumn],
) -> NodePreviewResult | None:
if result.node_preview is not None:
return result.node_preview
if preview_node_id is None:
return None
return NodePreviewResult(
node_id=preview_node_id,
columns=columns,
rows=result.rows,
total_rows=result.contract.row_count,
truncated=result.contract.truncated,
)
def list_pipeline_runs(
session: Session,
*,
tenant_id: str,
pipeline_id: str | None = None,
limit: int = 100,
) -> list[DataflowRun]:
statement = (
select(DataflowRun)
.where(DataflowRun.tenant_id == tenant_id)
.order_by(DataflowRun.created_at.desc(), DataflowRun.id.desc())
.limit(max(1, min(int(limit), 100)))
)
if pipeline_id:
get_pipeline(
session,
tenant_id=tenant_id,
pipeline_id=pipeline_id,
)
statement = statement.where(DataflowRun.pipeline_id == pipeline_id)
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,
*,
tenant_id: str,
run_ref: str,
) -> DataflowRun:
run_id = _strip_ref(run_ref, "dataflow-run:")
if not run_id:
raise DataflowNotFoundError("Dataflow run not found")
run = session.scalar(
select(DataflowRun).where(
DataflowRun.id == run_id,
DataflowRun.tenant_id == tenant_id,
)
)
if run is None:
raise DataflowNotFoundError("Dataflow run not found")
return run
def start_pipeline_run(
session: Session,
*,
tenant_id: str,
actor_id: str | None,
principal: ApiPrincipal,
registry: object | None,
request: DataflowRunRequest,
defer_execution: bool = False,
) -> tuple[DataflowRun, bool]:
pipeline, revision = _run_definition(
session,
tenant_id=tenant_id,
principal=principal,
registry=registry,
request=request,
)
idempotency_key, request_hash = _validated_run_identity(request)
existing = _existing_pipeline_run(
session,
tenant_id=tenant_id,
pipeline_id=pipeline.id,
idempotency_key=idempotency_key,
request_hash=request_hash,
)
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,
pipeline=pipeline,
revision=revision,
request=request,
idempotency_key=idempotency_key,
request_hash=request_hash,
principal=principal,
defer_execution=defer_execution,
)
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
_execute_pipeline_run(
session,
run=run,
pipeline=pipeline,
revision=revision,
request=request,
principal=principal,
registry=registry,
recovery=recovery,
)
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
def _run_definition(
session: Session,
*,
tenant_id: str,
principal: ApiPrincipal,
registry: object | None,
request: DataflowRunRequest,
) -> tuple[DataflowPipeline, DataflowPipelineRevision]:
pipeline_id = _strip_ref(request.pipeline_ref, "pipeline:")
if not pipeline_id:
raise DataflowNotFoundError("Dataflow pipeline not found")
pipeline = get_pipeline(
session,
tenant_id=tenant_id,
pipeline_id=pipeline_id,
)
action = (
"run"
if request.invocation.kind in {"manual", "api", "backfill"}
else "automate"
)
try:
require_definition_action(
pipeline,
principal=principal,
registry=registry,
action=action,
)
except PermissionError as exc:
raise DataflowConflictError(str(exc)) from exc
revision = get_pipeline_revision(
session,
pipeline=pipeline,
revision=request.revision,
)
_require_deployed_revision(
session,
tenant_id=tenant_id,
pipeline=pipeline,
revision=revision,
environment=request.environment,
)
return pipeline, revision
def _validated_run_identity(request: DataflowRunRequest) -> tuple[str, str]:
idempotency_key = request.idempotency_key.strip()
if not idempotency_key or len(idempotency_key) > 255:
raise DataflowConflictError(
"A Dataflow run idempotency key of at most 255 characters is required."
)
if request.row_limit < 1 or request.row_limit > MAX_PRODUCTION_ROWS:
raise DataflowConflictError(
"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,
*,
tenant_id: str,
pipeline_id: str,
idempotency_key: str,
request_hash: str,
) -> DataflowRun | None:
existing = session.scalar(
select(DataflowRun).where(
DataflowRun.tenant_id == tenant_id,
DataflowRun.pipeline_id == pipeline_id,
DataflowRun.idempotency_key == idempotency_key,
)
)
if existing is not None:
if existing.request_hash != request_hash:
raise DataflowConflictError(
"The Dataflow run idempotency key was already used with "
"different parameters."
)
return existing
def _new_pipeline_run(
*,
tenant_id: str,
actor_id: str | None,
pipeline: DataflowPipeline,
revision: DataflowPipelineRevision,
request: DataflowRunRequest,
idempotency_key: str,
request_hash: str,
principal: ApiPrincipal,
defer_execution: bool,
) -> DataflowRun:
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,
run_type="published" if request.publication else "run",
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,
request_hash=request_hash,
request_=_run_request_payload(request),
invocation_kind=request.invocation.kind,
trigger_id=_strip_ref(
request.invocation.trigger_ref or "",
"dataflow-trigger:",
),
trigger_delivery_id=_strip_ref(
request.invocation.delivery_ref or "",
"dataflow-trigger-delivery:",
),
correlation_id=request.invocation.correlation_id,
causation_id=request.invocation.causation_id,
source_fingerprints=[],
result_schema=[],
diagnostics=[],
input_row_count=0,
output_row_count=0,
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,
*,
run: DataflowRun,
pipeline: DataflowPipeline,
revision: DataflowPipelineRevision,
request: DataflowRunRequest,
principal: ApiPrincipal,
registry: object | None,
recovery: DataflowRunRecovery | None = None,
) -> bool:
try:
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,
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(
session,
run=run,
pipeline=pipeline,
revision=revision,
request=request,
result=result,
principal=principal,
registry=registry,
recovery=recovery,
)
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)
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)
)
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(
run: DataflowRun,
result: PipelineExecutionResult,
) -> None:
run.source_fingerprints = result.source_fingerprints
run.result_schema = [
item.model_dump(mode="json") for item in result.columns
]
run.diagnostics = [
item.model_dump(mode="json") for item in result.diagnostics
]
run.input_row_count = result.input_row_count
run.output_row_count = result.total_rows
def _ensure_publishable(result: PipelineExecutionResult) -> None:
source_truncated = any(
bool(item.get("truncated"))
for item in result.source_fingerprints
)
if result.truncated or source_truncated:
raise PipelineExecutionError(
"The bounded runner cannot publish a truncated result or a "
"result calculated from truncated source data."
)
def _publish_pipeline_result(
session: Session,
*,
run: DataflowRun,
pipeline: DataflowPipeline,
revision: DataflowPipelineRevision,
request: DataflowRunRequest,
result: PipelineExecutionResult,
principal: ApiPrincipal,
registry: object | None,
recovery: DataflowRunRecovery | None = None,
) -> None:
publisher = datasource_publication(registry)
if publisher is None:
raise PipelineExecutionError(
"Publishing Dataflow output requires the Datasources "
"publication capability."
)
target = request.publication
if target is None:
return
if recovery is None:
raise PipelineExecutionError(
"Publishing Dataflow output requires a durable recovery operation."
)
recovery.prepare_publication(
session,
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
def _mark_pipeline_run_failed(
run: DataflowRun,
exc: DatasourceError | PipelineExecutionError,
) -> None:
run.status = "failed"
run.finished_at = utcnow()
run.error = str(exc)
run.progress_phase = "failed"
diagnostics = list(getattr(exc, "diagnostics", ()))
diagnostics.append(
DataflowDiagnostic(
severity="error",
code="run.execution",
message=str(exc),
node_id=getattr(exc, "node_id", None),
)
)
run.diagnostics = [
item.model_dump(mode="json") for item in diagnostics
]
run.source_fingerprints = list(
getattr(exc, "source_fingerprints", ())
)
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,
*,
tenant_id: str,
run_ref: str,
) -> DataflowRun:
run = get_pipeline_run(
session,
tenant_id=tenant_id,
run_ref=run_ref,
)
if run.status not in {"queued", "retrying", "running"}:
raise DataflowConflictError(
f"Dataflow run is already {run.status} and cannot be cancelled."
)
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
def pipeline_run_response(
session: Session,
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,
revision=revision.revision,
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),
result_schema=list(run.result_schema),
diagnostics=list(run.diagnostics),
input_row_count=run.input_row_count,
output_row_count=run.output_row_count,
output_publication_ref=run.output_publication_ref,
output_datasource_ref=run.output_datasource_ref,
output_materialization_ref=run.output_materialization_ref,
invocation_kind=run.invocation_kind,
trigger_ref=(
f"dataflow-trigger:{run.trigger_id}" if run.trigger_id else None
),
delivery_ref=(
f"dataflow-trigger-delivery:{run.trigger_delivery_id}"
if run.trigger_delivery_id
else None
),
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,
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,
)
def pipeline_run_descriptor(
session: Session,
run: DataflowRun,
*,
replayed: bool = False,
) -> DataflowRunDescriptor:
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}",
revision=revision.revision,
status=run.status,
definition_hash=run.definition_hash,
executor_version=run.executor_version,
input_row_count=run.input_row_count,
output_row_count=run.output_row_count,
output_publication_ref=run.output_publication_ref,
output_datasource_ref=run.output_datasource_ref,
output_materialization_ref=run.output_materialization_ref,
invocation_kind=run.invocation_kind,
trigger_ref=(
f"dataflow-trigger:{run.trigger_id}" if run.trigger_id else None
),
delivery_ref=(
f"dataflow-trigger-delivery:{run.trigger_delivery_id}"
if run.trigger_delivery_id
else None
),
error=run.error,
started_at=run.started_at,
finished_at=run.finished_at,
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),
"recovery": dict(recovery) if recovery is not None else None,
},
)
class SqlDataflowRunLifecycleProvider:
def __init__(self, *, registry: object | None = None) -> None:
self._registry = registry
def start_run(
self,
session: object,
principal: object,
*,
request: DataflowRunRequest,
) -> DataflowRunDescriptor:
db, api_principal = _run_context(session, principal)
run, replayed = start_pipeline_run(
db,
tenant_id=api_principal.tenant_id,
actor_id=_principal_actor_id(api_principal),
principal=api_principal,
registry=self._registry,
request=request,
defer_execution=True,
)
return pipeline_run_descriptor(db, run, replayed=replayed)
def get_run(
self,
session: object,
principal: object,
*,
run_ref: str,
) -> DataflowRunDescriptor | None:
db, api_principal = _run_context(session, principal)
try:
run = get_pipeline_run(
db,
tenant_id=api_principal.tenant_id,
run_ref=run_ref,
)
except DataflowNotFoundError:
return None
return pipeline_run_descriptor(db, run)
def cancel_run(
self,
session: object,
principal: object,
*,
run_ref: str,
) -> DataflowRunDescriptor:
db, api_principal = _run_context(session, principal)
run = cancel_pipeline_run(
db,
tenant_id=api_principal.tenant_id,
run_ref=run_ref,
)
return pipeline_run_descriptor(db, run)
def normalize_definition(
*,
graph: PipelineGraph,
sql_text: str | None,
editor_mode: str,
) -> NormalizedDefinition:
if editor_mode == "sql":
try:
compiled_graph, normalized_sql, diagnostics = compile_sql(
sql_text or "",
source_nodes=_source_nodes(graph, ()),
)
except SqlCompilationError as exc:
raise DataflowValidationError(exc.diagnostics) from exc
return NormalizedDefinition(
graph=compiled_graph,
sql_text=normalized_sql,
diagnostics=diagnostics,
)
diagnostics = validate_graph(graph)
errors = [item for item in diagnostics if item.severity == "error"]
if errors:
raise DataflowValidationError(diagnostics)
try:
rendered_sql, render_diagnostics = render_sql(graph)
diagnostics.extend(render_diagnostics)
except SqlCompilationError:
rendered_sql = None
return NormalizedDefinition(graph=graph, sql_text=rendered_sql, diagnostics=diagnostics)
def _source_nodes(
graph: PipelineGraph | None,
explicit_nodes: list[GraphNode] | tuple[()],
) -> list[GraphNode]:
nodes = list(explicit_nodes)
if graph is not None:
known = {node.id for node in nodes}
nodes.extend(
node
for node in graph.nodes
if node.type.startswith("source.") and node.id not in known
)
return nodes
def _datasource_source_resolver(
*,
session: Session,
principal: ApiPrincipal,
registry: object | None,
):
provider = datasource_catalogue(registry)
def resolve_source(node: GraphNode, limit: int) -> ResolvedSource:
if provider is None:
raise PipelineExecutionError(
"Datasource-backed execution requires the Datasources "
"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:
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(rows),
source_ref=resolved.datasource.ref,
provider=resolved.datasource.provider or "datasources",
fingerprint=resolved.datasource.fingerprint,
total_rows=resolved.total_rows,
truncated=len(rows) < resolved.total_rows,
)
return resolve_source
def _run_context(
session: object,
principal: object,
) -> tuple[Session, ApiPrincipal]:
if not isinstance(session, Session):
raise TypeError("Dataflow run providers require a SQLAlchemy session.")
if not isinstance(principal, ApiPrincipal):
raise DataflowConflictError("A tenant API principal is required.")
if not principal.tenant_id:
raise DataflowConflictError("A tenant API principal is required.")
return session, principal
def _principal_actor_id(principal: ApiPrincipal) -> str | None:
return principal.account_id or principal.membership_id or principal.identity_id
def _strip_ref(value: str, prefix: str) -> str | None:
cleaned = str(value or "").strip()
if not cleaned:
return None
if cleaned.startswith(prefix):
return cleaned[len(prefix) :]
return cleaned if ":" not in cleaned else None
def _run_request_payload(request: DataflowRunRequest) -> dict[str, object]:
target = request.publication
return {
"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,
"name": target.name,
"source_name": target.source_name,
"description": target.description,
"freeze": target.freeze,
"frozen_label": target.frozen_label,
"set_current": target.set_current,
"metadata": dict(target.metadata),
}
if target
else None
),
"invocation": _invocation_payload(request.invocation),
}
def _invocation_payload(
invocation: AutomationInvocation,
) -> dict[str, object]:
return {
"kind": invocation.kind,
"trigger_ref": invocation.trigger_ref,
"delivery_ref": invocation.delivery_ref,
"event_id": invocation.event_id,
"event_type": invocation.event_type,
"correlation_id": invocation.correlation_id,
"causation_id": invocation.causation_id,
"scheduled_for": (
invocation.scheduled_for.isoformat()
if invocation.scheduled_for
else None
),
"requested_by": invocation.requested_by,
"metadata": dict(invocation.metadata),
}
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),
sort_keys=True,
separators=(",", ":"),
default=str,
)
return hashlib.sha256(encoded.encode("utf-8")).hexdigest()
def _clean_optional(value: object | None) -> str | None:
if value is None:
return None
cleaned = str(value).strip()
return cleaned or None
def _ancestor_governance_limits(
provenance: Mapping[str, object],
) -> dict[str, bool]:
raw = provenance.get("source_effective_limits")
limits = raw if isinstance(raw, Mapping) else {}
return {
key: value if isinstance((value := limits.get(key)), bool) else True
for key in (
"inherit_to_lower_scopes",
"allow_run",
"allow_reuse",
"allow_automation",
)
}
def _effective_governance_limits(
pipeline: DataflowPipeline,
*,
decision_details: Mapping[str, object] | None = None,
) -> dict[str, bool]:
ancestor = _ancestor_governance_limits(
pipeline.derivation_provenance
)
effective = {
"inherit_to_lower_scopes": (
pipeline.inherit_to_lower_scopes
and ancestor["inherit_to_lower_scopes"]
),
"allow_run": pipeline.allow_run and ancestor["allow_run"],
"allow_reuse": pipeline.allow_reuse and ancestor["allow_reuse"],
"allow_automation": (
pipeline.allow_automation and ancestor["allow_automation"]
),
}
policy_limits = (
decision_details.get("effective_limits")
if decision_details is not None
else None
)
if isinstance(policy_limits, Mapping):
for key in effective:
value = policy_limits.get(key)
if isinstance(value, bool):
effective[key] = effective[key] and value
return effective
__all__ = [
"DataflowConflictError",
"DataflowError",
"DataflowNotFoundError",
"DataflowValidationError",
"SqlDataflowRunLifecycleProvider",
"cancel_pipeline_run",
"compile_sql_draft",
"create_pipeline",
"derive_pipeline",
"delete_pipeline",
"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",
"update_pipeline",
"validate_draft",
]