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

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",
]