Expand governed Dataflow editor and node library

src/govoplan_dataflow/backend/executor.py

@@ -6,26 +6,42 @@ import time
 from collections import defaultdict
 from dataclasses import dataclass
 from decimal import Decimal
-from typing import Any
+from typing import Any, Callable
 
-from govoplan_dataflow.backend.graph import graph_input_map, topological_order, validate_graph
+from govoplan_dataflow.backend.graph import graph_inputs_by_port, topological_order, validate_graph
 from govoplan_dataflow.backend.schemas import (
     DataflowDiagnostic,
+    GraphNode,
     NodePreviewDiagnostic,
     PipelineGraph,
     PreviewColumn,
 )
 
 
-EXECUTOR_VERSION = "dataflow-preview-v1"
+EXECUTOR_VERSION = "dataflow-preview-v2"
 MAX_EXECUTION_SECONDS = 2.0
 MAX_RESULT_BYTES = 1_000_000
+MAX_SOURCE_ROWS = 250
+MAX_INTERMEDIATE_ROWS = 10_000
 
 
 class PipelineExecutionError(RuntimeError):
-    def __init__(self, message: str, *, node_id: str | None = None) -> None:
+    def __init__(
+        self,
+        message: str,
+        *,
+        node_id: str | None = None,
+        node_diagnostics: tuple[NodePreviewDiagnostic, ...] = (),
+        source_fingerprints: tuple[dict[str, Any], ...] = (),
+        input_row_count: int = 0,
+        diagnostics: tuple[DataflowDiagnostic, ...] = (),
+    ) -> None:
         super().__init__(message)
         self.node_id = node_id
+        self.node_diagnostics = node_diagnostics
+        self.source_fingerprints = source_fingerprints
+        self.input_row_count = input_row_count
+        self.diagnostics = diagnostics
 
 
 @dataclass(frozen=True)
@@ -40,14 +56,32 @@ class PipelineExecutionResult:
     input_row_count: int
 
 
-def execute_preview(graph: PipelineGraph, *, row_limit: int) -> PipelineExecutionResult:
+@dataclass(frozen=True)
+class ResolvedSource:
+    rows: tuple[dict[str, Any], ...]
+    source_ref: str
+    provider: str
+    fingerprint: str
+    total_rows: int
+    truncated: bool = False
+
+
+SourceResolver = Callable[[GraphNode, int], ResolvedSource]
+
+
+def execute_preview(
+    graph: PipelineGraph,
+    *,
+    row_limit: int,
+    source_resolver: SourceResolver | None = None,
+) -> PipelineExecutionResult:
     validation = validate_graph(graph)
     errors = [item for item in validation if item.severity == "error"]
     if errors:
         raise PipelineExecutionError(errors[0].message, node_id=errors[0].node_id)
 
     node_by_id = {node.id: node for node in graph.nodes}
-    inputs = graph_input_map(graph)
+    inputs = graph_inputs_by_port(graph)
     ordered, cyclic = topological_order(graph)
     if cyclic:
         raise PipelineExecutionError("Pipeline graph contains a cycle")
@@ -59,11 +93,34 @@ def execute_preview(graph: PipelineGraph, *, row_limit: int) -> PipelineExecutio
     input_row_count = 0
 
     for node_id in ordered:
-        if time.monotonic() - started > MAX_EXECUTION_SECONDS:
-            raise PipelineExecutionError("Preview exceeded the two-second execution limit", node_id=node_id)
         node = node_by_id[node_id]
         node_started = time.monotonic()
-        input_rows = [] if node.type.startswith("source.") else outputs[inputs[node_id]]
+        node_inputs = inputs.get(node_id, {})
+        input_sets = [
+            outputs[source_id]
+            for port_sources in node_inputs.values()
+            for source_id in port_sources
+        ]
+        input_rows = input_sets[0] if len(input_sets) == 1 else []
+        node_messages: list[str] = []
+        if time.monotonic() - started > MAX_EXECUTION_SECONDS:
+            failed = NodePreviewDiagnostic(
+                node_id=node.id,
+                status="failed",
+                input_rows=sum(len(rows) for rows in input_sets),
+                output_rows=0,
+                duration_ms=0,
+                columns=[],
+                messages=["Preview exceeded the two-second execution limit"],
+            )
+            raise PipelineExecutionError(
+                "Preview exceeded the two-second execution limit",
+                node_id=node.id,
+                node_diagnostics=tuple([*node_diagnostics, failed]),
+                source_fingerprints=tuple(source_fingerprints),
+                input_row_count=input_row_count,
+                diagnostics=tuple(item for item in validation if item.severity != "error"),
+            )
         try:
             if node.type == "source.inline":
                 output_rows = [dict(row) for row in node.config.get("rows", [])]
