586 lines
20 KiB
Python
586 lines
20 KiB
Python
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)
|
|
|
|
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()
|