feat: run pinned pipelines and publish outputs
This commit is contained in:
@@ -1,15 +1,26 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from govoplan_core.auth import ApiPrincipal
|
||||
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 (
|
||||
@@ -34,6 +45,7 @@ from govoplan_dataflow.backend.schemas import (
|
||||
PipelinePreviewResponse,
|
||||
PipelineResponse,
|
||||
PipelineRevisionResponse,
|
||||
PipelineRunResponse,
|
||||
PipelineSqlResponse,
|
||||
PipelineUpdateRequest,
|
||||
PipelineValidationResponse,
|
||||
@@ -45,15 +57,15 @@ from govoplan_dataflow.backend.sql_compiler import (
|
||||
)
|
||||
|
||||
|
||||
class DataflowError(RuntimeError):
|
||||
class DataflowError(DataflowRunError):
|
||||
pass
|
||||
|
||||
|
||||
class DataflowNotFoundError(DataflowError):
|
||||
class DataflowNotFoundError(DataflowError, DataflowRunNotFoundError):
|
||||
pass
|
||||
|
||||
|
||||
class DataflowConflictError(DataflowError):
|
||||
class DataflowConflictError(DataflowError, DataflowRunConflictError):
|
||||
pass
|
||||
|
||||
|
||||
@@ -495,6 +507,363 @@ def preview_pipeline(
|
||||
)
|
||||
|
||||
|
||||
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_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,
|
||||
)
|
||||
revision = get_pipeline_revision(
|
||||
session,
|
||||
pipeline=pipeline,
|
||||
revision=request.revision,
|
||||
)
|
||||
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."
|
||||
)
|
||||
request_hash = _run_request_hash(request)
|
||||
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, True
|
||||
|
||||
graph = PipelineGraph.model_validate(revision.graph)
|
||||
started_at = utcnow()
|
||||
run = 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),
|
||||
source_fingerprints=[],
|
||||
result_schema=[],
|
||||
diagnostics=[],
|
||||
input_row_count=0,
|
||||
output_row_count=0,
|
||||
started_at=started_at,
|
||||
created_by=actor_id,
|
||||
)
|
||||
session.add(run)
|
||||
session.flush()
|
||||
|
||||
try:
|
||||
result = execute_preview(
|
||||
graph,
|
||||
row_limit=request.row_limit,
|
||||
source_resolver=_datasource_source_resolver(
|
||||
session=session,
|
||||
principal=principal,
|
||||
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
|
||||
|
||||
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(
|
||||
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.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),
|
||||
)
|
||||
)
|
||||
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))
|
||||
session.flush()
|
||||
return run, False
|
||||
|
||||
|
||||
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,
|
||||
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,
|
||||
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,
|
||||
@@ -541,6 +910,109 @@ def _source_nodes(
|
||||
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
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
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
|
||||
@@ -553,16 +1025,23 @@ __all__ = [
|
||||
"DataflowError",
|
||||
"DataflowNotFoundError",
|
||||
"DataflowValidationError",
|
||||
"SqlDataflowRunLifecycleProvider",
|
||||
"cancel_pipeline_run",
|
||||
"compile_sql_draft",
|
||||
"create_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",
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user