@@ -78,14 +135,64 @@ def execute_preview(graph: PipelineGraph, *, row_limit: int) -> PipelineExecutio
                     }
                 )
             elif node.type == "source.reference":
-                raise PipelineExecutionError(
-                    "This connector-backed source is not available to the local preview executor.",
+                if source_resolver is None:
+                    raise PipelineExecutionError(
+                        "Connector-backed preview requires the Connectors tabular-source capability.",
+                        node_id=node.id,
+                    )
+                resolved = source_resolver(node, MAX_SOURCE_ROWS)
+                output_rows = [dict(row) for row in resolved.rows]
+                input_row_count += len(output_rows)
+                source_fingerprints.append(
+                    {
+                        "node_id": node.id,
+                        "source_ref": resolved.source_ref,
+                        "source_name": node.config.get("source_name"),
+                        "kind": "connector",
+                        "provider": resolved.provider,
+                        "fingerprint": resolved.fingerprint,
+                        "row_count": resolved.total_rows,
+                        "preview_rows": len(output_rows),
+                        "truncated": resolved.truncated,
+                    }
+                )
+                if resolved.truncated:
+                    message = (
+                        f"Source preview used {len(output_rows):,} of "
+                        f"{resolved.total_rows:,} rows."
+                    )
+                    node_messages.append(message)
+                    validation.append(
+                        DataflowDiagnostic(
+                            severity="warning",
+                            code="source.preview_truncated",
+                            message=message,
+                            node_id=node.id,
+                        )
+                    )
+            elif node.type == "combine.union":
+                ordered_inputs = [
+                    outputs[source_id]
+                    for source_id in node_inputs.get("input", [])
+                ]
+                output_rows = _union_rows(ordered_inputs, node.config)
+            elif node.type == "combine.join":
+                left_rows = outputs[node_inputs["left"][0]]
+                right_rows = outputs[node_inputs["right"][0]]
+                output_rows = _join_rows(
+                    left_rows,
+                    right_rows,
+                    node.config,
                     node_id=node.id,
                 )
             elif node.type == "filter":
                 output_rows = _filter_rows(input_rows, node.config, node_id=node.id)
+            elif node.type == "distinct":
+                output_rows = _distinct_rows(input_rows, node.config)
             elif node.type == "select":
                 output_rows = _select_rows(input_rows, node.config)
+            elif node.type == "derive":
+                output_rows = _derive_rows(input_rows, node.config, node_id=node.id)
             elif node.type == "aggregate":
                 output_rows = _aggregate_rows(input_rows, node.config, node_id=node.id)
             elif node.type == "sort":
@@ -96,25 +203,45 @@ def execute_preview(graph: PipelineGraph, *, row_limit: int) -> PipelineExecutio
                 output_rows = [dict(row) for row in input_rows]
             else:
                 raise PipelineExecutionError(f"Unsupported node type: {node.type}", node_id=node.id)
-        except PipelineExecutionError:
-            raise
-        except (ArithmeticError, TypeError, ValueError) as exc:
-            raise PipelineExecutionError(str(exc), node_id=node.id) from exc
-
-        if len(json.dumps(output_rows, default=str).encode("utf-8")) > MAX_RESULT_BYTES:
-            raise PipelineExecutionError(
-                "A preview node exceeded the one-megabyte result limit.",
-                node_id=node.id,
+            if len(json.dumps(output_rows, default=str).encode("utf-8")) > MAX_RESULT_BYTES:
+                raise PipelineExecutionError(
+                    "A preview node exceeded the one-megabyte result limit.",
+                    node_id=node.id,
+                )
+        except (PipelineExecutionError, ArithmeticError, KeyError, TypeError, ValueError) as exc:
+            execution_error = (
+                exc
+                if isinstance(exc, PipelineExecutionError)
+                else PipelineExecutionError(str(exc), node_id=node.id)
             )
+            failed = NodePreviewDiagnostic(
+                node_id=execution_error.node_id or node.id,
+                status="failed",
+                input_rows=sum(len(rows) for rows in input_sets),
+                output_rows=0,
+                duration_ms=round((time.monotonic() - node_started) * 1000, 3),
+                columns=[],
+                messages=[str(execution_error)],
+            )
+            raise PipelineExecutionError(
+                str(execution_error),
+                node_id=execution_error.node_id or node.id,
+                node_diagnostics=tuple([*node_diagnostics, failed]),
+                source_fingerprints=tuple(source_fingerprints),
+                input_row_count=input_row_count,
+                diagnostics=tuple(item for item in validation if item.severity != "error"),
+            ) from exc
+
         outputs[node.id] = output_rows
         node_diagnostics.append(
             NodePreviewDiagnostic(
                 node_id=node.id,
                 status="succeeded",
-                input_rows=len(input_rows),
+                input_rows=sum(len(rows) for rows in input_sets),
                 output_rows=len(output_rows),
                 duration_ms=round((time.monotonic() - node_started) * 1000, 3),
                 columns=infer_columns(output_rows),
+                messages=node_messages,
             )
         )
 
@@ -126,7 +253,7 @@ def execute_preview(graph: PipelineGraph, *, row_limit: int) -> PipelineExecutio
         total_rows=len(all_rows),
         truncated=len(rows) < len(all_rows),
         columns=infer_columns(all_rows),
-        diagnostics=[],
+        diagnostics=[item for item in validation if item.severity != "error"],
         node_diagnostics=node_diagnostics,
         source_fingerprints=source_fingerprints,
         input_row_count=input_row_count,
@@ -156,6 +283,202 @@ def infer_columns(rows: list[dict[str, Any]]) -> list[PreviewColumn]:
     return columns
 
 
+def _union_rows(
+    inputs: list[list[dict[str, Any]]],
+    config: dict[str, Any],
+) -> list[dict[str, Any]]:
+    rows = [dict(row) for input_rows in inputs for row in input_rows]
+    if config.get("mode", "all") == "distinct":
+        return _distinct_rows(rows, {})
+    return rows
+
+
+def _join_rows(
+    left_rows: list[dict[str, Any]],
+    right_rows: list[dict[str, Any]],
+    config: dict[str, Any],
+    *,
+    node_id: str,
+) -> list[dict[str, Any]]:
+    left_keys = [str(item) for item in config["left_keys"]]
+    right_keys = [str(item) for item in config["right_keys"]]
+    join_type = str(config.get("join_type", "inner"))
+    right_prefix = str(config.get("right_prefix", "right_"))
+    left_columns = _ordered_columns(left_rows)
+    right_columns = _ordered_columns(right_rows)
+    right_index: dict[tuple[Any, ...], list[tuple[int, dict[str, Any]]]] = defaultdict(list)
+    for index, row in enumerate(right_rows):
+        key = _join_key(row, right_keys)
+        if key is not None:
+            right_index[key].append((index, row))
+
+    output: list[dict[str, Any]] = []
+    matched_right: set[int] = set()
+    for left_row in left_rows:
+        key = _join_key(left_row, left_keys)
+        matches = right_index.get(key, ()) if key is not None else ()
+        if matches:
+            for right_index_value, right_row in matches:
+                matched_right.add(right_index_value)
+                output.append(
+                    _merge_join_rows(
+                        left_row,
+                        right_row,
+                        left_columns=left_columns,
+                        right_columns=right_columns,
+                        right_prefix=right_prefix,
+                    )
+                )
+                _guard_intermediate_size(output, node_id=node_id)
+        elif join_type in {"left", "full"}:
+            output.append(
+                _merge_join_rows(
+                    left_row,
+                    None,
+                    left_columns=left_columns,
+                    right_columns=right_columns,
+                    right_prefix=right_prefix,
+                )
+            )
+
+    if join_type in {"right", "full"}:
+        for index, right_row in enumerate(right_rows):
+            if index in matched_right:
+                continue
+            output.append(
+                _merge_join_rows(
+                    None,
+                    right_row,
+                    left_columns=left_columns,
+                    right_columns=right_columns,
+                    right_prefix=right_prefix,
+                )
+            )
+            _guard_intermediate_size(output, node_id=node_id)
+    return output
+
+
+def _distinct_rows(
+    rows: list[dict[str, Any]],
+    config: dict[str, Any],
+) -> list[dict[str, Any]]:
+    columns = [str(item) for item in config.get("columns", [])]
+    seen: set[str] = set()
+    output: list[dict[str, Any]] = []
+    for row in rows:
+        value = {column: row.get(column) for column in columns} if columns else row
+        identity = json.dumps(value, sort_keys=True, separators=(",", ":"), default=str)
+        if identity in seen:
+            continue
+        seen.add(identity)
+        output.append(dict(row))
+    return output
+
+
+def _derive_rows(
+    rows: list[dict[str, Any]],
+    config: dict[str, Any],
+    *,
+    node_id: str,
+) -> list[dict[str, Any]]:
+    target = str(config["target_column"])
+    operation = str(config["operation"])
+    columns = [str(item) for item in config["source_columns"]]
+    separator = str(config.get("separator", " "))
+    output: list[dict[str, Any]] = []
+    for row in rows:
+        values = [row.get(column) for column in columns]
+        try:
+            derived = _derive_value(operation, values, separator=separator)
+        except (ArithmeticError, TypeError, ValueError) as exc:
+            raise PipelineExecutionError(
+                f"Cannot derive {target!r} with {operation!r}: {exc}",
+                node_id=node_id,
+            ) from exc
+        result = dict(row)
+        result[target] = derived
+        output.append(result)
+    return output
+
+
+def _derive_value(operation: str, values: list[Any], *, separator: str) -> Any:
+    if operation == "copy":
+        return values[0]
+    if operation == "upper":
+        return None if values[0] is None else str(values[0]).upper()
+    if operation == "lower":
+        return None if values[0] is None else str(values[0]).lower()
+    if operation == "trim":
+        return None if values[0] is None else str(values[0]).strip()
+    if operation == "concat":
+        return separator.join(str(value) for value in values if value is not None)
+    if operation == "coalesce":
+        return next((value for value in values if value not in (None, "")), None)
+    if any(value is None for value in values):
+        return None
+    left, right = values
+    if isinstance(left, bool) or isinstance(right, bool):
+        raise TypeError("boolean values are not numeric inputs")
+    if operation == "add":
+        return left + right
+    if operation == "subtract":
+        return left - right
+    if operation == "multiply":
+        return left * right
+    if operation == "divide":
+        return left / right
+    raise ValueError(f"unknown derive operation {operation!r}")
+
+
+def _join_key(row: dict[str, Any], columns: list[str]) -> tuple[Any, ...] | None:
+    values = tuple(row.get(column) for column in columns)
+    if any(value is None for value in values):
+        return None
+    try:
+        hash(values)
+    except TypeError:
+        return tuple(
+            json.dumps(value, sort_keys=True, separators=(",", ":"), default=str)
+            for value in values
+        )
+    return values
+
+
+def _ordered_columns(rows: list[dict[str, Any]]) -> list[str]:
+    columns: list[str] = []
+    for row in rows:
+        for column in row:
+            if column not in columns:
+                columns.append(column)
+    return columns
+
+
+def _merge_join_rows(
+    left_row: dict[str, Any] | None,
+    right_row: dict[str, Any] | None,
+    *,
+    left_columns: list[str],
+    right_columns: list[str],
+    right_prefix: str,
+) -> dict[str, Any]:
+    output = {
+        column: left_row.get(column) if left_row is not None else None
+        for column in left_columns
+    }
+    for column in right_columns:
+        output_column = f"{right_prefix}{column}"
+        output[output_column] = right_row.get(column) if right_row is not None else None
+    return output
+
+
+def _guard_intermediate_size(rows: list[dict[str, Any]], *, node_id: str) -> None:
+    if len(rows) > MAX_INTERMEDIATE_ROWS:
+        raise PipelineExecutionError(
+            f"Preview join exceeded the {MAX_INTERMEDIATE_ROWS:,}-row intermediate limit.",
+            node_id=node_id,
+        )
+
+
 def _filter_rows(
     rows: list[dict[str, Any]],
     config: dict[str, Any],
@@ -305,8 +628,11 @@ def _type_name(value: Any) -> str:
 
 __all__ = [
     "EXECUTOR_VERSION",
+    "MAX_SOURCE_ROWS",
     "PipelineExecutionError",
     "PipelineExecutionResult",
+    "ResolvedSource",
+    "SourceResolver",
     "execute_preview",
     "infer_columns",
 ]