Refactor dataflow operators around runtime registry
This commit is contained in:
@@ -33,6 +33,7 @@ from govoplan_dataflow.backend.db.models import (
|
||||
from govoplan_dataflow.backend.executor import (
|
||||
EXECUTOR_VERSION,
|
||||
PipelineExecutionError,
|
||||
PipelineExecutionResult,
|
||||
ResolvedSource,
|
||||
execute_preview,
|
||||
)
|
||||
@@ -757,6 +758,55 @@ def start_pipeline_run(
|
||||
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")
|
||||
@@ -784,6 +834,10 @@ def start_pipeline_run(
|
||||
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(
|
||||
@@ -793,11 +847,21 @@ def start_pipeline_run(
|
||||
raise DataflowConflictError(
|
||||
"The bounded Dataflow runner supports between 1 and 500 output rows."
|
||||
)
|
||||
request_hash = _run_request_hash(request)
|
||||
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.pipeline_id == pipeline_id,
|
||||
DataflowRun.idempotency_key == idempotency_key,
|
||||
)
|
||||
)
|
||||
@@ -807,11 +871,20 @@ def start_pipeline_run(
|
||||
"The Dataflow run idempotency key was already used with "
|
||||
"different parameters."
|
||||
)
|
||||
return existing, True
|
||||
return existing
|
||||
|
||||
graph = PipelineGraph.model_validate(revision.graph)
|
||||
started_at = utcnow()
|
||||
run = DataflowRun(
|
||||
|
||||
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,
|
||||
@@ -838,15 +911,24 @@ def start_pipeline_run(
|
||||
diagnostics=[],
|
||||
input_row_count=0,
|
||||
output_row_count=0,
|
||||
started_at=started_at,
|
||||
started_at=utcnow(),
|
||||
created_by=actor_id,
|
||||
)
|
||||
session.add(run)
|
||||
session.flush()
|
||||
|
||||
|
||||
def _execute_pipeline_run(
|
||||
session: Session,
|
||||
*,
|
||||
run: DataflowRun,
|
||||
pipeline: DataflowPipeline,
|
||||
revision: DataflowPipelineRevision,
|
||||
request: DataflowRunRequest,
|
||||
principal: ApiPrincipal,
|
||||
registry: object | None,
|
||||
) -> None:
|
||||
try:
|
||||
result = execute_preview(
|
||||
graph,
|
||||
PipelineGraph.model_validate(revision.graph),
|
||||
row_limit=request.row_limit,
|
||||
source_resolver=_datasource_source_resolver(
|
||||
session=session,
|
||||
@@ -854,91 +936,128 @@ def start_pipeline_run(
|
||||
registry=registry,
|
||||
),
|
||||
)
|
||||
if request.publication and (
|
||||
result.truncated
|
||||
or any(
|
||||
bool(item.get("truncated"))
|
||||
for item in result.source_fingerprints
|
||||
)
|
||||
):
|
||||
raise PipelineExecutionError(
|
||||
"The bounded runner cannot publish a truncated result or a "
|
||||
"result calculated from truncated source data."
|
||||
)
|
||||
|
||||
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
|
||||
|
||||
_apply_pipeline_result(run, result)
|
||||
if request.publication:
|
||||
publisher = datasource_publication(registry)
|
||||
if publisher is None:
|
||||
raise PipelineExecutionError(
|
||||
"Publishing Dataflow output requires the Datasources "
|
||||
"publication capability."
|
||||
)
|
||||
target = request.publication
|
||||
publication = publisher.publish_rows(
|
||||
_ensure_publishable(result)
|
||||
_publish_pipeline_result(
|
||||
session,
|
||||
principal,
|
||||
request=DatasourcePublicationRequest(
|
||||
producer_module="dataflow",
|
||||
producer_run_ref=f"dataflow-run:{run.id}",
|
||||
idempotency_key=f"{pipeline.id}:{idempotency_key}",
|
||||
rows=tuple(dict(row) for row in result.rows),
|
||||
target_datasource_ref=target.target_datasource_ref,
|
||||
name=target.name or f"{pipeline.name} output",
|
||||
source_name=target.source_name,
|
||||
description=target.description,
|
||||
freeze=target.freeze,
|
||||
frozen_label=target.frozen_label,
|
||||
set_current=target.set_current,
|
||||
provenance={
|
||||
"pipeline_ref": f"pipeline:{pipeline.id}",
|
||||
"pipeline_revision": revision.revision,
|
||||
"definition_hash": revision.content_hash,
|
||||
"source_fingerprints": result.source_fingerprints,
|
||||
},
|
||||
metadata={
|
||||
**dict(target.metadata),
|
||||
"dataflow_run_ref": f"dataflow-run:{run.id}",
|
||||
},
|
||||
),
|
||||
run=run,
|
||||
pipeline=pipeline,
|
||||
revision=revision,
|
||||
request=request,
|
||||
result=result,
|
||||
principal=principal,
|
||||
registry=registry,
|
||||
)
|
||||
run.output_publication_ref = publication.ref
|
||||
run.output_datasource_ref = publication.datasource.ref
|
||||
run.output_materialization_ref = publication.materialization.ref
|
||||
run.status = "succeeded"
|
||||
run.finished_at = utcnow()
|
||||
run.error = None
|
||||
except (DatasourceError, PipelineExecutionError) as exc:
|
||||
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),
|
||||
)
|
||||
_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."
|
||||
)
|
||||
run.diagnostics = [
|
||||
item.model_dump(mode="json") for item in diagnostics
|
||||
]
|
||||
run.source_fingerprints = list(
|
||||
getattr(exc, "source_fingerprints", ())
|
||||
|
||||
|
||||
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."
|
||||
)
|
||||
run.input_row_count = int(getattr(exc, "input_row_count", 0))
|
||||
session.flush()
|
||||
return run, False
|
||||
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(
|
||||
|
||||
Reference in New Issue
Block a user