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

1481 lines
46 KiB
Python

from __future__ import annotations
from collections.abc import Mapping
import hashlib
import json
from dataclasses import dataclass
from sqlalchemy import 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 (
DataflowRunConflictError,
DataflowRunDescriptor,
DataflowRunError,
DataflowRunNotFoundError,
DataflowRunRequest,
)
from govoplan_core.core.datasources import (
DatasourceError,
DatasourcePublicationRequest,
DatasourceReadRequest,
datasource_catalogue,
datasource_publication,
)
from govoplan_core.db.base import utcnow
from govoplan_dataflow.backend.db.models import (
DataflowPipeline,
DataflowPipelineRevision,
DataflowRun,
)
from govoplan_dataflow.backend.executor import (
EXECUTOR_VERSION,
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,
PipelineCreateRequest,
PipelineDeriveRequest,
PipelineDraftRequest,
PipelineGraph,
PipelinePreviewRequest,
PipelinePreviewResponse,
PipelineResponse,
PipelineRevisionResponse,
PipelineRunResponse,
PipelineSqlResponse,
PipelineUpdateRequest,
PipelineValidationResponse,
)
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
@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:
provider = datasource_catalogue(registry)
def resolve_source(node: GraphNode, limit: int) -> ResolvedSource:
if provider is None:
raise PipelineExecutionError(
"Datasource-backed preview requires the Datasources catalogue capability.",
node_id=node.id,
)
if principal is None:
raise PipelineExecutionError(
"Datasource-backed preview requires a tenant API principal.",
node_id=node.id,
)
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")
),
),
)
except DatasourceError as exc:
raise PipelineExecutionError(str(exc), node_id=node.id) from exc
return ResolvedSource(
rows=tuple(dict(row) for row in resolved.rows),
source_ref=resolved.datasource.ref,
provider=resolved.datasource.provider or "datasources",
fingerprint=resolved.datasource.fingerprint,
total_rows=resolved.total_rows,
truncated=resolved.truncated,
)
result = execute_preview(
graph,
row_limit=payload.row_limit,
source_resolver=resolve_source,
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:
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 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 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,
) -> 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
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,
)
session.add(run)
session.flush()
_execute_pipeline_run(
session,
run=run,
pipeline=pipeline,
revision=revision,
request=request,
principal=principal,
registry=registry,
)
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,
)
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 > 500:
raise DataflowConflictError(
"The bounded Dataflow runner supports between 1 and 500 output rows."
)
return idempotency_key, _run_request_hash(request)
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,
) -> DataflowRun:
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",
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,
started_at=utcnow(),
created_by=actor_id,
)
def _execute_pipeline_run(
session: Session,
*,
run: DataflowRun,
pipeline: DataflowPipeline,
revision: DataflowPipelineRevision,
request: DataflowRunRequest,
principal: ApiPrincipal,
registry: object | None,
) -> None:
try:
result = execute_preview(
PipelineGraph.model_validate(revision.graph),
row_limit=request.row_limit,
source_resolver=_datasource_source_resolver(
session=session,
principal=principal,
registry=registry,
),
)
_apply_pipeline_result(run, result)
if request.publication:
_ensure_publishable(result)
_publish_pipeline_result(
session,
run=run,
pipeline=pipeline,
revision=revision,
request=request,
result=result,
principal=principal,
registry=registry,
)
run.status = "succeeded"
run.finished_at = utcnow()
run.error = None
except (DatasourceError, PipelineExecutionError) as exc:
_mark_pipeline_run_failed(run, exc)
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,
) -> 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
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}",
},
),
)
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)
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 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", "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."
session.flush()
return run
def pipeline_run_response(
session: Session,
run: DataflowRun,
*,
replayed: bool = False,
) -> PipelineRunResponse:
revision = session.get(DataflowPipelineRevision, run.pipeline_revision_id)
if revision is None:
raise DataflowNotFoundError("Dataflow pipeline revision not found")
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,
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,
error=run.error,
started_at=run.started_at,
finished_at=run.finished_at,
created_by=run.created_by,
created_at=run.created_at,
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")
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,
"source_fingerprints": list(run.source_fingerprints),
"diagnostics": list(run.diagnostics),
},
)
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,
)
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,
)
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")
),
),
)
except DatasourceError as exc:
raise PipelineExecutionError(str(exc), node_id=node.id) from exc
return ResolvedSource(
rows=tuple(dict(row) for row in resolved.rows),
source_ref=resolved.datasource.ref,
provider=resolved.datasource.provider or "datasources",
fingerprint=resolved.datasource.fingerprint,
total_rows=resolved.total_rows,
truncated=resolved.truncated,
)
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,
"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 _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: str | None) -> str | None:
if value is None:
return None
cleaned = 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_runs",
"list_pipelines",
"normalize_definition",
"pipeline_response",
"pipeline_run_descriptor",
"pipeline_run_response",
"preview_pipeline",
"render_graph_sql",
"start_pipeline_run",
"update_pipeline",
"validate_draft",
]