Add intermediate previews and graph reconnection

src/govoplan_dataflow/backend/executor.py

@@ -13,6 +13,7 @@ from govoplan_dataflow.backend.schemas import (
     DataflowDiagnostic,
     GraphNode,
     NodePreviewDiagnostic,
+    NodePreviewResult,
     PipelineGraph,
     PreviewColumn,
 )
@@ -35,6 +36,7 @@ class PipelineExecutionError(RuntimeError):
         source_fingerprints: tuple[dict[str, Any], ...] = (),
         input_row_count: int = 0,
         diagnostics: tuple[DataflowDiagnostic, ...] = (),
+        node_preview: NodePreviewResult | None = None,
     ) -> None:
         super().__init__(message)
         self.node_id = node_id
@@ -42,6 +44,7 @@ class PipelineExecutionError(RuntimeError):
         self.source_fingerprints = source_fingerprints
         self.input_row_count = input_row_count
         self.diagnostics = diagnostics
+        self.node_preview = node_preview
 
 
 @dataclass(frozen=True)
@@ -52,6 +55,7 @@ class PipelineExecutionResult:
     columns: list[PreviewColumn]
     diagnostics: list[DataflowDiagnostic]
     node_diagnostics: list[NodePreviewDiagnostic]
+    node_preview: NodePreviewResult | None
     source_fingerprints: list[dict[str, Any]]
     input_row_count: int
 
@@ -74,6 +78,7 @@ def execute_preview(
     *,
     row_limit: int,
     source_resolver: SourceResolver | None = None,
+    preview_node_id: str | None = None,
 ) -> PipelineExecutionResult:
     validation = validate_graph(graph)
     errors = [item for item in validation if item.severity == "error"]
@@ -81,6 +86,11 @@ def execute_preview(
         raise PipelineExecutionError(errors[0].message, node_id=errors[0].node_id)
 
     node_by_id = {node.id: node for node in graph.nodes}
+    if preview_node_id is not None and preview_node_id not in node_by_id:
+        raise PipelineExecutionError(
+            "The requested preview node is not part of this pipeline.",
+            node_id=preview_node_id,
+        )
     inputs = graph_inputs_by_port(graph)
     ordered, cyclic = topological_order(graph)
     if cyclic:
@@ -120,6 +130,7 @@ def execute_preview(
                 source_fingerprints=tuple(source_fingerprints),
                 input_row_count=input_row_count,
                 diagnostics=tuple(item for item in validation if item.severity != "error"),
+                node_preview=_node_preview(outputs, preview_node_id, row_limit),
             )
         try:
             if node.type == "source.inline":
@@ -230,6 +241,7 @@ def execute_preview(
                 source_fingerprints=tuple(source_fingerprints),
                 input_row_count=input_row_count,
                 diagnostics=tuple(item for item in validation if item.severity != "error"),
+                node_preview=_node_preview(outputs, preview_node_id, row_limit),
             ) from exc
 
         outputs[node.id] = output_rows
@@ -255,11 +267,30 @@ def execute_preview(
         columns=infer_columns(all_rows),
         diagnostics=[item for item in validation if item.severity != "error"],
         node_diagnostics=node_diagnostics,
+        node_preview=_node_preview(outputs, preview_node_id, row_limit),
         source_fingerprints=source_fingerprints,
         input_row_count=input_row_count,
     )
 
 
+def _node_preview(
+    outputs: dict[str, list[dict[str, Any]]],
+    node_id: str | None,
+    row_limit: int,
+) -> NodePreviewResult | None:
+    if node_id is None or node_id not in outputs:
+        return None
+    all_rows = outputs[node_id]
+    rows = [dict(row) for row in all_rows[:row_limit]]
+    return NodePreviewResult(
+        node_id=node_id,
+        columns=infer_columns(all_rows),
+        rows=rows,
+        total_rows=len(all_rows),
+        truncated=len(rows) < len(all_rows),
+    )
+
+
 def infer_columns(rows: list[dict[str, Any]]) -> list[PreviewColumn]:
     names: list[str] = []
     for row in rows: