from __future__ import annotations from dataclasses import dataclass from sqlalchemy import select from sqlalchemy.orm import Session 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, execute_preview, ) from govoplan_dataflow.backend.graph import canonical_graph_payload, definition_hash, validate_graph from govoplan_dataflow.backend.schemas import ( DataflowDiagnostic, GraphNode, PipelineCreateRequest, PipelineDraftRequest, PipelineGraph, PipelinePreviewRequest, PipelinePreviewResponse, PipelineResponse, PipelineRevisionResponse, PipelineSqlResponse, PipelineUpdateRequest, PipelineValidationResponse, ) from govoplan_dataflow.backend.sql_compiler import ( SqlCompilationError, compile_sql, render_sql, ) class DataflowError(RuntimeError): pass class DataflowNotFoundError(DataflowError): pass class DataflowConflictError(DataflowError): 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( DataflowPipeline.tenant_id == tenant_id, 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, DataflowPipeline.tenant_id == tenant_id, 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) pipeline = DataflowPipeline( tenant_id=tenant_id, 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=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}" ) 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 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=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 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) -> 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), ) 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), ) 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, ) -> 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) 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=[], 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 = execute_preview(graph, row_limit=payload.row_limit) 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 source_fingerprints = result.source_fingerprints input_row_count = result.input_row_count except PipelineExecutionError as exc: status = "failed" error = str(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 = [] source_fingerprints = [] input_row_count = 0 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, definition_hash=graph_hash, executor_version=EXECUTOR_VERSION, ) 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 _clean_optional(value: str | None) -> str | None: if value is None: return None cleaned = value.strip() return cleaned or None __all__ = [ "DataflowConflictError", "DataflowError", "DataflowNotFoundError", "DataflowValidationError", "compile_sql_draft", "create_pipeline", "delete_pipeline", "get_pipeline", "get_pipeline_revision", "list_pipelines", "normalize_definition", "pipeline_response", "preview_pipeline", "render_graph_sql", "update_pipeline", "validate_draft", ]