from __future__ import annotations import unittest from govoplan_dataflow.backend.executor import ( MAX_SOURCE_ROWS, PipelineExecutionError, ResolvedSource, execute_preview, ) from govoplan_dataflow.backend.graph import validate_graph from govoplan_dataflow.backend.schemas import ( GraphEdge, GraphNode, GraphPosition, PipelineGraph, ) from govoplan_dataflow.backend.sql_compiler import SqlCompilationError, compile_sql, render_sql def inline_source() -> GraphNode: return GraphNode( id="source", type="source.inline", label="Monthly input", position=GraphPosition(x=40, y=160), config={ "source_name": "monthly_files", "rows": [ {"department": "A", "status": "open", "amount": 10}, {"department": "A", "status": "closed", "amount": 5}, {"department": "B", "status": "open", "amount": 20}, ], }, ) def lookup_source() -> GraphNode: return GraphNode( id="lookup", type="source.inline", label="Department lookup", position=GraphPosition(x=40, y=300), config={ "source_name": "department_lookup", "rows": [ {"department": "A", "label": "Administration"}, {"department": "B", "label": "Building services"}, ], }, ) class DataflowGraphAndSqlTests(unittest.TestCase): def test_rejects_invalid_or_duplicate_logical_source_names(self) -> None: invalid_source = inline_source().model_copy(deep=True) invalid_source.config["source_name"] = "monthly cases" duplicate_source = lookup_source().model_copy(deep=True) duplicate_source.config["source_name"] = "MONTHLY_FILES" graph = PipelineGraph( nodes=[ invalid_source, duplicate_source, GraphNode( id="union", type="combine.union", label="Append rows", position=GraphPosition(x=250, y=160), config={"mode": "all"}, ), GraphNode( id="output", type="output", label="Output", position=GraphPosition(x=450, y=160), config={}, ), ], edges=[ GraphEdge(id="edge-1", source="source", target="union"), GraphEdge(id="edge-2", source="lookup", target="union"), GraphEdge(id="edge-3", source="union", target="output"), ], ) diagnostics = validate_graph(graph) self.assertTrue(any(item.code == "source.name_invalid" for item in diagnostics)) invalid_source.config["source_name"] = "monthly_files" diagnostics = validate_graph( graph.model_copy( update={"nodes": [invalid_source, *graph.nodes[1:]]}, deep=True, ) ) self.assertTrue(any(item.code == "source.duplicate_name" for item in diagnostics)) def test_compiles_renders_and_executes_grouped_query(self) -> None: sql = """ SELECT department, COUNT(*) AS records, SUM(amount) AS total FROM monthly_files WHERE status = 'open' GROUP BY department ORDER BY total DESC LIMIT 50 """ graph, normalized, diagnostics = compile_sql(sql, source_nodes=[inline_source()]) self.assertEqual([], diagnostics) self.assertEqual( ["source.inline", "filter", "aggregate", "sort", "limit", "output"], [node.type for node in graph.nodes], ) self.assertEqual(normalized, render_sql(graph)[0]) result = execute_preview(graph, row_limit=100) self.assertEqual( [ {"department": "B", "records": 1, "total": 20}, {"department": "A", "records": 1, "total": 10}, ], result.rows, ) self.assertEqual(2, result.total_rows) self.assertFalse(result.truncated) self.assertEqual(3, result.input_row_count) self.assertEqual(6, len(result.node_diagnostics)) self.assertEqual(1, len(result.source_fingerprints)) def test_rejects_effectful_or_multi_statement_sql(self) -> None: for sql in ( "DELETE FROM monthly_files", "SELECT * FROM monthly_files; DROP TABLE monthly_files", "INSERT INTO target SELECT * FROM monthly_files", ): with self.subTest(sql=sql), self.assertRaises(SqlCompilationError): compile_sql(sql, source_nodes=[inline_source()]) def test_compiles_renders_and_executes_two_source_join(self) -> None: sql = """ SELECT monthly_files.department, lookup.label FROM monthly_files LEFT JOIN department_lookup AS lookup ON monthly_files.department = lookup.department ORDER BY monthly_files.department """ graph, _, diagnostics = compile_sql( sql, source_nodes=[inline_source(), lookup_source()], ) self.assertEqual([], diagnostics) self.assertEqual( [ "source.inline", "source.inline", "combine.join", "select", "sort", "output", ], [node.type for node in graph.nodes], ) self.assertEqual("lookup_", graph.nodes[2].config["right_prefix"]) rendered, render_diagnostics = render_sql(graph) self.assertEqual([], render_diagnostics) roundtrip, _, _ = compile_sql( rendered, source_nodes=[inline_source(), lookup_source()], ) self.assertEqual( graph.model_dump(mode="json"), roundtrip.model_dump(mode="json"), ) result = execute_preview(graph, row_limit=100) self.assertEqual( [ {"department": "A", "lookup_label": "Administration"}, {"department": "A", "lookup_label": "Administration"}, {"department": "B", "lookup_label": "Building services"}, ], result.rows, ) with self.assertRaisesRegex(SqlCompilationError, "source-qualified"): compile_sql( """ SELECT label FROM monthly_files JOIN department_lookup AS lookup ON monthly_files.department = lookup.department """, source_nodes=[inline_source(), lookup_source()], ) def test_distinct_and_derived_columns_are_deterministic(self) -> None: source = inline_source().model_copy( update={ "config": { "source_name": "monthly_files", "rows": [ {"first": " Ada ", "last": "Lovelace"}, {"first": " Ada ", "last": "Lovelace"}, ], } } ) derive = GraphNode( id="derive", type="derive", label="Display name", position=GraphPosition(x=260, y=160), config={ "target_column": "display_name", "operation": "concat", "source_columns": ["first", "last"], "separator": " ", }, ) distinct = GraphNode( id="distinct", type="distinct", label="Unique people", position=GraphPosition(x=480, y=160), config={"columns": ["display_name"]}, ) output = GraphNode( id="output", type="output", label="Output", position=GraphPosition(x=700, y=160), config={}, ) graph = PipelineGraph( nodes=[source, derive, distinct, output], edges=[ GraphEdge(id="source-derive", source=source.id, target=derive.id), GraphEdge(id="derive-distinct", source=derive.id, target=distinct.id), GraphEdge(id="distinct-output", source=distinct.id, target=output.id), ], ) result = execute_preview(graph, row_limit=100) self.assertEqual( [{"first": " Ada ", "last": "Lovelace", "display_name": " Ada Lovelace"}], result.rows, ) def test_union_can_keep_or_remove_duplicate_rows(self) -> None: first = inline_source().model_copy( update={ "id": "first", "config": {"source_name": "first", "rows": [{"id": 1}, {"id": 2}]}, } ) second = inline_source().model_copy( update={ "id": "second", "config": {"source_name": "second", "rows": [{"id": 2}, {"id": 3}]}, } ) union = GraphNode( id="union", type="combine.union", label="Append", position=GraphPosition(x=300, y=200), config={"mode": "distinct"}, ) output = GraphNode( id="output", type="output", label="Output", position=GraphPosition(x=520, y=200), config={}, ) graph = PipelineGraph( nodes=[first, second, union, output], edges=[ GraphEdge(id="first-union", source=first.id, target=union.id), GraphEdge(id="second-union", source=second.id, target=union.id), GraphEdge(id="union-output", source=union.id, target=output.id), ], ) result = execute_preview(graph, row_limit=100) self.assertEqual([{"id": 1}, {"id": 2}, {"id": 3}], result.rows) rendered, diagnostics = render_sql(graph) self.assertEqual([], diagnostics) self.assertIn("UNION BY NAME", rendered) roundtrip, normalized, diagnostics = compile_sql( rendered, source_nodes=[first, second], ) self.assertEqual(rendered, normalized) self.assertEqual([], diagnostics) self.assertEqual( ["source.inline", "source.inline", "combine.union", "output"], [node.type for node in roundtrip.nodes], ) self.assertEqual(result.rows, execute_preview(roundtrip, row_limit=100).rows) def test_legacy_composite_edge_ids_remain_readable(self) -> None: edge_id = f"edge-{'a' * 100}-{'b' * 100}" edge = GraphEdge(id=edge_id, source="source", target="output") self.assertEqual(edge_id, edge.id) def test_connector_source_uses_bounded_resolver_and_records_lineage(self) -> None: source = GraphNode( id="connector", type="source.reference", label="Monthly import", position=GraphPosition(x=40, y=160), config={ "source_ref": "snapshot:source-1", "source_name": "monthly_import", "expected_fingerprint": "sha256:expected", }, ) output = GraphNode( id="output", type="output", label="Output", position=GraphPosition(x=280, y=160), config={}, ) graph = PipelineGraph( nodes=[source, output], edges=[GraphEdge(id="source-output", source=source.id, target=output.id)], ) calls: list[tuple[str, int]] = [] def resolve(node: GraphNode, limit: int) -> ResolvedSource: calls.append((node.id, limit)) return ResolvedSource( rows=({"id": 1}, {"id": 2}), source_ref="snapshot:source-1", provider="snapshot", fingerprint="sha256:expected", total_rows=20, truncated=True, ) result = execute_preview( graph, row_limit=100, source_resolver=resolve, ) self.assertEqual([("connector", MAX_SOURCE_ROWS)], calls) self.assertEqual([{"id": 1}, {"id": 2}], result.rows) self.assertEqual(20, result.source_fingerprints[0]["row_count"]) self.assertTrue(result.source_fingerprints[0]["truncated"]) self.assertEqual("source.preview_truncated", result.diagnostics[0].code) def test_graph_validation_rejects_cycle(self) -> None: source = inline_source() output = GraphNode( id="output", type="output", label="Output", position=GraphPosition(x=300, y=160), config={}, ) graph = PipelineGraph( nodes=[source, output], edges=[ GraphEdge(id="forward", source="source", target="output"), GraphEdge(id="back", source="output", target="source"), ], ) codes = {item.code for item in validate_graph(graph)} self.assertIn("graph.cycle", codes) def test_schema_validation_reports_unknown_and_colliding_columns(self) -> None: source = inline_source().model_copy( update={ "config": { "source_name": "monthly_files", "rows": [{"id": 1, "right_id": "reserved"}], } } ) lookup = lookup_source().model_copy( update={ "config": { "source_name": "department_lookup", "rows": [{"id": 1}], } } ) join = GraphNode( id="join", type="combine.join", label="Join", position=GraphPosition(x=280, y=200), config={ "join_type": "inner", "left_keys": ["missing_id"], "right_keys": ["id"], "right_prefix": "right_", }, ) output = GraphNode( id="output", type="output", label="Output", position=GraphPosition(x=500, y=200), config={}, ) graph = PipelineGraph( nodes=[source, lookup, join, output], edges=[ GraphEdge( id="source-join", source=source.id, target=join.id, target_port="left", ), GraphEdge( id="lookup-join", source=lookup.id, target=join.id, target_port="right", ), GraphEdge(id="join-output", source=join.id, target=output.id), ], ) diagnostics = validate_graph(graph) self.assertTrue( any( item.code == "schema.unknown_column" and item.node_id == join.id and item.field == "left_keys" for item in diagnostics ) ) self.assertTrue( any(item.code == "join.output_collision" for item in diagnostics) ) def test_preview_truncates_response_without_changing_pipeline_limit(self) -> None: source = inline_source() output = GraphNode( id="output", type="output", label="Output", position=GraphPosition(x=300, y=160), config={}, ) graph = PipelineGraph( nodes=[source, output], edges=[GraphEdge(id="edge", source="source", target="output")], ) result = execute_preview(graph, row_limit=2) self.assertEqual(2, len(result.rows)) self.assertEqual(3, result.total_rows) self.assertTrue(result.truncated) def test_preview_can_return_a_bounded_intermediate_node_result(self) -> None: source = inline_source() filter_node = GraphNode( id="filter", type="filter", label="Open cases", position=GraphPosition(x=240, y=160), config={"column": "status", "operator": "eq", "value": "open"}, ) output = GraphNode( id="output", type="output", label="Output", position=GraphPosition(x=440, y=160), config={}, ) graph = PipelineGraph( nodes=[source, filter_node, output], edges=[ GraphEdge(id="source-filter", source="source", target="filter"), GraphEdge(id="filter-output", source="filter", target="output"), ], ) result = execute_preview( graph, row_limit=1, preview_node_id="source", ) self.assertEqual(2, result.total_rows) self.assertIsNotNone(result.node_preview) self.assertEqual("source", result.node_preview.node_id) self.assertEqual(3, result.node_preview.total_rows) self.assertEqual(1, len(result.node_preview.rows)) self.assertTrue(result.node_preview.truncated) def test_empty_aggregate_and_null_sort_are_deterministic(self) -> None: graph, _, _ = compile_sql( """ SELECT status, COUNT(*) AS records FROM monthly_files GROUP BY status ORDER BY status DESC """, source_nodes=[ inline_source().model_copy( update={ "config": { "source_name": "monthly_files", "rows": [ {"status": None}, {"status": "open"}, {"status": "closed"}, ], } } ) ], ) first = execute_preview(graph, row_limit=100) second = execute_preview(graph, row_limit=100) self.assertEqual(first.rows, second.rows) self.assertEqual( [ {"status": "open", "records": 1}, {"status": "closed", "records": 1}, {"status": None, "records": 1}, ], first.rows, ) self.assertEqual( [ diagnostic.model_dump(exclude={"duration_ms"}) for diagnostic in first.node_diagnostics ], [ diagnostic.model_dump(exclude={"duration_ms"}) for diagnostic in second.node_diagnostics ], ) def test_empty_global_aggregate_returns_a_count_row(self) -> None: graph, _, _ = compile_sql( "SELECT COUNT(*) AS records FROM monthly_files", source_nodes=[ inline_source().model_copy( update={"config": {"source_name": "monthly_files", "rows": []}} ) ], ) result = execute_preview(graph, row_limit=100) self.assertEqual([{"records": 0}], result.rows) self.assertEqual(0, result.input_row_count) def test_invalid_runtime_comparison_is_scoped_to_the_node(self) -> None: graph, _, _ = compile_sql( "SELECT * FROM monthly_files WHERE amount > 10", source_nodes=[ inline_source().model_copy( update={ "config": { "source_name": "monthly_files", "rows": [{"amount": "not-a-number"}], } } ) ], ) with self.assertRaises(PipelineExecutionError) as raised: execute_preview(graph, row_limit=100, preview_node_id="source") self.assertEqual("filter-1", raised.exception.node_id) self.assertIn("Cannot apply", str(raised.exception)) self.assertEqual( [("source", "succeeded"), ("filter-1", "failed")], [ (diagnostic.node_id, diagnostic.status) for diagnostic in raised.exception.node_diagnostics ], ) self.assertEqual(1, len(raised.exception.source_fingerprints)) self.assertIsNotNone(raised.exception.node_preview) self.assertEqual("source", raised.exception.node_preview.node_id) self.assertEqual( [{"amount": "not-a-number"}], raised.exception.node_preview.rows, ) if __name__ == "__main__": unittest.main()