diff --git a/src/govoplan_dataflow/backend/executor.py b/src/govoplan_dataflow/backend/executor.py index 41308de..ca5951f 100644 --- a/src/govoplan_dataflow/backend/executor.py +++ b/src/govoplan_dataflow/backend/executor.py @@ -4,9 +4,9 @@ import hashlib import json import time from collections import defaultdict -from dataclasses import dataclass +from dataclasses import dataclass, field from decimal import Decimal -from typing import Any, Callable +from typing import Any, Callable, NoReturn from govoplan_dataflow.backend.expressions import ( convert_value, @@ -84,6 +84,26 @@ class ResolvedSource: SourceResolver = Callable[[GraphNode, int], ResolvedSource] +@dataclass(frozen=True, slots=True) +class _PreviewPlan: + graph: PipelineGraph + node_by_id: dict[str, GraphNode] + inputs_by_port: dict[str, dict[str, list[str]]] + ordered_node_ids: list[str] + validation: list[DataflowDiagnostic] + preview_node_id: str | None + row_limit: int + + +@dataclass(slots=True) +class _PreviewState: + started: float = field(default_factory=time.monotonic) + outputs: dict[str, list[dict[str, Any]]] = field(default_factory=dict) + node_diagnostics: list[NodePreviewDiagnostic] = field(default_factory=list) + source_fingerprints: list[dict[str, Any]] = field(default_factory=list) + input_row_count: int = 0 + + def execute_preview( graph: PipelineGraph, *, @@ -92,151 +112,315 @@ def execute_preview( preview_node_id: str | None = None, _execution_depth: int = 0, ) -> PipelineExecutionResult: - if _execution_depth > 5: + plan = _prepare_preview( + graph, + row_limit=row_limit, + preview_node_id=preview_node_id, + execution_depth=_execution_depth, + ) + state = _PreviewState() + for node_id in plan.ordered_node_ids: + _execute_preview_node( + plan, + state, + plan.node_by_id[node_id], + source_resolver=source_resolver, + execution_depth=_execution_depth, + ) + return _preview_result(plan, state) + + +def _prepare_preview( + graph: PipelineGraph, + *, + row_limit: int, + preview_node_id: str | None, + execution_depth: int, +) -> _PreviewPlan: + if execution_depth > 5: raise PipelineExecutionError("Subflows are limited to five nested levels.") 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) - + first_error = next( + (item for item in validation if item.severity == "error"), + None, + ) + if first_error is not None: + raise PipelineExecutionError( + first_error.message, + node_id=first_error.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: raise PipelineExecutionError("Pipeline graph contains a cycle") + return _PreviewPlan( + graph=graph, + node_by_id=node_by_id, + inputs_by_port=graph_inputs_by_port(graph), + ordered_node_ids=ordered, + validation=validation, + preview_node_id=preview_node_id, + row_limit=row_limit, + ) - outputs: dict[str, list[dict[str, Any]]] = {} - node_diagnostics: list[NodePreviewDiagnostic] = [] - source_fingerprints: list[dict[str, Any]] = [] - started = time.monotonic() - input_row_count = 0 - for node_id in ordered: - node = node_by_id[node_id] - node_started = time.monotonic() - 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"), - node_preview=_node_preview(outputs, preview_node_id, row_limit), - ) - try: - executor = OPERATOR_REGISTRY.executor(node.type) - if executor is None: - raise PipelineExecutionError( - f"Node type {node.type!r} has no registered executor.", - node_id=node.id, - ) - execution = executor( - OperatorExecutionContext( - node=node, - inputs_by_port=node_inputs, - outputs=outputs, - input_sets=input_sets, - input_rows=input_rows, - source_resolver=source_resolver, - execution_depth=_execution_depth, - ) - ) - output_rows = execution.rows - input_row_count += execution.input_row_count - source_fingerprints.extend(execution.source_fingerprints) - node_messages.extend(execution.messages) - if node.type == "source.reference" and execution.messages: - validation.extend( - DataflowDiagnostic( - severity="warning", - code="source.preview_truncated", - message=message, - node_id=node.id, - ) - for message in execution.messages - ) - 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"), - node_preview=_node_preview(outputs, preview_node_id, row_limit), - ) from exc +def _execute_preview_node( + plan: _PreviewPlan, + state: _PreviewState, + node: GraphNode, + *, + source_resolver: SourceResolver | None, + execution_depth: int, +) -> None: + node_started = time.monotonic() + node_inputs = plan.inputs_by_port.get(node.id, {}) + input_sets = [ + state.outputs[source_id] + for port_sources in node_inputs.values() + for source_id in port_sources + ] + _enforce_preview_deadline(plan, state, node, input_sets) + try: + execution = _run_operator( + node, + node_inputs=node_inputs, + input_sets=input_sets, + outputs=state.outputs, + source_resolver=source_resolver, + execution_depth=execution_depth, + ) + _record_execution(plan, state, node, execution) + _guard_result_size(execution.rows, node_id=node.id) + except ( + PipelineExecutionError, + ArithmeticError, + KeyError, + TypeError, + ValueError, + ) as exc: + _raise_preview_node_failure( + plan, + state, + node, + input_sets=input_sets, + node_started=node_started, + exc=exc, + ) + state.outputs[node.id] = execution.rows + state.node_diagnostics.append( + _successful_node_diagnostic( + node, + execution, + input_sets=input_sets, + node_started=node_started, + ) + ) - outputs[node.id] = output_rows - node_diagnostics.append( - NodePreviewDiagnostic( + + +def _run_operator( + node: GraphNode, + *, + node_inputs: dict[str, list[str]], + input_sets: list[list[dict[str, Any]]], + outputs: dict[str, list[dict[str, Any]]], + source_resolver: SourceResolver | None, + execution_depth: int, +) -> OperatorExecutionResult: + executor = OPERATOR_REGISTRY.executor(node.type) + if executor is None: + raise PipelineExecutionError( + f"Node type {node.type!r} has no registered executor.", + node_id=node.id, + ) + return executor( + OperatorExecutionContext( + node=node, + inputs_by_port=node_inputs, + outputs=outputs, + input_sets=input_sets, + input_rows=input_sets[0] if len(input_sets) == 1 else [], + source_resolver=source_resolver, + execution_depth=execution_depth, + ) + ) + + +def _record_execution( + plan: _PreviewPlan, + state: _PreviewState, + node: GraphNode, + execution: OperatorExecutionResult, +) -> None: + state.input_row_count += execution.input_row_count + state.source_fingerprints.extend(execution.source_fingerprints) + if node.type == "source.reference": + plan.validation.extend( + DataflowDiagnostic( + severity="warning", + code="source.preview_truncated", + message=message, node_id=node.id, - status="succeeded", - 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, ) + for message in execution.messages ) - output_node = next(node for node in graph.nodes if node.type == "output") - all_rows = outputs[output_node.id] - rows = all_rows[:row_limit] + +def _guard_result_size( + rows: list[dict[str, Any]], + *, + node_id: str, +) -> None: + if len(json.dumps(rows, default=str).encode("utf-8")) <= MAX_RESULT_BYTES: + return + raise PipelineExecutionError( + "A preview node exceeded the one-megabyte result limit.", + node_id=node_id, + ) + + +def _enforce_preview_deadline( + plan: _PreviewPlan, + state: _PreviewState, + node: GraphNode, + input_sets: list[list[dict[str, Any]]], +) -> None: + if time.monotonic() - state.started <= MAX_EXECUTION_SECONDS: + return + message = "Preview exceeded the two-second execution limit" + failed = NodePreviewDiagnostic( + node_id=node.id, + status="failed", + input_rows=_input_row_total(input_sets), + output_rows=0, + duration_ms=0, + columns=[], + messages=[message], + ) + raise PipelineExecutionError( + message, + node_id=node.id, + node_diagnostics=tuple([*state.node_diagnostics, failed]), + source_fingerprints=tuple(state.source_fingerprints), + input_row_count=state.input_row_count, + diagnostics=tuple(_warnings(plan.validation)), + node_preview=_node_preview( + state.outputs, + plan.preview_node_id, + plan.row_limit, + ), + ) + + +def _raise_preview_node_failure( + plan: _PreviewPlan, + state: _PreviewState, + node: GraphNode, + *, + input_sets: list[list[dict[str, Any]]], + node_started: float, + exc: BaseException, +) -> NoReturn: + 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=_input_row_total(input_sets), + output_rows=0, + duration_ms=_elapsed_ms(node_started), + columns=[], + messages=[str(execution_error)], + ) + raise PipelineExecutionError( + str(execution_error), + node_id=execution_error.node_id or node.id, + node_diagnostics=tuple([*state.node_diagnostics, failed]), + source_fingerprints=tuple(state.source_fingerprints), + input_row_count=state.input_row_count, + diagnostics=tuple(_warnings(plan.validation)), + node_preview=_node_preview( + state.outputs, + plan.preview_node_id, + plan.row_limit, + ), + ) from exc + + +def _successful_node_diagnostic( + node: GraphNode, + execution: OperatorExecutionResult, + *, + input_sets: list[list[dict[str, Any]]], + node_started: float, +) -> NodePreviewDiagnostic: + return NodePreviewDiagnostic( + node_id=node.id, + status="succeeded", + input_rows=_input_row_total(input_sets), + output_rows=len(execution.rows), + duration_ms=_elapsed_ms(node_started), + columns=infer_columns(execution.rows), + messages=list(execution.messages), + ) + + +def _preview_result( + plan: _PreviewPlan, + state: _PreviewState, +) -> PipelineExecutionResult: + output_node = next( + node + for node in plan.graph.nodes + if node.type == "output" + ) + all_rows = state.outputs[output_node.id] + rows = all_rows[: plan.row_limit] return PipelineExecutionResult( rows=rows, total_rows=len(all_rows), truncated=len(rows) < len(all_rows), 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, + diagnostics=_warnings(plan.validation), + node_diagnostics=state.node_diagnostics, + node_preview=_node_preview( + state.outputs, + plan.preview_node_id, + plan.row_limit, + ), + source_fingerprints=state.source_fingerprints, + input_row_count=state.input_row_count, ) +def _warnings( + diagnostics: list[DataflowDiagnostic], +) -> list[DataflowDiagnostic]: + return [ + item + for item in diagnostics + if item.severity != "error" + ] + + +def _input_row_total( + input_sets: list[list[dict[str, Any]]], +) -> int: + return sum(len(rows) for rows in input_sets) + + +def _elapsed_ms(started: float) -> float: + return round((time.monotonic() - started) * 1000, 3) + + def _node_preview( outputs: dict[str, list[dict[str, Any]]], node_id: str | None, @@ -397,32 +581,67 @@ def _derive_rows( 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) + handler = _DERIVE_HANDLERS.get(operation) + if handler is None: + raise ValueError(f"unknown derive operation {operation!r}") + return handler(values, separator) + + +def _unary_text( + values: list[Any], + _separator: str, + operation: Callable[[str], str], +) -> str | None: + return None if values[0] is None else operation(str(values[0])) + + +def _derive_concat(values: list[Any], separator: str) -> str: + return separator.join(str(value) for value in values if value is not None) + + +def _derive_coalesce(values: list[Any], _separator: str) -> Any: + return next((value for value in values if value not in (None, "")), None) + + +def _derive_numeric( + values: list[Any], + _separator: str, + operation: Callable[[Any, Any], Any], +) -> Any: 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}") + return operation(left, right) + + +_DERIVE_HANDLERS: dict[str, Callable[[list[Any], str], Any]] = { + "copy": lambda values, _separator: values[0], + "upper": lambda values, separator: _unary_text( + values, separator, str.upper + ), + "lower": lambda values, separator: _unary_text( + values, separator, str.lower + ), + "trim": lambda values, separator: _unary_text( + values, separator, str.strip + ), + "concat": _derive_concat, + "coalesce": _derive_coalesce, + "add": lambda values, separator: _derive_numeric( + values, separator, lambda left, right: left + right + ), + "subtract": lambda values, separator: _derive_numeric( + values, separator, lambda left, right: left - right + ), + "multiply": lambda values, separator: _derive_numeric( + values, separator, lambda left, right: left * right + ), + "divide": lambda values, separator: _derive_numeric( + values, separator, lambda left, right: left / right + ), +} def _expression_filter_rows( @@ -848,46 +1067,95 @@ def _aggregate_rows( ) -> list[dict[str, Any]]: group_by = [str(item) for item in config.get("group_by", [])] aggregates = list(config["aggregates"]) + grouped = _group_rows(rows, group_by=group_by) + return [ + _aggregate_group( + key, + group_rows, + group_by=group_by, + aggregates=aggregates, + node_id=node_id, + ) + for key, group_rows in grouped.items() + ] + + +def _group_rows( + rows: list[dict[str, Any]], + *, + group_by: list[str], +) -> dict[tuple[Any, ...], list[dict[str, Any]]]: grouped: dict[tuple[Any, ...], list[dict[str, Any]]] = defaultdict(list) if rows: for row in rows: grouped[tuple(row.get(column) for column in group_by)].append(row) elif not group_by: grouped[()] = [] + return grouped - output: list[dict[str, Any]] = [] - for key, group_rows in grouped.items(): - result = {column: value for column, value in zip(group_by, key, strict=True)} - for aggregate in aggregates: - function = str(aggregate["function"]) - column = aggregate.get("column") - alias = str(aggregate["alias"]) - values = [row.get(column) for row in group_rows if column is not None and row.get(column) is not None] - try: - if function == "count": - result[alias] = len(group_rows) if column in (None, "", "*") else len(values) - elif function == "sum": - result[alias] = sum(values) if values else 0 - elif function == "avg": - result[alias] = sum(values) / len(values) if values else None - elif function == "min": - result[alias] = min(values) if values else None - elif function == "max": - result[alias] = max(values) if values else None - except (ArithmeticError, TypeError, ValueError) as exc: - raise PipelineExecutionError( - f"Cannot calculate {function.upper()} for {column!r}: {exc}", - node_id=node_id, - ) from exc - output.append(result) - return output + +def _aggregate_group( + key: tuple[Any, ...], + rows: list[dict[str, Any]], + *, + group_by: list[str], + aggregates: list[dict[str, Any]], + node_id: str, +) -> dict[str, Any]: + result = { + column: value + for column, value in zip(group_by, key, strict=True) + } + for aggregate in aggregates: + alias = str(aggregate["alias"]) + result[alias] = _aggregate_value( + str(aggregate["function"]), + aggregate.get("column"), + rows, + node_id=node_id, + ) + return result + + +def _aggregate_value( + function: str, + column: Any, + rows: list[dict[str, Any]], + *, + node_id: str, +) -> Any: + values = [ + row.get(column) + for row in rows + if column is not None and row.get(column) is not None + ] + try: + if function == "count": + return len(rows) if column in (None, "", "*") else len(values) + handler = _AGGREGATE_HANDLERS.get(function) + if handler is None: + raise ValueError(f"unknown aggregate function {function!r}") + return handler(values) + except (ArithmeticError, TypeError, ValueError) as exc: + raise PipelineExecutionError( + f"Cannot calculate {function.upper()} for {column!r}: {exc}", + node_id=node_id, + ) from exc + + +_AGGREGATE_HANDLERS: dict[str, Callable[[list[Any]], Any]] = { + "sum": lambda values: sum(values) if values else 0, + "avg": lambda values: sum(values) / len(values) if values else None, + "min": lambda values: min(values) if values else None, + "max": lambda values: max(values) if values else None, +} def _sort_rows(rows: list[dict[str, Any]], config: dict[str, Any]) -> list[dict[str, Any]]: result = [dict(row) for row in rows] - for field in reversed(config["fields"]): - column = str(field["column"]) - reverse = field.get("direction", "asc") == "desc" + for field_config in reversed(config["fields"]): + column = str(field_config["column"]) + reverse = field_config.get("direction", "asc") == "desc" concrete = [row for row in result if row.get(column) is not None] nulls = [row for row in result if row.get(column) is None] concrete.sort(key=lambda row: _sortable_value(row[column]), reverse=reverse) diff --git a/src/govoplan_dataflow/backend/expressions.py b/src/govoplan_dataflow/backend/expressions.py index d347521..55aa778 100644 --- a/src/govoplan_dataflow/backend/expressions.py +++ b/src/govoplan_dataflow/backend/expressions.py @@ -1,8 +1,9 @@ from __future__ import annotations +import operator from dataclasses import dataclass from decimal import Decimal -from typing import Any, Literal +from typing import Any, Callable, Literal import sqlglot from sqlglot import exp @@ -141,133 +142,255 @@ def infer_expression_type( def _evaluate(expression: exp.Expression, row: dict[str, Any]) -> Any: - if isinstance(expression, exp.Paren): - return _evaluate(expression.this, row) - if isinstance(expression, exp.Column): - return row.get(expression.name) - if isinstance(expression, exp.Null): - return None - if isinstance(expression, exp.Boolean): - return bool(expression.this) - if isinstance(expression, exp.Literal): - return _literal(expression) - if isinstance(expression, exp.Neg): - value = _evaluate(expression.this, row) - return None if value is None else -value - if isinstance(expression, exp.Not): - return not bool(_evaluate(expression.this, row)) + evaluator = _EVALUATORS.get(type(expression)) + if evaluator is None: + raise ExpressionError( + f"Expression element {type(expression).__name__} cannot be evaluated." + ) + return evaluator(expression, row) + + +def _evaluate_child(expression: exp.Expression, row: dict[str, Any]) -> Any: + return _evaluate(expression.this, row) + + +def _evaluate_negation(expression: exp.Expression, row: dict[str, Any]) -> Any: + value = _evaluate_child(expression, row) + return None if value is None else -value + + +def _evaluate_boolean_not( + expression: exp.Expression, + row: dict[str, Any], +) -> bool: + return not bool(_evaluate_child(expression, row)) + + +def _evaluate_boolean_binary( + expression: exp.Expression, + row: dict[str, Any], +) -> bool: + left = bool(_evaluate(expression.this, row)) if isinstance(expression, exp.And): - return bool(_evaluate(expression.this, row)) and bool( - _evaluate(expression.expression, row) - ) - if isinstance(expression, exp.Or): - return bool(_evaluate(expression.this, row)) or bool( - _evaluate(expression.expression, row) - ) - if isinstance(expression, exp.Is): - left = _evaluate(expression.this, row) - right = _evaluate(expression.expression, row) - return left is right if right is None else left == right - if isinstance(expression, (exp.Add, exp.Sub, exp.Mul, exp.Div, exp.Mod)): - left = _evaluate(expression.this, row) - right = _evaluate(expression.expression, row) - if left is None or right is None: - return None - if isinstance(expression, exp.Add): - return left + right - if isinstance(expression, exp.Sub): - return left - right - if isinstance(expression, exp.Mul): - return left * right - if isinstance(expression, exp.Div): - return left / right - return left % right - if isinstance(expression, (exp.EQ, exp.NEQ, exp.GT, exp.GTE, exp.LT, exp.LTE)): - left = _evaluate(expression.this, row) - right = _evaluate(expression.expression, row) - if isinstance(expression, exp.EQ): - return left == right - if isinstance(expression, exp.NEQ): - return left != right - if left is None or right is None: - return False - if isinstance(expression, exp.GT): - return left > right - if isinstance(expression, exp.GTE): - return left >= right - if isinstance(expression, exp.LT): - return left < right - return left <= right - if isinstance(expression, exp.Lower): - return _string_call(expression.this, row, str.lower) - if isinstance(expression, exp.Upper): - return _string_call(expression.this, row, str.upper) - if isinstance(expression, exp.Trim): - return _string_call(expression.this, row, str.strip) - if isinstance(expression, exp.Length): - value = _evaluate(expression.this, row) - return None if value is None else len(value) - if isinstance(expression, exp.Abs): - value = _evaluate(expression.this, row) - return None if value is None else abs(value) - if isinstance(expression, exp.Round): - value = _evaluate(expression.this, row) - decimals = expression.args.get("decimals") - places = int(_evaluate(decimals, row)) if decimals is not None else 0 - return None if value is None else round(value, places) - if isinstance(expression, exp.Coalesce): - arguments = [expression.this, *expression.expressions] - return next( - ( - value - for argument in arguments - if (value := _evaluate(argument, row)) is not None - ), - None, - ) - if isinstance(expression, exp.Concat): - return "".join( - "" if (value := _evaluate(item, row)) is None else str(value) - for item in expression.expressions - ) - if isinstance(expression, exp.Replace): - value = _evaluate(expression.this, row) - old = _evaluate(expression.expression, row) - new = _evaluate(expression.args["replacement"], row) - return None if value is None else str(value).replace(str(old), str(new)) - if isinstance(expression, exp.Substring): - value = _evaluate(expression.this, row) - if value is None: - return None - start = int(_evaluate(expression.args["start"], row)) - 1 - length = expression.args.get("length") - return ( - str(value)[start:] - if length is None - else str(value)[start : start + int(_evaluate(length, row))] - ) - if isinstance(expression, exp.Cast): - return convert_value( - _evaluate(expression.this, row), - _cast_type(expression), - on_error="fail", - ) - if isinstance(expression, exp.Case): - base = _evaluate(expression.this, row) if expression.this is not None else None - for condition in expression.args.get("ifs") or []: - if not isinstance(condition, exp.If): - continue - candidate = _evaluate(condition.this, row) - matched = candidate == base if expression.this is not None else bool(candidate) - if matched: - return _evaluate(condition.args["true"], row) - default = expression.args.get("default") - return _evaluate(default, row) if default is not None else None - raise ExpressionError( - f"Expression element {type(expression).__name__} cannot be evaluated." + return left and bool(_evaluate(expression.expression, row)) + return left or bool(_evaluate(expression.expression, row)) + + +def _evaluate_is(expression: exp.Expression, row: dict[str, Any]) -> bool: + left = _evaluate(expression.this, row) + right = _evaluate(expression.expression, row) + return left is right if right is None else left == right + + +def _evaluate_binary( + expression: exp.Expression, + row: dict[str, Any], + operations: dict[type[exp.Expression], Callable[[Any, Any], Any]], + *, + null_result: Any, +) -> Any: + left = _evaluate(expression.this, row) + right = _evaluate(expression.expression, row) + if left is None or right is None: + return null_result + return operations[type(expression)](left, right) + + +def _evaluate_arithmetic(expression: exp.Expression, row: dict[str, Any]) -> Any: + return _evaluate_binary( + expression, + row, + _ARITHMETIC_OPERATIONS, + null_result=None, ) +def _evaluate_comparison( + expression: exp.Expression, + row: dict[str, Any], +) -> bool: + if isinstance(expression, (exp.EQ, exp.NEQ)): + left = _evaluate(expression.this, row) + right = _evaluate(expression.expression, row) + return _COMPARISON_OPERATIONS[type(expression)](left, right) + return _evaluate_binary( + expression, + row, + _COMPARISON_OPERATIONS, + null_result=False, + ) + + +def _evaluate_string_function( + expression: exp.Expression, + row: dict[str, Any], +) -> str | None: + return _string_call( + expression.this, + row, + _STRING_OPERATIONS[type(expression)], + ) + + +def _evaluate_length(expression: exp.Expression, row: dict[str, Any]) -> int | None: + value = _evaluate_child(expression, row) + return None if value is None else len(value) + + +def _evaluate_abs(expression: exp.Expression, row: dict[str, Any]) -> Any: + value = _evaluate_child(expression, row) + return None if value is None else abs(value) + + +def _evaluate_round(expression: exp.Expression, row: dict[str, Any]) -> Any: + value = _evaluate_child(expression, row) + decimals = expression.args.get("decimals") + places = int(_evaluate(decimals, row)) if decimals is not None else 0 + return None if value is None else round(value, places) + + +def _evaluate_coalesce(expression: exp.Expression, row: dict[str, Any]) -> Any: + arguments = [expression.this, *expression.expressions] + return next( + ( + value + for argument in arguments + if (value := _evaluate(argument, row)) is not None + ), + None, + ) + + +def _evaluate_concat(expression: exp.Expression, row: dict[str, Any]) -> str: + return "".join( + "" if (value := _evaluate(item, row)) is None else str(value) + for item in expression.expressions + ) + + +def _evaluate_replace(expression: exp.Expression, row: dict[str, Any]) -> Any: + value = _evaluate_child(expression, row) + old = _evaluate(expression.expression, row) + new = _evaluate(expression.args["replacement"], row) + return None if value is None else str(value).replace(str(old), str(new)) + + +def _evaluate_substring(expression: exp.Expression, row: dict[str, Any]) -> Any: + value = _evaluate_child(expression, row) + if value is None: + return None + start = int(_evaluate(expression.args["start"], row)) - 1 + length = expression.args.get("length") + if length is None: + return str(value)[start:] + return str(value)[start : start + int(_evaluate(length, row))] + + +def _evaluate_cast(expression: exp.Expression, row: dict[str, Any]) -> Any: + return convert_value( + _evaluate_child(expression, row), + _cast_type(expression), # type: ignore[arg-type] + on_error="fail", + ) + + +def _evaluate_case(expression: exp.Expression, row: dict[str, Any]) -> Any: + base_expression = expression.this + base = _evaluate(base_expression, row) if base_expression is not None else None + for condition in expression.args.get("ifs") or []: + if isinstance(condition, exp.If) and _case_matches( + condition, + row, + base=base, + has_base=base_expression is not None, + ): + return _evaluate(condition.args["true"], row) + default = expression.args.get("default") + return _evaluate(default, row) if default is not None else None + + +def _case_matches( + condition: exp.If, + row: dict[str, Any], + *, + base: Any, + has_base: bool, +) -> bool: + candidate = _evaluate(condition.this, row) + return candidate == base if has_base else bool(candidate) + + +def _evaluate_if(expression: exp.Expression, row: dict[str, Any]) -> Any: + branch = "true" if bool(_evaluate(expression.this, row)) else "false" + selected = expression.args.get(branch) + return _evaluate(selected, row) if selected is not None else None + + +_ARITHMETIC_OPERATIONS: dict[ + type[exp.Expression], + Callable[[Any, Any], Any], +] = { + exp.Add: operator.add, + exp.Sub: operator.sub, + exp.Mul: operator.mul, + exp.Div: operator.truediv, + exp.Mod: operator.mod, +} +_COMPARISON_OPERATIONS: dict[ + type[exp.Expression], + Callable[[Any, Any], bool], +] = { + exp.EQ: operator.eq, + exp.NEQ: operator.ne, + exp.GT: operator.gt, + exp.GTE: operator.ge, + exp.LT: operator.lt, + exp.LTE: operator.le, +} +_STRING_OPERATIONS: dict[type[exp.Expression], Callable[[str], str]] = { + exp.Lower: str.lower, + exp.Upper: str.upper, + exp.Trim: str.strip, +} +_EVALUATORS: dict[ + type[exp.Expression], + Callable[[exp.Expression, dict[str, Any]], Any], +] = { + exp.Paren: _evaluate_child, + exp.Column: lambda expression, row: row.get(expression.name), + exp.Null: lambda _expression, _row: None, + exp.Boolean: lambda expression, _row: bool(expression.this), + exp.Literal: lambda expression, _row: _literal(expression), # type: ignore[arg-type] + exp.Neg: _evaluate_negation, + exp.Not: _evaluate_boolean_not, + exp.And: _evaluate_boolean_binary, + exp.Or: _evaluate_boolean_binary, + exp.Is: _evaluate_is, + **{ + expression_type: _evaluate_arithmetic + for expression_type in _ARITHMETIC_OPERATIONS + }, + **{ + expression_type: _evaluate_comparison + for expression_type in _COMPARISON_OPERATIONS + }, + **{ + expression_type: _evaluate_string_function + for expression_type in _STRING_OPERATIONS + }, + exp.Length: _evaluate_length, + exp.Abs: _evaluate_abs, + exp.Round: _evaluate_round, + exp.Coalesce: _evaluate_coalesce, + exp.Concat: _evaluate_concat, + exp.Replace: _evaluate_replace, + exp.Substring: _evaluate_substring, + exp.Cast: _evaluate_cast, + exp.Case: _evaluate_case, + exp.If: _evaluate_if, +} + + def convert_value( value: Any, target_type: str, @@ -312,19 +435,115 @@ def _infer( expression: exp.Expression, schema: dict[str, ExpressionDataType], ) -> ExpressionDataType: - if isinstance(expression, exp.Column): - return schema.get(expression.name, "unknown") - if isinstance(expression, exp.Null): - return "null" - if isinstance(expression, exp.Boolean): - return "boolean" - if isinstance(expression, exp.Literal): - if expression.is_string: - return "string" - return "number" if "." in str(expression.this) else "integer" - if isinstance( - expression, - ( + inferer = _TYPE_INFERERS.get(type(expression)) + return inferer(expression, schema) if inferer is not None else "unknown" + + +def _infer_literal( + expression: exp.Expression, + _schema: dict[str, ExpressionDataType], +) -> ExpressionDataType: + literal = expression + if not isinstance(literal, exp.Literal): + return "unknown" + if literal.is_string: + return "string" + return "number" if "." in str(literal.this) else "integer" + + +def _infer_cast( + expression: exp.Expression, + _schema: dict[str, ExpressionDataType], +) -> ExpressionDataType: + return _cast_type(expression) # type: ignore[arg-type,return-value] + + +def _infer_numeric( + expression: exp.Expression, + schema: dict[str, ExpressionDataType], +) -> ExpressionDataType: + child_types = { + _infer(item, schema) + for item in expression.iter_expressions() + } + if "number" in child_types or isinstance(expression, exp.Div): + return "number" + return "integer" + + +def _infer_coalesce( + expression: exp.Expression, + schema: dict[str, ExpressionDataType], +) -> ExpressionDataType: + candidates = ( + _infer(item, schema) + for item in (expression.this, *expression.expressions) + ) + return next( + (item for item in candidates if item not in {"null", "unknown"}), + "unknown", + ) + + +def _infer_case( + expression: exp.Expression, + schema: dict[str, ExpressionDataType], +) -> ExpressionDataType: + types = [ + _infer(item.args["true"], schema) + for item in expression.args.get("ifs") or [] + if isinstance(item, exp.If) + ] + default = expression.args.get("default") + if default is not None: + types.append(_infer(default, schema)) + return _common_expression_type(types) + + +def _infer_if( + expression: exp.Expression, + schema: dict[str, ExpressionDataType], +) -> ExpressionDataType: + return _common_expression_type( + [ + _infer(branch, schema) + for key in ("true", "false") + if (branch := expression.args.get(key)) is not None + ] + ) + + +def _common_expression_type( + types: list[ExpressionDataType], +) -> ExpressionDataType: + concrete = {item for item in types if item not in {"null", "unknown"}} + return concrete.pop() if len(concrete) == 1 else "unknown" + + +def _infer_child( + expression: exp.Expression, + schema: dict[str, ExpressionDataType], +) -> ExpressionDataType: + return _infer(expression.this, schema) + + +_TYPE_INFERERS: dict[ + type[exp.Expression], + Callable[ + [exp.Expression, dict[str, ExpressionDataType]], + ExpressionDataType, + ], +] = { + exp.Column: lambda expression, schema: schema.get( + expression.name, + "unknown", + ), + exp.Null: lambda _expression, _schema: "null", + exp.Boolean: lambda _expression, _schema: "boolean", + exp.Literal: _infer_literal, + **{ + expression_type: lambda _expression, _schema: "boolean" + for expression_type in ( exp.EQ, exp.NEQ, exp.GT, @@ -335,41 +554,39 @@ def _infer( exp.Or, exp.Is, exp.Not, - ), - ): - return "boolean" - if isinstance(expression, (exp.Length,)): - return "integer" - if isinstance(expression, (exp.Lower, exp.Upper, exp.Trim, exp.Concat, exp.Replace, exp.Substring)): - return "string" - if isinstance(expression, exp.Cast): - return _cast_type(expression) # type: ignore[return-value] - if isinstance(expression, (exp.Add, exp.Sub, exp.Mul, exp.Div, exp.Mod, exp.Abs, exp.Round)): - child_types = { - _infer(item, schema) - for item in expression.iter_expressions() - } - return "number" if "number" in child_types or isinstance(expression, exp.Div) else "integer" - if isinstance(expression, exp.Coalesce): - types = [ - _infer(item, schema) - for item in (expression.this, *expression.expressions) - ] - return next((item for item in types if item not in {"null", "unknown"}), "unknown") - if isinstance(expression, exp.Case): - types = [ - _infer(item.args["true"], schema) - for item in expression.args.get("ifs") or [] - if isinstance(item, exp.If) - ] - default = expression.args.get("default") - if default is not None: - types.append(_infer(default, schema)) - concrete = {item for item in types if item not in {"null", "unknown"}} - return concrete.pop() if len(concrete) == 1 else "unknown" - if isinstance(expression, (exp.Paren, exp.Neg)): - return _infer(expression.this, schema) - return "unknown" + ) + }, + exp.Length: lambda _expression, _schema: "integer", + **{ + expression_type: lambda _expression, _schema: "string" + for expression_type in ( + exp.Lower, + exp.Upper, + exp.Trim, + exp.Concat, + exp.Replace, + exp.Substring, + ) + }, + exp.Cast: _infer_cast, + **{ + expression_type: _infer_numeric + for expression_type in ( + exp.Add, + exp.Sub, + exp.Mul, + exp.Div, + exp.Mod, + exp.Abs, + exp.Round, + ) + }, + exp.Coalesce: _infer_coalesce, + exp.Case: _infer_case, + exp.If: _infer_if, + exp.Paren: _infer_child, + exp.Neg: _infer_child, +} def _literal(expression: exp.Literal) -> Any: diff --git a/src/govoplan_dataflow/backend/graph.py b/src/govoplan_dataflow/backend/graph.py index 54934eb..6bf6417 100644 --- a/src/govoplan_dataflow/backend/graph.py +++ b/src/govoplan_dataflow/backend/graph.py @@ -4,7 +4,6 @@ import hashlib import json import re from collections import deque -from dataclasses import dataclass, field from typing import Any from govoplan_core.core.definition_graphs import ( @@ -17,12 +16,11 @@ from govoplan_dataflow.backend.node_library import ( NODE_TYPES, node_definition, ) -from govoplan_dataflow.backend.expressions import ( - ExpressionError, - infer_expression_type, - parse_expression, -) from govoplan_dataflow.backend.operator_registry import OPERATOR_REGISTRY +from govoplan_dataflow.backend.schema_validation import ( + register_schema_propagators, + validate_graph_schemas, +) from govoplan_dataflow.backend.subflows import ( SubflowParameterError, substitute_parameters, @@ -140,7 +138,40 @@ def preserve_compatible_graph_layout( def validate_graph(graph: PipelineGraph) -> list[DataflowDiagnostic]: - diagnostics = [ + diagnostics = _definition_graph_diagnostics(graph) + nodes, incoming, outgoing = _valid_adjacency(graph) + source_nodes = [ + node + for node in graph.nodes + if node.type.startswith("source.") + ] + output_nodes = [ + node + for node in graph.nodes + if node.type == "output" + ] + diagnostics.extend(_source_name_diagnostics(source_nodes)) + diagnostics.extend(_node_config_diagnostics(graph.nodes)) + ordered, cyclic = topological_order(graph) + if not cyclic and source_nodes and len(output_nodes) == 1: + diagnostics.extend( + _topology_and_schema_diagnostics( + graph, + nodes=nodes, + incoming=incoming, + outgoing=outgoing, + source_nodes=source_nodes, + output_node=output_nodes[0], + ordered=ordered, + ) + ) + return _dedupe_diagnostics(diagnostics) + + +def _definition_graph_diagnostics( + graph: PipelineGraph, +) -> list[DataflowDiagnostic]: + return [ DataflowDiagnostic( severity=item.severity, code=item.code, @@ -163,26 +194,51 @@ def validate_graph(graph: PipelineGraph) -> list[DataflowDiagnostic]: ), ) ] + + +def _valid_adjacency( + graph: PipelineGraph, +) -> tuple[ + dict[str, GraphNode], + dict[str, list[GraphEdge]], + dict[str, list[str]], +]: nodes = {node.id: node for node in graph.nodes} incoming: dict[str, list[GraphEdge]] = {node_id: [] for node_id in nodes} outgoing: dict[str, list[str]] = {node_id: [] for node_id in nodes} for edge in graph.edges: - if edge.source not in nodes or edge.target not in nodes or edge.source == edge.target: - continue - source_definition = node_definition(nodes[edge.source].type) - target_definition = node_definition(nodes[edge.target].type) - if ( - source_definition is None - or target_definition is None - or edge.source_port not in {port.id for port in source_definition.output_ports} - or edge.target_port not in {port.id for port in target_definition.input_ports} - ): + if not _valid_graph_edge(edge, nodes): continue outgoing[edge.source].append(edge.target) incoming[edge.target].append(edge) + return nodes, incoming, outgoing - source_nodes = [node for node in graph.nodes if node.type.startswith("source.")] - output_nodes = [node for node in graph.nodes if node.type == "output"] + +def _valid_graph_edge( + edge: GraphEdge, + nodes: dict[str, GraphNode], +) -> bool: + if ( + edge.source not in nodes + or edge.target not in nodes + or edge.source == edge.target + ): + return False + source_definition = node_definition(nodes[edge.source].type) + target_definition = node_definition(nodes[edge.target].type) + if source_definition is None or target_definition is None: + return False + return ( + edge.source_port + in {port.id for port in source_definition.output_ports} + and edge.target_port + in {port.id for port in target_definition.input_ports} + ) + + +def _source_name_diagnostics( + source_nodes: list[GraphNode], +) -> list[DataflowDiagnostic]: source_names = [ str(node.config.get("source_name", "")).strip() for node in source_nodes @@ -196,15 +252,22 @@ def validate_graph(graph: PipelineGraph) -> list[DataflowDiagnostic]: }, key=str.casefold, ) - if duplicate_source_names: - diagnostics.append( - _error( - "source.duplicate_name", - "Logical source names must be unique: " - f"{', '.join(duplicate_source_names)}.", - ) + if not duplicate_source_names: + return [] + return [ + _error( + "source.duplicate_name", + "Logical source names must be unique: " + f"{', '.join(duplicate_source_names)}.", ) - for node in graph.nodes: + ] + + +def _node_config_diagnostics( + nodes: list[GraphNode], +) -> list[DataflowDiagnostic]: + diagnostics: list[DataflowDiagnostic] = [] + for node in nodes: if node.type not in SUPPORTED_NODE_TYPES: continue validator = OPERATOR_REGISTRY.config_validator(node.type) @@ -218,31 +281,53 @@ def validate_graph(graph: PipelineGraph) -> list[DataflowDiagnostic]: ) continue diagnostics.extend(validator(node)) + return diagnostics - ordered, cyclic = topological_order(graph) - if not cyclic and source_nodes and len(output_nodes) == 1: - reachable: set[str] = set() - for source in source_nodes: - reachable.update(_reachable_from(source.id, outgoing)) - if len(reachable) != len(nodes): - diagnostics.append( - _error("graph.disconnected", "Every node must be connected to a pipeline source.") + +def _topology_and_schema_diagnostics( + graph: PipelineGraph, + *, + nodes: dict[str, GraphNode], + incoming: dict[str, list[GraphEdge]], + outgoing: dict[str, list[str]], + source_nodes: list[GraphNode], + output_node: GraphNode, + ordered: list[str], +) -> list[DataflowDiagnostic]: + diagnostics: list[DataflowDiagnostic] = [] + reachable = set().union( + *( + _reachable_from(source.id, outgoing) + for source in source_nodes + ) + ) + if len(reachable) != len(nodes): + diagnostics.append( + _error( + "graph.disconnected", + "Every node must be connected to a pipeline source.", ) - reverse_adjacency = { - node_id: [edge.source for edge in edges] - for node_id, edges in incoming.items() - } - reaches_output = _reachable_from(output_nodes[0].id, reverse_adjacency) - if len(reaches_output) != len(nodes): - diagnostics.append( - _error("graph.dead_end", "Every node must lead to the pipeline output.") + ) + reverse_adjacency = { + node_id: [edge.source for edge in edges] + for node_id, edges in incoming.items() + } + if len(_reachable_from(output_node.id, reverse_adjacency)) != len(nodes): + diagnostics.append( + _error( + "graph.dead_end", + "Every node must lead to the pipeline output.", ) - if ordered and ordered[-1] != output_nodes[0].id: - diagnostics.append( - _error("graph.output_not_terminal", "The output node must be the terminal transform.") + ) + if ordered and ordered[-1] != output_node.id: + diagnostics.append( + _error( + "graph.output_not_terminal", + "The output node must be the terminal transform.", ) - diagnostics.extend(_validate_graph_schemas(graph, ordered=ordered)) - return _dedupe_diagnostics(diagnostics) + ) + diagnostics.extend(validate_graph_schemas(graph, ordered=ordered)) + return diagnostics def _structural_node_keys( @@ -363,1220 +448,802 @@ def _reachable_from(start: str, adjacency: dict[str, list[str]]) -> set[str]: return seen -@dataclass(frozen=True) -class _SchemaState: - columns: frozenset[str] - open: bool = False - types: dict[str, str] = field(default_factory=dict) - - def knows(self, column: str) -> bool: - return self.open or column in self.columns - - def type_of(self, column: str) -> str: - return self.types.get(column, "unknown") +def _validate_source(node: GraphNode) -> list[DataflowDiagnostic]: + diagnostics = _validate_source_name(node) + if node.type == "source.reference": + diagnostics.extend(_validate_reference_source(node)) + else: + diagnostics.extend(_validate_inline_source(node)) + return diagnostics -def _validate_graph_schemas( - graph: PipelineGraph, +def _validate_source_name(node: GraphNode) -> list[DataflowDiagnostic]: + source_name = node.config.get("source_name") + if not _non_empty_text(source_name): + return [ + _node_field_error( + node, + "source.name_required", + "A source needs a logical SQL name.", + "source_name", + ) + ] + if re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", str(source_name).strip()): + return [] + return [ + _node_field_error( + node, + "source.name_invalid", + "Logical source names must be SQL identifiers, such as monthly_cases.", + "source_name", + ) + ] + + +def _validate_reference_source(node: GraphNode) -> list[DataflowDiagnostic]: + diagnostics: list[DataflowDiagnostic] = [] + if not _non_empty_text(node.config.get("source_ref")): + diagnostics.append( + _node_field_error( + node, + "source.reference_required", + "Choose a datasource.", + "source_ref", + ) + ) + expected = node.config.get("expected_fingerprint") + if expected is not None and not isinstance(expected, str): + diagnostics.append( + _node_field_error( + node, + "source.fingerprint", + "The expected source fingerprint must be text.", + "expected_fingerprint", + ) + ) + return diagnostics + + +def _validate_inline_source(node: GraphNode) -> list[DataflowDiagnostic]: + diagnostics: list[DataflowDiagnostic] = [] + rows = node.config.get("rows") + if not isinstance(rows, list): + diagnostics.append( + _node_field_error( + node, + "source.rows_required", + "Inline source rows must be a list.", + "rows", + ) + ) + elif len(rows) > 250: + diagnostics.append( + _node_field_error( + node, + "source.row_limit", + "Inline sources are limited to 250 rows.", + "rows", + ) + ) + elif any(not isinstance(row, dict) for row in rows): + diagnostics.append( + _node_field_error( + node, + "source.row_shape", + "Every inline source row must be an object.", + "rows", + ) + ) + if len(json.dumps(node.config, default=str).encode("utf-8")) > 512_000: + diagnostics.append( + _error( + "source.byte_limit", + "Inline source configuration is limited to 512 KiB.", + node_id=node.id, + ) + ) + return diagnostics + + +def _validate_filter(node: GraphNode) -> list[DataflowDiagnostic]: + diagnostics: list[DataflowDiagnostic] = [] + if not _non_empty_text(node.config.get("column")): + diagnostics.append( + _node_field_error(node, "filter.column", "Choose a column.", "column") + ) + operator = node.config.get("operator") + if operator not in FILTER_OPERATORS: + diagnostics.append( + _node_field_error( + node, + "filter.operator", + "Choose a supported comparison operator.", + "operator", + ) + ) + if operator not in {"is_null", "not_null"} and "value" not in node.config: + diagnostics.append( + _node_field_error( + node, + "filter.value", + "Enter a comparison value.", + "value", + ) + ) + return diagnostics + + +def _validate_filter_expression(node: GraphNode) -> list[DataflowDiagnostic]: + if _non_empty_text(node.config.get("expression")): + return [] + return [ + _node_field_error( + node, + "expression.required", + "Enter a filter expression.", + "expression", + ) + ] + + +def _validate_distinct(node: GraphNode) -> list[DataflowDiagnostic]: + columns = node.config.get("columns", []) + if isinstance(columns, list) and all(_non_empty_text(item) for item in columns): + return [] + return [ + _node_field_error( + node, + "distinct.columns", + "Distinct key columns must be named.", + "columns", + ) + ] + + +def _validate_union(node: GraphNode) -> list[DataflowDiagnostic]: + if node.config.get("mode", "all") in {"all", "distinct"}: + return [] + return [ + _node_field_error( + node, + "union.mode", + "Choose whether duplicate rows are kept or removed.", + "mode", + ) + ] + + +def _validate_join(node: GraphNode) -> list[DataflowDiagnostic]: + diagnostics: list[DataflowDiagnostic] = [] + if node.config.get("join_type", "inner") not in JOIN_TYPES: + diagnostics.append( + _node_field_error( + node, + "join.type", + "Choose a supported join type.", + "join_type", + ) + ) + diagnostics.extend( + _validate_key_pair( + node, + left_field="left_keys", + right_field="right_keys", + code_prefix="join", + left_label="left", + right_label="right", + ) + ) + diagnostics.extend( + _validate_identifier_prefix( + node, + field_name="right_prefix", + default="right_", + code="join.right_prefix", + message=( + "Enter an identifier prefix ending in an underscore, " + "such as right_." + ), + ) + ) + return diagnostics + + +def _validate_key_pair( + node: GraphNode, *, - ordered: list[str], + left_field: str, + right_field: str, + code_prefix: str, + left_label: str, + right_label: str, ) -> list[DataflowDiagnostic]: diagnostics: list[DataflowDiagnostic] = [] - node_by_id = {node.id: node for node in graph.nodes} - inputs = graph_inputs_by_port(graph) - schemas: dict[str, _SchemaState] = {} - for node_id in ordered: - node = node_by_id[node_id] - node_inputs = inputs.get(node.id, {}) - input_states = [ - schemas[source_id] - for port_sources in node_inputs.values() - for source_id in port_sources - if source_id in schemas - ] - input_state = input_states[0] if input_states else _SchemaState(frozenset(), open=True) - if node.type == "source.inline": - rows = node.config.get("rows") - schemas[node.id] = _inline_schema(rows) - continue - if node.type == "source.reference": - schemas[node.id] = _configured_source_schema( - node.config.get("source_columns") + left_keys = node.config.get(left_field) + right_keys = node.config.get(right_field) + if not _named_list(left_keys): + diagnostics.append( + _node_field_error( + node, + f"{code_prefix}.{left_field}", + f"Add at least one {left_label} key.", + left_field, ) - continue - if node.type == "combine.union": - if len(input_states) > 1: - closed_shapes = { - state.columns - for state in input_states - if not state.open - } - if len(closed_shapes) > 1: - diagnostics.append( - _warning( - "union.schema_mismatch", - "Appended inputs use different columns; missing values will be null.", - node_id=node.id, - ) - ) - schemas[node.id] = _SchemaState( - frozenset().union(*(state.columns for state in input_states)), - open=any(state.open for state in input_states), - types=_merged_schema_types(input_states), + ) + if not _named_list(right_keys): + diagnostics.append( + _node_field_error( + node, + f"{code_prefix}.{right_field}", + f"Add at least one {right_label} key.", + right_field, ) - continue - if node.type == "combine.join": - left_state = _port_schema(node_inputs, schemas, "left") - right_state = _port_schema(node_inputs, schemas, "right") - _validate_columns( - diagnostics, - node=node, - state=left_state, - columns=node.config.get("left_keys"), - field="left_keys", + ) + if ( + isinstance(left_keys, list) + and isinstance(right_keys, list) + and len(left_keys) != len(right_keys) + ): + diagnostics.append( + _node_field_error( + node, + f"{code_prefix}.key_count", + f"{left_label.title()} and {right_label} keys must have the same length.", + right_field, ) - _validate_columns( - diagnostics, - node=node, - state=right_state, - columns=node.config.get("right_keys"), - field="right_keys", - ) - prefix = str(node.config.get("right_prefix", "right_")) - prefixed_right = {f"{prefix}{column}" for column in right_state.columns} - collisions = left_state.columns & prefixed_right - if collisions: - diagnostics.append( - _error( - "join.output_collision", - f"Join output columns collide: {', '.join(sorted(collisions))}.", - node_id=node.id, - field="right_prefix", - ) - ) - schemas[node.id] = _SchemaState( - frozenset(left_state.columns | prefixed_right), - open=left_state.open or right_state.open, - types={ - **left_state.types, - **{ - f"{prefix}{column}": right_state.type_of(column) - for column in right_state.columns - }, - }, - ) - continue - if node.type == "filter": - _validate_columns( - diagnostics, - node=node, - state=input_state, - columns=[node.config.get("column")], - field="column", - ) - schemas[node.id] = input_state - continue - if node.type == "filter.expression": - parsed = _parsed_node_expression( - diagnostics, - node=node, - field_name="expression", - ) - if parsed is not None: - _validate_columns( - diagnostics, - node=node, - state=input_state, - columns=list(parsed.columns), - field="expression", - ) - schemas[node.id] = input_state - continue - if node.type == "distinct": - _validate_columns( - diagnostics, - node=node, - state=input_state, - columns=node.config.get("columns"), - field="columns", - ) - schemas[node.id] = input_state - continue - if node.type == "select": - fields = node.config.get("fields") - selected_columns: list[str] = [] - output_columns: list[str] = [] - if isinstance(fields, list): - for field in fields: - if isinstance(field, str): - selected_columns.append(field) - output_columns.append(field) - elif isinstance(field, dict): - column = field.get("column") - alias = field.get("alias") or column - if isinstance(column, str): - selected_columns.append(column) - if isinstance(alias, str): - output_columns.append(alias) - _validate_columns( - diagnostics, - node=node, - state=input_state, - columns=selected_columns, - field="fields", - ) - _validate_output_names( - diagnostics, - node=node, - columns=output_columns, - field="fields", - ) - schemas[node.id] = _SchemaState( - frozenset(output_columns), - types={ - output: input_state.type_of(source) - for source, output in zip( - selected_columns, - output_columns, - strict=False, - ) - }, - ) - continue - if node.type == "derive": - _validate_columns( - diagnostics, - node=node, - state=input_state, - columns=node.config.get("source_columns"), - field="source_columns", - ) - target = node.config.get("target_column") - if isinstance(target, str) and target: - if target in input_state.columns: - diagnostics.append( - _warning( - "derive.overwrites_column", - f"Derived column {target!r} replaces an existing value.", - node_id=node.id, - field="target_column", - ) - ) - schemas[node.id] = _SchemaState( - input_state.columns | frozenset((target,)), - open=input_state.open, - types={ - **input_state.types, - target: _derive_result_type( - str(node.config.get("operation") or ""), - [ - input_state.type_of(column) - for column in _text_items( - node.config.get("source_columns") - ) - ], - ), - }, - ) - else: - schemas[node.id] = input_state - continue - if node.type == "expression": - parsed = _parsed_node_expression( - diagnostics, - node=node, - field_name="expression", - ) - target = str(node.config.get("target_column") or "") - result_type = "unknown" - if parsed is not None: - _validate_columns( - diagnostics, - node=node, - state=input_state, - columns=list(parsed.columns), - field="expression", - ) - result_type = infer_expression_type( - parsed, - { - name: input_state.type_of(name) # type: ignore[dict-item] - for name in input_state.columns - }, - ) - expected = str(node.config.get("result_type") or "unknown") - if ( - expected != "unknown" - and result_type not in {"unknown", "null", expected} - ): - diagnostics.append( - _warning( - "expression.type_mismatch", - f"Expression infers {result_type}, not {expected}.", - node_id=node.id, - field="result_type", - ) - ) - schemas[node.id] = _SchemaState( - input_state.columns | ({target} if target else set()), - open=input_state.open, - types={ - **input_state.types, - **( - {target: expected if expected != "unknown" else result_type} - if target - else {} - ), - }, - ) - continue - if node.type in {"convert", "replace"}: - source = str(node.config.get("source_column") or "") - target = str(node.config.get("target_column") or "") - _validate_columns( - diagnostics, - node=node, - state=input_state, - columns=[source], - field="source_column", - ) - target_type = ( - str(node.config.get("target_type") or "unknown") - if node.type == "convert" - else input_state.type_of(source) - ) - schemas[node.id] = _SchemaState( - input_state.columns | ({target} if target else set()), - open=input_state.open, - types={ - **input_state.types, - **({target: target_type} if target else {}), - }, - ) - continue - if node.type == "aggregate": - group_by = _text_items(node.config.get("group_by")) - aggregates = node.config.get("aggregates") - aggregate_columns = [ - str(item.get("column")) - for item in aggregates - if isinstance(aggregates, list) - and isinstance(item, dict) - and item.get("column") not in (None, "", "*") - ] if isinstance(aggregates, list) else [] - _validate_columns( - diagnostics, - node=node, - state=input_state, - columns=[*group_by, *aggregate_columns], - field="aggregates", - ) - aliases = [ - str(item.get("alias")) - for item in aggregates - if isinstance(aggregates, list) - and isinstance(item, dict) - and item.get("alias") - ] if isinstance(aggregates, list) else [] - output_columns = [*group_by, *aliases] - _validate_output_names( - diagnostics, - node=node, - columns=output_columns, - field="aggregates", - ) - aggregate_types = { - str(item.get("alias")): ( - "integer" - if item.get("function") == "count" - else input_state.type_of(str(item.get("column") or "")) - ) - for item in aggregates - if isinstance(aggregates, list) - and isinstance(item, dict) - and item.get("alias") - } if isinstance(aggregates, list) else {} - schemas[node.id] = _SchemaState( - frozenset(output_columns), - types={ - **{ - column: input_state.type_of(column) - for column in group_by - }, - **aggregate_types, - }, - ) - continue - if node.type == "sort": - fields = node.config.get("fields") - columns = [ - str(item.get("column")) - for item in fields - if isinstance(fields, list) - and isinstance(item, dict) - and item.get("column") - ] if isinstance(fields, list) else [] - _validate_columns( - diagnostics, - node=node, - state=input_state, - columns=columns, - field="fields", - ) - schemas[node.id] = input_state - continue - if node.type == "quality.rules": - rules = node.config.get("rules") - _validate_columns( - diagnostics, - node=node, - state=input_state, - columns=[ - str(rule.get("column") or "") - for rule in rules - if isinstance(rules, list) and isinstance(rule, dict) - ] if isinstance(rules, list) else [], - field="rules", - ) - if node.config.get("action", "annotate") == "annotate": - schemas[node.id] = _SchemaState( - input_state.columns - | frozenset(("_quality_valid", "_quality_errors")), - open=input_state.open, - types={ - **input_state.types, - "_quality_valid": "boolean", - "_quality_errors": "array", - }, - ) - else: - schemas[node.id] = input_state - continue - if node.type == "reconcile.compare": - left_state = _port_schema(node_inputs, schemas, "left") - right_state = _port_schema(node_inputs, schemas, "right") - _validate_columns( - diagnostics, - node=node, - state=left_state, - columns=node.config.get("left_keys"), - field="left_keys", - ) - _validate_columns( - diagnostics, - node=node, - state=right_state, - columns=node.config.get("right_keys"), - field="right_keys", - ) - comparison_columns = node.config.get("compare_columns") - if isinstance(comparison_columns, list): - left_columns: list[str] = [] - right_columns: list[str] = [] - for item in comparison_columns: - if isinstance(item, str): - left_columns.append(item) - right_columns.append(item) - elif isinstance(item, dict): - left_columns.append( - str(item.get("left") or item.get("column") or "") - ) - right_columns.append( - str( - item.get("right") - or item.get("left") - or item.get("column") - or "" - ) - ) - _validate_columns( - diagnostics, - node=node, - state=left_state, - columns=left_columns, - field="compare_columns", - ) - _validate_columns( - diagnostics, - node=node, - state=right_state, - columns=right_columns, - field="compare_columns", - ) - prefix = str(node.config.get("right_prefix") or "observed_") - schemas[node.id] = _SchemaState( - left_state.columns - | frozenset(f"{prefix}{column}" for column in right_state.columns) - | frozenset( - ( - "_reconciliation_status", - "_reconciliation_differences", - ) - ), - open=left_state.open or right_state.open, - types={ - **left_state.types, - **{ - f"{prefix}{column}": right_state.type_of(column) - for column in right_state.columns - }, - "_reconciliation_status": "string", - "_reconciliation_differences": "array", - }, - ) - continue - if node.type == "subflow": - output_schema = _configured_source_schema( - node.config.get("output_schema") - ) - schemas[node.id] = ( - output_schema - if output_schema.columns - else _SchemaState(frozenset(), open=True) - ) - continue - schemas[node.id] = input_state + ) return diagnostics -def _configured_source_schema(value: object) -> _SchemaState: - if not isinstance(value, list): - return _SchemaState(frozenset(), open=True) - columns = { - item if isinstance(item, str) else str(item.get("name")) - for item in value - if ( - isinstance(item, str) and item - or isinstance(item, dict) and item.get("name") - ) - } - types = { - str(item.get("name")): str( - item.get("data_type") or item.get("type") or "unknown" - ) - for item in value - if isinstance(item, dict) and item.get("name") - } - return _SchemaState( - frozenset(columns), - open=not columns, - types=types, - ) - - -def _inline_schema(value: object) -> _SchemaState: - if not isinstance(value, list): - return _SchemaState(frozenset(), open=True) - columns = { - str(column) - for row in value - if isinstance(row, dict) - for column in row - } - types: dict[str, str] = {} - for column in columns: - observed = { - _schema_value_type(row.get(column)) - for row in value - if isinstance(row, dict) and row.get(column) is not None - } - types[column] = observed.pop() if len(observed) == 1 else "unknown" - return _SchemaState( - frozenset(columns), - open=not columns, - types=types, - ) - - -def _schema_value_type(value: object) -> str: - if isinstance(value, bool): - return "boolean" - if isinstance(value, int): - return "integer" - if isinstance(value, float): - return "number" - if isinstance(value, str): - return "string" - if isinstance(value, list): - return "array" - if isinstance(value, dict): - return "object" - return "unknown" - - -def _merged_schema_types(states: list[_SchemaState]) -> dict[str, str]: - columns = frozenset().union(*(state.columns for state in states)) - types: dict[str, str] = {} - for column in columns: - observed = { - state.type_of(column) - for state in states - if column in state.columns and state.type_of(column) != "unknown" - } - types[column] = observed.pop() if len(observed) == 1 else "unknown" - return types - - -def _derive_result_type(operation: str, source_types: list[str]) -> str: - if operation in {"upper", "lower", "trim", "concat"}: - return "string" - if operation in {"add", "subtract", "multiply", "divide"}: - return "number" if "number" in source_types or operation == "divide" else "integer" - if operation in {"copy", "coalesce"}: - concrete = { - item for item in source_types if item not in {"unknown", "null"} - } - return concrete.pop() if len(concrete) == 1 else "unknown" - return "unknown" - - -def _parsed_node_expression( - diagnostics: list[DataflowDiagnostic], - *, +def _validate_identifier_prefix( node: GraphNode, + *, field_name: str, -): - try: - return parse_expression(str(node.config.get(field_name) or "")) - except ExpressionError as exc: - diagnostics.append( - _error( - "expression.invalid", - str(exc), - node_id=node.id, - field=field_name, - ) - ) - return None - - -def _port_schema( - inputs: dict[str, list[str]], - schemas: dict[str, _SchemaState], - port: str, -) -> _SchemaState: - source_ids = inputs.get(port, []) - return schemas.get(source_ids[0], _SchemaState(frozenset(), open=True)) if source_ids else _SchemaState(frozenset(), open=True) - - -def _validate_columns( - diagnostics: list[DataflowDiagnostic], - *, - node: GraphNode, - state: _SchemaState, - columns: object, - field: str, -) -> None: - for column in _text_items(columns): - if not state.knows(column): - diagnostics.append( - _error( - "schema.unknown_column", - f"Column {column!r} is not available at this node.", - node_id=node.id, - field=field, - ) - ) - - -def _validate_output_names( - diagnostics: list[DataflowDiagnostic], - *, - node: GraphNode, - columns: list[str], - field: str, -) -> None: - duplicates = sorted({column for column in columns if columns.count(column) > 1}) - if duplicates: - diagnostics.append( - _error( - "schema.duplicate_output", - f"Output column names must be unique: {', '.join(duplicates)}.", - node_id=node.id, - field=field, - ) - ) - - -def _text_items(value: object) -> list[str]: - if not isinstance(value, list): + default: str, + code: str, + message: str, +) -> list[DataflowDiagnostic]: + prefix = node.config.get(field_name, default) + if _non_empty_text(prefix) and re.fullmatch( + r"[A-Za-z_][A-Za-z0-9_]*_", + str(prefix), + ): return [] - return [item for item in value if isinstance(item, str) and item] + return [_node_field_error(node, code, message, field_name)] -def _validate_node_config(node: GraphNode) -> list[DataflowDiagnostic]: - config = node.config +def _validate_select(node: GraphNode) -> list[DataflowDiagnostic]: + fields = node.config.get("fields") + if not isinstance(fields, list) or not fields: + return [ + _node_field_error( + node, + "select.fields", + "Select at least one output column.", + "fields", + ) + ] + invalid = any( + not _non_empty_text( + field + if isinstance(field, str) + else field.get("column") + if isinstance(field, dict) + else None + ) + for field in fields + ) + if not invalid: + return [] + return [ + _node_field_error( + node, + "select.field", + "Selected fields need a source column.", + "fields", + ) + ] + + +def _validate_aggregate(node: GraphNode) -> list[DataflowDiagnostic]: diagnostics: list[DataflowDiagnostic] = [] - if node.type in {"source.inline", "source.reference"}: - source_name = config.get("source_name") - if not isinstance(source_name, str) or not source_name.strip(): - diagnostics.append( - _error( - "source.name_required", - "A source needs a logical SQL name.", - node_id=node.id, - field="source_name", - ) + group_by = node.config.get("group_by", []) + if not isinstance(group_by, list) or any( + not _non_empty_text(item) + for item in group_by + ): + diagnostics.append( + _node_field_error( + node, + "aggregate.group_by", + "Group-by columns must be named.", + "group_by", ) - elif re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", source_name.strip()) is None: - diagnostics.append( - _error( - "source.name_invalid", - "Logical source names must be SQL identifiers, such as monthly_cases.", - node_id=node.id, - field="source_name", - ) + ) + aggregates = node.config.get("aggregates") + if not isinstance(aggregates, list) or not aggregates: + diagnostics.append( + _node_field_error( + node, + "aggregate.required", + "Add at least one aggregate.", + "aggregates", ) - if node.type == "source.reference": - if not _non_empty_text(config.get("source_ref")): - diagnostics.append( - _error( - "source.reference_required", - "Choose a datasource.", - node_id=node.id, - field="source_ref", - ) - ) - expected_fingerprint = config.get("expected_fingerprint") - if expected_fingerprint is not None and not isinstance(expected_fingerprint, str): - diagnostics.append( - _error( - "source.fingerprint", - "The expected source fingerprint must be text.", - node_id=node.id, - field="expected_fingerprint", - ) - ) - return diagnostics - rows = config.get("rows") - if not isinstance(rows, list): - diagnostics.append( - _error("source.rows_required", "Inline source rows must be a list.", node_id=node.id, field="rows") - ) - elif len(rows) > 250: - diagnostics.append( - _error("source.row_limit", "Inline sources are limited to 250 rows.", node_id=node.id, field="rows") - ) - elif any(not isinstance(row, dict) for row in rows): - diagnostics.append( - _error("source.row_shape", "Every inline source row must be an object.", node_id=node.id, field="rows") - ) - if len(json.dumps(config, default=str).encode("utf-8")) > 512_000: - diagnostics.append( - _error("source.byte_limit", "Inline source configuration is limited to 512 KiB.", node_id=node.id) - ) - elif node.type == "filter": - if not _non_empty_text(config.get("column")): - diagnostics.append(_node_field_error(node, "filter.column", "Choose a column.", "column")) - operator = config.get("operator") - if operator not in FILTER_OPERATORS: - diagnostics.append( - _node_field_error(node, "filter.operator", "Choose a supported comparison operator.", "operator") - ) - if operator not in {"is_null", "not_null"} and "value" not in config: - diagnostics.append(_node_field_error(node, "filter.value", "Enter a comparison value.", "value")) - elif node.type == "filter.expression": - if not _non_empty_text(config.get("expression")): - diagnostics.append( - _node_field_error( - node, - "expression.required", - "Enter a filter expression.", - "expression", - ) - ) - elif node.type == "distinct": - columns = config.get("columns", []) - if not isinstance(columns, list) or any(not _non_empty_text(item) for item in columns): - diagnostics.append( - _node_field_error( - node, - "distinct.columns", - "Distinct key columns must be named.", - "columns", - ) - ) - elif node.type == "combine.union": - if config.get("mode", "all") not in {"all", "distinct"}: - diagnostics.append( - _node_field_error( - node, - "union.mode", - "Choose whether duplicate rows are kept or removed.", - "mode", - ) - ) - elif node.type == "combine.join": - if config.get("join_type", "inner") not in JOIN_TYPES: - diagnostics.append( - _node_field_error(node, "join.type", "Choose a supported join type.", "join_type") - ) - left_keys = config.get("left_keys") - right_keys = config.get("right_keys") - if ( - not isinstance(left_keys, list) - or not left_keys - or any(not _non_empty_text(item) for item in left_keys) - ): - diagnostics.append( - _node_field_error(node, "join.left_keys", "Add at least one left key.", "left_keys") - ) - if ( - not isinstance(right_keys, list) - or not right_keys - or any(not _non_empty_text(item) for item in right_keys) - ): - diagnostics.append( - _node_field_error(node, "join.right_keys", "Add at least one right key.", "right_keys") - ) - if isinstance(left_keys, list) and isinstance(right_keys, list) and len(left_keys) != len(right_keys): - diagnostics.append( - _node_field_error( - node, - "join.key_count", - "Left and right joins need the same number of keys.", - "right_keys", - ) - ) - right_prefix = config.get("right_prefix", "right_") - if ( - not _non_empty_text(right_prefix) - or re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*_", str(right_prefix)) is None - ): - diagnostics.append( - _node_field_error( - node, - "join.right_prefix", - "Enter an identifier prefix ending in an underscore, such as right_.", - "right_prefix", - ) - ) - elif node.type == "select": - fields = config.get("fields") - if not isinstance(fields, list) or not fields: - diagnostics.append( - _node_field_error(node, "select.fields", "Select at least one output column.", "fields") - ) - else: - for field in fields: - column = field if isinstance(field, str) else field.get("column") if isinstance(field, dict) else None - if not _non_empty_text(column): - diagnostics.append( - _node_field_error(node, "select.field", "Selected fields need a source column.", "fields") - ) - break - elif node.type == "aggregate": - group_by = config.get("group_by", []) - aggregates = config.get("aggregates") - if not isinstance(group_by, list) or any(not _non_empty_text(item) for item in group_by): - diagnostics.append( - _node_field_error(node, "aggregate.group_by", "Group-by columns must be named.", "group_by") - ) - if not isinstance(aggregates, list) or not aggregates: - diagnostics.append( - _node_field_error(node, "aggregate.required", "Add at least one aggregate.", "aggregates") - ) - else: - for aggregate in aggregates: - if not isinstance(aggregate, dict) or aggregate.get("function") not in AGGREGATE_FUNCTIONS: - diagnostics.append( - _node_field_error( - node, - "aggregate.function", - "Choose COUNT, SUM, AVG, MIN, or MAX.", - "aggregates", - ) - ) - break - if aggregate.get("function") != "count" and not _non_empty_text(aggregate.get("column")): - diagnostics.append( - _node_field_error( - node, - "aggregate.column", - "This aggregate needs a source column.", - "aggregates", - ) - ) - break - if not _non_empty_text(aggregate.get("alias")): - diagnostics.append( - _node_field_error(node, "aggregate.alias", "Every aggregate needs an alias.", "aggregates") - ) - break - elif node.type == "derive": - if not _non_empty_text(config.get("target_column")): - diagnostics.append( - _node_field_error( - node, - "derive.target_column", - "Choose an output column.", - "target_column", - ) - ) - operation = config.get("operation") - if operation not in DERIVE_OPERATIONS: - diagnostics.append( - _node_field_error( - node, - "derive.operation", - "Choose a supported derive operation.", - "operation", - ) - ) - source_columns = config.get("source_columns") - if ( - not isinstance(source_columns, list) - or not source_columns - or any(not _non_empty_text(item) for item in source_columns) - ): - diagnostics.append( - _node_field_error( - node, - "derive.source_columns", - "Choose at least one source column.", - "source_columns", - ) - ) - if operation in {"copy", "upper", "lower", "trim"} and isinstance(source_columns, list) and len(source_columns) != 1: - diagnostics.append( - _node_field_error( - node, - "derive.source_count", - "This operation requires exactly one source column.", - "source_columns", - ) - ) - if operation in {"add", "subtract", "multiply", "divide"} and isinstance(source_columns, list) and len(source_columns) != 2: - diagnostics.append( - _node_field_error( - node, - "derive.numeric_source_count", - "Numeric operations require exactly two source columns.", - "source_columns", - ) - ) - elif node.type == "expression": - if not _non_empty_text(config.get("target_column")): - diagnostics.append( - _node_field_error( - node, - "expression.target_column", - "Choose an output column.", - "target_column", - ) - ) - if not _non_empty_text(config.get("expression")): - diagnostics.append( - _node_field_error( - node, - "expression.required", - "Enter an expression.", - "expression", - ) - ) - if str(config.get("result_type") or "unknown") not in { - "unknown", - *DATA_TYPES, - }: - diagnostics.append( - _node_field_error( - node, - "expression.result_type", - "Choose a supported expression result type.", - "result_type", - ) - ) - elif node.type == "convert": - for field_name, message in ( - ("source_column", "Choose a source column."), - ("target_column", "Choose an output column."), - ): - if not _non_empty_text(config.get(field_name)): - diagnostics.append( - _node_field_error( - node, - f"convert.{field_name}", - message, - field_name, - ) - ) - if config.get("target_type") not in DATA_TYPES: - diagnostics.append( - _node_field_error( - node, - "convert.target_type", - "Choose a supported conversion type.", - "target_type", - ) - ) - if config.get("on_error", "fail") not in {"fail", "null", "keep"}: - diagnostics.append( - _node_field_error( - node, - "convert.on_error", - "Choose how conversion errors are handled.", - "on_error", - ) - ) - elif node.type == "replace": - for field_name, message in ( - ("source_column", "Choose a source column."), - ("target_column", "Choose an output column."), - ): - if not _non_empty_text(config.get(field_name)): - diagnostics.append( - _node_field_error( - node, - f"replace.{field_name}", - message, - field_name, - ) - ) - if config.get("mode", "exact") not in {"exact", "text"}: - diagnostics.append( - _node_field_error( - node, - "replace.mode", - "Choose exact-value or text replacement.", - "mode", - ) - ) - elif node.type == "sort": - fields = config.get("fields") - if not isinstance(fields, list) or not fields: - diagnostics.append(_node_field_error(node, "sort.fields", "Add at least one sort field.", "fields")) - elif any( - not isinstance(item, dict) - or not _non_empty_text(item.get("column")) - or item.get("direction", "asc") not in {"asc", "desc"} - for item in fields - ): - diagnostics.append( - _node_field_error(node, "sort.field", "Sort fields need a column and direction.", "fields") - ) - elif node.type == "limit": - count = config.get("count") - if not isinstance(count, int) or isinstance(count, bool) or not 1 <= count <= 100_000: - diagnostics.append( - _node_field_error(node, "limit.count", "Limit must be between 1 and 100,000.", "count") - ) - elif node.type == "quality.rules": - rules = config.get("rules") - if not isinstance(rules, list) or not rules or len(rules) > 100: - diagnostics.append( - _node_field_error( - node, - "quality.rules", - "Add between one and 100 quality rules.", - "rules", - ) - ) - else: - rule_ids: list[str] = [] - for rule in rules: - if not isinstance(rule, dict): - diagnostics.append( - _node_field_error( - node, - "quality.rule_shape", - "Every quality rule must be an object.", - "rules", - ) - ) - break - rule_id = str(rule.get("id") or "").strip() - rule_ids.append(rule_id) - operator = rule.get("operator") - if not rule_id or not _non_empty_text(rule.get("column")): - diagnostics.append( - _node_field_error( - node, - "quality.rule_identity", - "Quality rules need an ID and column.", - "rules", - ) - ) - break - if operator not in QUALITY_OPERATORS: - diagnostics.append( - _node_field_error( - node, - "quality.operator", - "Choose a supported quality operator.", - "rules", - ) - ) - break - if operator in {"type", "min", "max"} and "value" not in rule: - diagnostics.append( - _node_field_error( - node, - "quality.value", - "This quality rule requires a value.", - "rules", - ) - ) - break - if operator == "allowed" and not isinstance(rule.get("values"), list): - diagnostics.append( - _node_field_error( - node, - "quality.values", - "Allowed-value rules require a list of values.", - "rules", - ) - ) - break - duplicates = sorted( - {rule_id for rule_id in rule_ids if rule_ids.count(rule_id) > 1} - ) - if duplicates: - diagnostics.append( - _node_field_error( - node, - "quality.duplicate_id", - f"Quality rule IDs must be unique: {', '.join(duplicates)}.", - "rules", - ) - ) - if config.get("action", "annotate") not in {"annotate", "drop", "fail"}: - diagnostics.append( - _node_field_error( - node, - "quality.action", - "Choose how invalid rows are handled.", - "action", - ) - ) - elif node.type == "reconcile.compare": - left_keys = config.get("left_keys") - right_keys = config.get("right_keys") - if ( - not isinstance(left_keys, list) - or not left_keys - or any(not _non_empty_text(item) for item in left_keys) - ): - diagnostics.append( - _node_field_error( - node, - "reconcile.left_keys", - "Add at least one expected-table key.", - "left_keys", - ) - ) - if ( - not isinstance(right_keys, list) - or not right_keys - or any(not _non_empty_text(item) for item in right_keys) - ): - diagnostics.append( - _node_field_error( - node, - "reconcile.right_keys", - "Add at least one observed-table key.", - "right_keys", - ) - ) - if ( - isinstance(left_keys, list) - and isinstance(right_keys, list) - and len(left_keys) != len(right_keys) - ): - diagnostics.append( - _node_field_error( - node, - "reconcile.key_count", - "Expected and observed keys must have the same length.", - "right_keys", - ) - ) - prefix = config.get("right_prefix", "observed_") - if ( - not _non_empty_text(prefix) - or re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*_", str(prefix)) is None - ): - diagnostics.append( - _node_field_error( - node, - "reconcile.right_prefix", - "Enter an identifier prefix ending in an underscore.", - "right_prefix", - ) - ) - elif node.type == "subflow": - if not _non_empty_text(config.get("template_ref")): - diagnostics.append( - _node_field_error( - node, - "subflow.template_ref", - "Choose a reusable template.", - "template_ref", - ) - ) - if not _non_empty_text(config.get("template_version")): - diagnostics.append( - _node_field_error( - node, - "subflow.template_version", - "Pin a template version.", - "template_version", - ) - ) - parameters = config.get("parameters") - if not isinstance(parameters, dict) or len(parameters) > 100: - diagnostics.append( - _node_field_error( - node, - "subflow.parameters", - "Subflow parameters must be an object with at most 100 entries.", - "parameters", - ) - ) - try: - nested_payload = substitute_parameters( - config.get("graph"), - parameters if isinstance(parameters, dict) else {}, - ) - nested = PipelineGraph.model_validate(nested_payload) - except (SubflowParameterError, ValueError) as exc: - diagnostics.append( - _node_field_error( - node, - "subflow.graph", - f"Pinned subflow graph is invalid: {exc}", - "graph", - ) - ) - else: - bound_inputs = [ - item - for item in nested.nodes - if item.type == "source.inline" - and item.config.get("input_binding") is True - ] - if len(bound_inputs) != 1: - diagnostics.append( - _node_field_error( - node, - "subflow.input_binding", - "Pinned subflows need exactly one inline input binding.", - "graph", - ) - ) - for nested_diagnostic in validate_graph(nested): - if nested_diagnostic.severity == "error": - diagnostics.append( - _node_field_error( - node, - "subflow.graph_validation", - f"Pinned subflow: {nested_diagnostic.message}", - "graph", - ) - ) - break + ) + return diagnostics + invalid = next( + ( + diagnostic + for aggregate in aggregates + if ( + diagnostic := _validate_aggregate_item(node, aggregate) + ) is not None + ), + None, + ) + if invalid is not None: + diagnostics.append(invalid) return diagnostics +def _validate_aggregate_item( + node: GraphNode, + aggregate: object, +) -> DataflowDiagnostic | None: + if ( + not isinstance(aggregate, dict) + or aggregate.get("function") not in AGGREGATE_FUNCTIONS + ): + return _node_field_error( + node, + "aggregate.function", + "Choose COUNT, SUM, AVG, MIN, or MAX.", + "aggregates", + ) + if ( + aggregate.get("function") != "count" + and not _non_empty_text(aggregate.get("column")) + ): + return _node_field_error( + node, + "aggregate.column", + "This aggregate needs a source column.", + "aggregates", + ) + if not _non_empty_text(aggregate.get("alias")): + return _node_field_error( + node, + "aggregate.alias", + "Every aggregate needs an alias.", + "aggregates", + ) + return None + + +def _validate_derive(node: GraphNode) -> list[DataflowDiagnostic]: + diagnostics: list[DataflowDiagnostic] = [] + if not _non_empty_text(node.config.get("target_column")): + diagnostics.append( + _node_field_error( + node, + "derive.target_column", + "Choose an output column.", + "target_column", + ) + ) + operation = node.config.get("operation") + if operation not in DERIVE_OPERATIONS: + diagnostics.append( + _node_field_error( + node, + "derive.operation", + "Choose a supported derive operation.", + "operation", + ) + ) + source_columns = node.config.get("source_columns") + if not _named_list(source_columns): + diagnostics.append( + _node_field_error( + node, + "derive.source_columns", + "Choose at least one source column.", + "source_columns", + ) + ) + return diagnostics + expected_count = ( + 1 + if operation in {"copy", "upper", "lower", "trim"} + else 2 + if operation in {"add", "subtract", "multiply", "divide"} + else None + ) + if expected_count is not None and len(source_columns) != expected_count: + code = ( + "derive.source_count" + if expected_count == 1 + else "derive.numeric_source_count" + ) + message = ( + "This operation requires exactly one source column." + if expected_count == 1 + else "Numeric operations require exactly two source columns." + ) + diagnostics.append( + _node_field_error(node, code, message, "source_columns") + ) + return diagnostics + + +def _validate_expression(node: GraphNode) -> list[DataflowDiagnostic]: + diagnostics: list[DataflowDiagnostic] = [] + for field_name, code, message in ( + ("target_column", "expression.target_column", "Choose an output column."), + ("expression", "expression.required", "Enter an expression."), + ): + if not _non_empty_text(node.config.get(field_name)): + diagnostics.append( + _node_field_error(node, code, message, field_name) + ) + if str(node.config.get("result_type") or "unknown") not in { + "unknown", + *DATA_TYPES, + }: + diagnostics.append( + _node_field_error( + node, + "expression.result_type", + "Choose a supported expression result type.", + "result_type", + ) + ) + return diagnostics + + +def _validate_convert(node: GraphNode) -> list[DataflowDiagnostic]: + diagnostics = _validate_source_target_columns(node, "convert") + if node.config.get("target_type") not in DATA_TYPES: + diagnostics.append( + _node_field_error( + node, + "convert.target_type", + "Choose a supported conversion type.", + "target_type", + ) + ) + if node.config.get("on_error", "fail") not in {"fail", "null", "keep"}: + diagnostics.append( + _node_field_error( + node, + "convert.on_error", + "Choose how conversion errors are handled.", + "on_error", + ) + ) + return diagnostics + + +def _validate_replace(node: GraphNode) -> list[DataflowDiagnostic]: + diagnostics = _validate_source_target_columns(node, "replace") + if node.config.get("mode", "exact") not in {"exact", "text"}: + diagnostics.append( + _node_field_error( + node, + "replace.mode", + "Choose exact-value or text replacement.", + "mode", + ) + ) + return diagnostics + + +def _validate_source_target_columns( + node: GraphNode, + code_prefix: str, +) -> list[DataflowDiagnostic]: + diagnostics: list[DataflowDiagnostic] = [] + for field_name, message in ( + ("source_column", "Choose a source column."), + ("target_column", "Choose an output column."), + ): + if not _non_empty_text(node.config.get(field_name)): + diagnostics.append( + _node_field_error( + node, + f"{code_prefix}.{field_name}", + message, + field_name, + ) + ) + return diagnostics + + +def _validate_sort(node: GraphNode) -> list[DataflowDiagnostic]: + fields = node.config.get("fields") + if not isinstance(fields, list) or not fields: + return [ + _node_field_error( + node, + "sort.fields", + "Add at least one sort field.", + "fields", + ) + ] + valid = all( + isinstance(item, dict) + and _non_empty_text(item.get("column")) + and item.get("direction", "asc") in {"asc", "desc"} + for item in fields + ) + if valid: + return [] + return [ + _node_field_error( + node, + "sort.field", + "Sort fields need a column and direction.", + "fields", + ) + ] + + +def _validate_limit(node: GraphNode) -> list[DataflowDiagnostic]: + count = node.config.get("count") + if ( + isinstance(count, int) + and not isinstance(count, bool) + and 1 <= count <= 100_000 + ): + return [] + return [ + _node_field_error( + node, + "limit.count", + "Limit must be between 1 and 100,000.", + "count", + ) + ] + + +def _validate_quality(node: GraphNode) -> list[DataflowDiagnostic]: + diagnostics: list[DataflowDiagnostic] = [] + rules = node.config.get("rules") + if not isinstance(rules, list) or not 1 <= len(rules) <= 100: + diagnostics.append( + _node_field_error( + node, + "quality.rules", + "Add between one and 100 quality rules.", + "rules", + ) + ) + else: + diagnostics.extend(_validate_quality_rules(node, rules)) + if node.config.get("action", "annotate") not in {"annotate", "drop", "fail"}: + diagnostics.append( + _node_field_error( + node, + "quality.action", + "Choose how invalid rows are handled.", + "action", + ) + ) + return diagnostics + + +def _validate_quality_rules( + node: GraphNode, + rules: list[object], +) -> list[DataflowDiagnostic]: + rule_ids: list[str] = [] + for rule in rules: + diagnostic, rule_id = _validate_quality_rule(node, rule) + if diagnostic is not None: + return [diagnostic] + rule_ids.append(rule_id) + duplicates = sorted( + rule_id + for rule_id in set(rule_ids) + if rule_ids.count(rule_id) > 1 + ) + if not duplicates: + return [] + return [ + _node_field_error( + node, + "quality.duplicate_id", + f"Quality rule IDs must be unique: {', '.join(duplicates)}.", + "rules", + ) + ] + + +def _validate_quality_rule( + node: GraphNode, + rule: object, +) -> tuple[DataflowDiagnostic | None, str]: + if not isinstance(rule, dict): + return ( + _node_field_error( + node, + "quality.rule_shape", + "Every quality rule must be an object.", + "rules", + ), + "", + ) + rule_id = str(rule.get("id") or "").strip() + operator = rule.get("operator") + if not rule_id or not _non_empty_text(rule.get("column")): + return ( + _node_field_error( + node, + "quality.rule_identity", + "Quality rules need an ID and column.", + "rules", + ), + rule_id, + ) + if operator not in QUALITY_OPERATORS: + return ( + _node_field_error( + node, + "quality.operator", + "Choose a supported quality operator.", + "rules", + ), + rule_id, + ) + if operator in {"type", "min", "max"} and "value" not in rule: + return ( + _node_field_error( + node, + "quality.value", + "This quality rule requires a value.", + "rules", + ), + rule_id, + ) + if operator == "allowed" and not isinstance(rule.get("values"), list): + return ( + _node_field_error( + node, + "quality.values", + "Allowed-value rules require a list of values.", + "rules", + ), + rule_id, + ) + return None, rule_id + + +def _validate_reconcile(node: GraphNode) -> list[DataflowDiagnostic]: + diagnostics = _validate_key_pair( + node, + left_field="left_keys", + right_field="right_keys", + code_prefix="reconcile", + left_label="expected-table", + right_label="observed-table", + ) + diagnostics.extend( + _validate_identifier_prefix( + node, + field_name="right_prefix", + default="observed_", + code="reconcile.right_prefix", + message="Enter an identifier prefix ending in an underscore.", + ) + ) + return diagnostics + + +def _validate_subflow(node: GraphNode) -> list[DataflowDiagnostic]: + diagnostics: list[DataflowDiagnostic] = [] + for field_name, code, message in ( + ("template_ref", "subflow.template_ref", "Choose a reusable template."), + ("template_version", "subflow.template_version", "Pin a template version."), + ): + if not _non_empty_text(node.config.get(field_name)): + diagnostics.append( + _node_field_error(node, code, message, field_name) + ) + parameters = node.config.get("parameters") + if not isinstance(parameters, dict) or len(parameters) > 100: + diagnostics.append( + _node_field_error( + node, + "subflow.parameters", + "Subflow parameters must be an object with at most 100 entries.", + "parameters", + ) + ) + nested, diagnostic = _validated_subflow_graph(node, parameters) + if diagnostic is not None: + diagnostics.append(diagnostic) + elif nested is not None: + diagnostics.extend(_validate_subflow_binding(node, nested)) + diagnostics.extend(_validate_nested_subflow(node, nested)) + return diagnostics + + +def _validated_subflow_graph( + node: GraphNode, + parameters: object, +) -> tuple[PipelineGraph | None, DataflowDiagnostic | None]: + try: + payload = substitute_parameters( + node.config.get("graph"), + parameters if isinstance(parameters, dict) else {}, + ) + return PipelineGraph.model_validate(payload), None + except (SubflowParameterError, ValueError) as exc: + return None, _node_field_error( + node, + "subflow.graph", + f"Pinned subflow graph is invalid: {exc}", + "graph", + ) + + +def _validate_subflow_binding( + node: GraphNode, + nested: PipelineGraph, +) -> list[DataflowDiagnostic]: + binding_count = sum( + item.type == "source.inline" + and item.config.get("input_binding") is True + for item in nested.nodes + ) + if binding_count == 1: + return [] + return [ + _node_field_error( + node, + "subflow.input_binding", + "Pinned subflows need exactly one inline input binding.", + "graph", + ) + ] + + +def _validate_nested_subflow( + node: GraphNode, + nested: PipelineGraph, +) -> list[DataflowDiagnostic]: + error = next( + ( + diagnostic + for diagnostic in validate_graph(nested) + if diagnostic.severity == "error" + ), + None, + ) + if error is None: + return [] + return [ + _node_field_error( + node, + "subflow.graph_validation", + f"Pinned subflow: {error.message}", + "graph", + ) + ] + + +def _validate_no_config(_node: GraphNode) -> list[DataflowDiagnostic]: + return [] + + +def _named_list(value: object) -> bool: + return ( + isinstance(value, list) + and bool(value) + and all(_non_empty_text(item) for item in value) + ) + + def _non_empty_text(value: object) -> bool: return isinstance(value, str) and bool(value.strip()) @@ -1618,15 +1285,37 @@ def _warning( def _register_config_validators() -> None: - for node_type in SUPPORTED_NODE_TYPES: + validators = { + "source.inline": _validate_source, + "source.reference": _validate_source, + "combine.union": _validate_union, + "combine.join": _validate_join, + "filter": _validate_filter, + "filter.expression": _validate_filter_expression, + "distinct": _validate_distinct, + "select": _validate_select, + "derive": _validate_derive, + "expression": _validate_expression, + "convert": _validate_convert, + "replace": _validate_replace, + "aggregate": _validate_aggregate, + "sort": _validate_sort, + "limit": _validate_limit, + "quality.rules": _validate_quality, + "reconcile.compare": _validate_reconcile, + "subflow": _validate_subflow, + "output": _validate_no_config, + } + for node_type, validator in validators.items(): if OPERATOR_REGISTRY.config_validator(node_type) is None: OPERATOR_REGISTRY.register_config_validator( node_type, - _validate_node_config, + validator, ) _register_config_validators() +register_schema_propagators() __all__ = [ diff --git a/src/govoplan_dataflow/backend/schema_validation.py b/src/govoplan_dataflow/backend/schema_validation.py new file mode 100644 index 0000000..632cbcf --- /dev/null +++ b/src/govoplan_dataflow/backend/schema_validation.py @@ -0,0 +1,892 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from govoplan_dataflow.backend.expressions import ( + ExpressionError, + infer_expression_type, + parse_expression, +) +from govoplan_dataflow.backend.operator_registry import OPERATOR_REGISTRY +from govoplan_dataflow.backend.schemas import ( + DataflowDiagnostic, + GraphNode, + PipelineGraph, +) + + +@dataclass(frozen=True, slots=True) +class SchemaState: + columns: frozenset[str] + open: bool = False + types: dict[str, str] = field(default_factory=dict) + + def knows(self, column: str) -> bool: + return self.open or column in self.columns + + def type_of(self, column: str) -> str: + return self.types.get(column, "unknown") + + +@dataclass(frozen=True, slots=True) +class SchemaPropagationContext: + node: GraphNode + input_state: SchemaState + input_states: tuple[SchemaState, ...] + inputs_by_port: dict[str, list[str]] + schemas: dict[str, SchemaState] + + def port_state(self, port: str) -> SchemaState: + source_ids = self.inputs_by_port.get(port, ()) + if not source_ids: + return unknown_schema() + return self.schemas.get(source_ids[0], unknown_schema()) + + +@dataclass(frozen=True, slots=True) +class SchemaPropagationResult: + state: SchemaState + diagnostics: tuple[DataflowDiagnostic, ...] = () + + +def validate_graph_schemas( + graph: PipelineGraph, + *, + ordered: list[str], +) -> list[DataflowDiagnostic]: + node_by_id = {node.id: node for node in graph.nodes} + inputs = _graph_inputs_by_port(graph) + schemas: dict[str, SchemaState] = {} + diagnostics: list[DataflowDiagnostic] = [] + for node_id in ordered: + node = node_by_id[node_id] + context = _propagation_context(node, inputs, schemas) + propagator = OPERATOR_REGISTRY.schema_propagator(node.type) + if propagator is None: + diagnostics.append( + _error( + "node.schema_propagator_missing", + f"Node type {node.type!r} has no schema propagator.", + node_id=node.id, + ) + ) + schemas[node.id] = context.input_state + continue + result = propagator(context) + if not isinstance(result, SchemaPropagationResult): + raise TypeError( + f"Schema propagator for {node.type!r} returned " + f"{type(result).__name__}, not SchemaPropagationResult." + ) + schemas[node.id] = result.state + diagnostics.extend(result.diagnostics) + return diagnostics + + +def _graph_inputs_by_port( + graph: PipelineGraph, +) -> dict[str, dict[str, list[str]]]: + result: dict[str, dict[str, list[str]]] = {} + for edge in graph.edges: + result.setdefault(edge.target, {}).setdefault( + edge.target_port, + [], + ).append(edge.source) + return result + + +def _propagation_context( + node: GraphNode, + inputs: dict[str, dict[str, list[str]]], + schemas: dict[str, SchemaState], +) -> SchemaPropagationContext: + node_inputs = inputs.get(node.id, {}) + input_states = tuple( + schemas[source_id] + for port_sources in node_inputs.values() + for source_id in port_sources + if source_id in schemas + ) + return SchemaPropagationContext( + node=node, + input_state=input_states[0] if input_states else unknown_schema(), + input_states=input_states, + inputs_by_port=node_inputs, + schemas=schemas, + ) + + +def _inline_source( + context: SchemaPropagationContext, +) -> SchemaPropagationResult: + return SchemaPropagationResult(_inline_schema(context.node.config.get("rows"))) + + +def _reference_source( + context: SchemaPropagationContext, +) -> SchemaPropagationResult: + return SchemaPropagationResult( + _configured_schema(context.node.config.get("source_columns")) + ) + + +def _union(context: SchemaPropagationContext) -> SchemaPropagationResult: + diagnostics: list[DataflowDiagnostic] = [] + closed_shapes = { + state.columns + for state in context.input_states + if not state.open + } + if len(closed_shapes) > 1: + diagnostics.append( + _warning( + "union.schema_mismatch", + "Appended inputs use different columns; missing values will be null.", + node_id=context.node.id, + ) + ) + return SchemaPropagationResult( + SchemaState( + frozenset().union( + *(state.columns for state in context.input_states) + ), + open=any(state.open for state in context.input_states), + types=_merged_schema_types(context.input_states), + ), + tuple(diagnostics), + ) + + +def _join(context: SchemaPropagationContext) -> SchemaPropagationResult: + node = context.node + left_state = context.port_state("left") + right_state = context.port_state("right") + diagnostics = [ + *_unknown_columns( + node, + left_state, + node.config.get("left_keys"), + field_name="left_keys", + ), + *_unknown_columns( + node, + right_state, + node.config.get("right_keys"), + field_name="right_keys", + ), + ] + prefix = str(node.config.get("right_prefix", "right_")) + prefixed_right = { + f"{prefix}{column}" + for column in right_state.columns + } + collisions = left_state.columns & prefixed_right + if collisions: + diagnostics.append( + _error( + "join.output_collision", + f"Join output columns collide: {', '.join(sorted(collisions))}.", + node_id=node.id, + field="right_prefix", + ) + ) + return SchemaPropagationResult( + SchemaState( + frozenset(left_state.columns | prefixed_right), + open=left_state.open or right_state.open, + types={ + **left_state.types, + **{ + f"{prefix}{column}": right_state.type_of(column) + for column in right_state.columns + }, + }, + ), + tuple(diagnostics), + ) + + +def _filter(context: SchemaPropagationContext) -> SchemaPropagationResult: + return _passthrough_with_columns( + context, + [context.node.config.get("column")], + field_name="column", + ) + + +def _filter_expression( + context: SchemaPropagationContext, +) -> SchemaPropagationResult: + parsed, diagnostics = _node_expression(context.node, "expression") + if parsed is not None: + diagnostics.extend( + _unknown_columns( + context.node, + context.input_state, + list(parsed.columns), + field_name="expression", + ) + ) + return SchemaPropagationResult( + context.input_state, + tuple(diagnostics), + ) + + +def _distinct(context: SchemaPropagationContext) -> SchemaPropagationResult: + return _passthrough_with_columns( + context, + context.node.config.get("columns"), + field_name="columns", + ) + + +def _select(context: SchemaPropagationContext) -> SchemaPropagationResult: + selected, output = _selected_columns(context.node.config.get("fields")) + diagnostics = [ + *_unknown_columns( + context.node, + context.input_state, + selected, + field_name="fields", + ), + *_duplicate_outputs( + context.node, + output, + field_name="fields", + ), + ] + return SchemaPropagationResult( + SchemaState( + frozenset(output), + types={ + target: context.input_state.type_of(source) + for source, target in zip(selected, output, strict=False) + }, + ), + tuple(diagnostics), + ) + + +def _selected_columns(value: object) -> tuple[list[str], list[str]]: + selected: list[str] = [] + output: list[str] = [] + if not isinstance(value, list): + return selected, output + for item in value: + if isinstance(item, str): + selected.append(item) + output.append(item) + elif isinstance(item, dict): + source = item.get("column") + target = item.get("alias") or source + if isinstance(source, str): + selected.append(source) + if isinstance(target, str): + output.append(target) + return selected, output + + +def _derive(context: SchemaPropagationContext) -> SchemaPropagationResult: + node = context.node + source_columns = _text_items(node.config.get("source_columns")) + diagnostics = _unknown_columns( + node, + context.input_state, + source_columns, + field_name="source_columns", + ) + target = node.config.get("target_column") + if not isinstance(target, str) or not target: + return SchemaPropagationResult( + context.input_state, + tuple(diagnostics), + ) + if target in context.input_state.columns: + diagnostics.append( + _warning( + "derive.overwrites_column", + f"Derived column {target!r} replaces an existing value.", + node_id=node.id, + field="target_column", + ) + ) + result_type = _derive_result_type( + str(node.config.get("operation") or ""), + [ + context.input_state.type_of(column) + for column in source_columns + ], + ) + return SchemaPropagationResult( + _with_column(context.input_state, target, result_type), + tuple(diagnostics), + ) + + +def _expression(context: SchemaPropagationContext) -> SchemaPropagationResult: + node = context.node + parsed, diagnostics = _node_expression(node, "expression") + inferred = "unknown" + if parsed is not None: + diagnostics.extend( + _unknown_columns( + node, + context.input_state, + list(parsed.columns), + field_name="expression", + ) + ) + inferred = infer_expression_type( + parsed, + { + name: context.input_state.type_of(name) # type: ignore[dict-item] + for name in context.input_state.columns + }, + ) + expected = str(node.config.get("result_type") or "unknown") + if expected != "unknown" and inferred not in {"unknown", "null", expected}: + diagnostics.append( + _warning( + "expression.type_mismatch", + f"Expression infers {inferred}, not {expected}.", + node_id=node.id, + field="result_type", + ) + ) + target = str(node.config.get("target_column") or "") + state = ( + _with_column( + context.input_state, + target, + expected if expected != "unknown" else inferred, + ) + if target + else context.input_state + ) + return SchemaPropagationResult(state, tuple(diagnostics)) + + +def _convert_or_replace( + context: SchemaPropagationContext, +) -> SchemaPropagationResult: + node = context.node + source = str(node.config.get("source_column") or "") + target = str(node.config.get("target_column") or "") + diagnostics = _unknown_columns( + node, + context.input_state, + [source], + field_name="source_column", + ) + target_type = ( + str(node.config.get("target_type") or "unknown") + if node.type == "convert" + else context.input_state.type_of(source) + ) + state = ( + _with_column(context.input_state, target, target_type) + if target + else context.input_state + ) + return SchemaPropagationResult(state, tuple(diagnostics)) + + +def _aggregate(context: SchemaPropagationContext) -> SchemaPropagationResult: + node = context.node + group_by = _text_items(node.config.get("group_by")) + aggregates = _mapping_items(node.config.get("aggregates")) + source_columns = [ + str(item.get("column")) + for item in aggregates + if item.get("column") not in (None, "", "*") + ] + aliases = [ + str(item.get("alias")) + for item in aggregates + if item.get("alias") + ] + output = [*group_by, *aliases] + diagnostics = [ + *_unknown_columns( + node, + context.input_state, + [*group_by, *source_columns], + field_name="aggregates", + ), + *_duplicate_outputs(node, output, field_name="aggregates"), + ] + aggregate_types = { + str(item["alias"]): ( + "integer" + if item.get("function") == "count" + else context.input_state.type_of(str(item.get("column") or "")) + ) + for item in aggregates + if item.get("alias") + } + return SchemaPropagationResult( + SchemaState( + frozenset(output), + types={ + **{ + column: context.input_state.type_of(column) + for column in group_by + }, + **aggregate_types, + }, + ), + tuple(diagnostics), + ) + + +def _sort(context: SchemaPropagationContext) -> SchemaPropagationResult: + columns = [ + str(item.get("column")) + for item in _mapping_items(context.node.config.get("fields")) + if item.get("column") + ] + return _passthrough_with_columns( + context, + columns, + field_name="fields", + ) + + +def _quality(context: SchemaPropagationContext) -> SchemaPropagationResult: + rules = _mapping_items(context.node.config.get("rules")) + diagnostics = _unknown_columns( + context.node, + context.input_state, + [str(rule.get("column") or "") for rule in rules], + field_name="rules", + ) + if context.node.config.get("action", "annotate") != "annotate": + return SchemaPropagationResult( + context.input_state, + tuple(diagnostics), + ) + state = _with_column(context.input_state, "_quality_valid", "boolean") + state = _with_column(state, "_quality_errors", "array") + return SchemaPropagationResult(state, tuple(diagnostics)) + + +def _reconcile(context: SchemaPropagationContext) -> SchemaPropagationResult: + node = context.node + left_state = context.port_state("left") + right_state = context.port_state("right") + left_compare, right_compare = _comparison_columns( + node.config.get("compare_columns") + ) + diagnostics = [ + *_unknown_columns( + node, + left_state, + node.config.get("left_keys"), + field_name="left_keys", + ), + *_unknown_columns( + node, + right_state, + node.config.get("right_keys"), + field_name="right_keys", + ), + *_unknown_columns( + node, + left_state, + left_compare, + field_name="compare_columns", + ), + *_unknown_columns( + node, + right_state, + right_compare, + field_name="compare_columns", + ), + ] + prefix = str(node.config.get("right_prefix") or "observed_") + state = SchemaState( + left_state.columns + | frozenset( + f"{prefix}{column}" + for column in right_state.columns + ) + | frozenset( + ( + "_reconciliation_status", + "_reconciliation_differences", + ) + ), + open=left_state.open or right_state.open, + types={ + **left_state.types, + **{ + f"{prefix}{column}": right_state.type_of(column) + for column in right_state.columns + }, + "_reconciliation_status": "string", + "_reconciliation_differences": "array", + }, + ) + return SchemaPropagationResult(state, tuple(diagnostics)) + + +def _comparison_columns(value: object) -> tuple[list[str], list[str]]: + left: list[str] = [] + right: list[str] = [] + if not isinstance(value, list): + return left, right + for item in value: + if isinstance(item, str): + left.append(item) + right.append(item) + elif isinstance(item, dict): + left_name = str(item.get("left") or item.get("column") or "") + right_name = str( + item.get("right") + or item.get("left") + or item.get("column") + or "" + ) + left.append(left_name) + right.append(right_name) + return left, right + + +def _subflow(context: SchemaPropagationContext) -> SchemaPropagationResult: + output_schema = _configured_schema( + context.node.config.get("output_schema") + ) + return SchemaPropagationResult( + output_schema + if output_schema.columns + else unknown_schema() + ) + + +def _identity(context: SchemaPropagationContext) -> SchemaPropagationResult: + return SchemaPropagationResult(context.input_state) + + +def _passthrough_with_columns( + context: SchemaPropagationContext, + columns: object, + *, + field_name: str, +) -> SchemaPropagationResult: + diagnostics = _unknown_columns( + context.node, + context.input_state, + columns, + field_name=field_name, + ) + return SchemaPropagationResult( + context.input_state, + tuple(diagnostics), + ) + + +def _with_column( + state: SchemaState, + name: str, + data_type: str, +) -> SchemaState: + return SchemaState( + state.columns | frozenset((name,)), + open=state.open, + types={**state.types, name: data_type}, + ) + + +def _configured_schema(value: object) -> SchemaState: + if not isinstance(value, list): + return unknown_schema() + columns = { + item + if isinstance(item, str) + else str(item.get("name")) + for item in value + if ( + isinstance(item, str) + and item + or isinstance(item, dict) + and item.get("name") + ) + } + types = { + str(item["name"]): str( + item.get("data_type") + or item.get("type") + or "unknown" + ) + for item in value + if isinstance(item, dict) and item.get("name") + } + return SchemaState( + frozenset(columns), + open=not columns, + types=types, + ) + + +def _inline_schema(value: object) -> SchemaState: + if not isinstance(value, list): + return unknown_schema() + rows = [row for row in value if isinstance(row, dict)] + columns = { + str(column) + for row in rows + for column in row + } + types = { + column: _observed_column_type(rows, column) + for column in columns + } + return SchemaState( + frozenset(columns), + open=not columns, + types=types, + ) + + +def _observed_column_type( + rows: list[dict[str, Any]], + column: str, +) -> str: + observed = { + _schema_value_type(row.get(column)) + for row in rows + if row.get(column) is not None + } + return observed.pop() if len(observed) == 1 else "unknown" + + +def _schema_value_type(value: object) -> str: + type_checks = ( + (bool, "boolean"), + (int, "integer"), + (float, "number"), + (str, "string"), + (list, "array"), + (dict, "object"), + ) + return next( + ( + name + for value_type, name in type_checks + if isinstance(value, value_type) + ), + "unknown", + ) + + +def _merged_schema_types( + states: tuple[SchemaState, ...], +) -> dict[str, str]: + columns = frozenset().union(*(state.columns for state in states)) + return { + column: _merged_column_type(states, column) + for column in columns + } + + +def _merged_column_type( + states: tuple[SchemaState, ...], + column: str, +) -> str: + observed = { + state.type_of(column) + for state in states + if ( + column in state.columns + and state.type_of(column) != "unknown" + ) + } + return observed.pop() if len(observed) == 1 else "unknown" + + +def _derive_result_type( + operation: str, + source_types: list[str], +) -> str: + if operation in {"upper", "lower", "trim", "concat"}: + return "string" + if operation in {"add", "subtract", "multiply", "divide"}: + return ( + "number" + if "number" in source_types or operation == "divide" + else "integer" + ) + if operation in {"copy", "coalesce"}: + concrete = { + item + for item in source_types + if item not in {"unknown", "null"} + } + return ( + concrete.pop() + if len(concrete) == 1 + else "unknown" + ) + return "unknown" + + +def _node_expression( + node: GraphNode, + field_name: str, +) -> tuple[object | None, list[DataflowDiagnostic]]: + try: + return ( + parse_expression(str(node.config.get(field_name) or "")), + [], + ) + except ExpressionError as exc: + return ( + None, + [ + _error( + "expression.invalid", + str(exc), + node_id=node.id, + field=field_name, + ) + ], + ) + + +def _unknown_columns( + node: GraphNode, + state: SchemaState, + columns: object, + *, + field_name: str, +) -> list[DataflowDiagnostic]: + return [ + _error( + "schema.unknown_column", + f"Column {column!r} is not available at this node.", + node_id=node.id, + field=field_name, + ) + for column in _text_items(columns) + if not state.knows(column) + ] + + +def _duplicate_outputs( + node: GraphNode, + columns: list[str], + *, + field_name: str, +) -> list[DataflowDiagnostic]: + duplicates = sorted( + column + for column in set(columns) + if columns.count(column) > 1 + ) + if not duplicates: + return [] + return [ + _error( + "schema.duplicate_output", + f"Output column names must be unique: {', '.join(duplicates)}.", + node_id=node.id, + field=field_name, + ) + ] + + +def _text_items(value: object) -> list[str]: + if not isinstance(value, list): + return [] + return [ + item + for item in value + if isinstance(item, str) and item + ] + + +def _mapping_items(value: object) -> list[dict[str, Any]]: + if not isinstance(value, list): + return [] + return [ + item + for item in value + if isinstance(item, dict) + ] + + +def unknown_schema() -> SchemaState: + return SchemaState(frozenset(), open=True) + + +def _error( + code: str, + message: str, + *, + node_id: str | None = None, + field: str | None = None, +) -> DataflowDiagnostic: + return DataflowDiagnostic( + severity="error", + code=code, + message=message, + node_id=node_id, + field=field, + ) + + +def _warning( + code: str, + message: str, + *, + node_id: str | None = None, + field: str | None = None, +) -> DataflowDiagnostic: + return DataflowDiagnostic( + severity="warning", + code=code, + message=message, + node_id=node_id, + field=field, + ) + + +def register_schema_propagators() -> None: + propagators = { + "source.inline": _inline_source, + "source.reference": _reference_source, + "combine.union": _union, + "combine.join": _join, + "filter": _filter, + "filter.expression": _filter_expression, + "distinct": _distinct, + "select": _select, + "derive": _derive, + "expression": _expression, + "convert": _convert_or_replace, + "replace": _convert_or_replace, + "aggregate": _aggregate, + "sort": _sort, + "limit": _identity, + "quality.rules": _quality, + "reconcile.compare": _reconcile, + "subflow": _subflow, + "output": _identity, + } + for node_type, propagator in propagators.items(): + if OPERATOR_REGISTRY.schema_propagator(node_type) is None: + OPERATOR_REGISTRY.register_schema_propagator( + node_type, + propagator, + ) + + +__all__ = [ + "SchemaPropagationContext", + "SchemaPropagationResult", + "SchemaState", + "register_schema_propagators", + "validate_graph_schemas", +] diff --git a/src/govoplan_dataflow/backend/service.py b/src/govoplan_dataflow/backend/service.py index 884ff1f..7ba3a16 100644 --- a/src/govoplan_dataflow/backend/service.py +++ b/src/govoplan_dataflow/backend/service.py @@ -33,6 +33,7 @@ from govoplan_dataflow.backend.db.models import ( from govoplan_dataflow.backend.executor import ( EXECUTOR_VERSION, PipelineExecutionError, + PipelineExecutionResult, ResolvedSource, execute_preview, ) @@ -757,6 +758,55 @@ def start_pipeline_run( registry: object | None, request: DataflowRunRequest, ) -> tuple[DataflowRun, bool]: + pipeline, revision = _run_definition( + session, + tenant_id=tenant_id, + principal=principal, + registry=registry, + request=request, + ) + idempotency_key, request_hash = _validated_run_identity(request) + existing = _existing_pipeline_run( + session, + tenant_id=tenant_id, + pipeline_id=pipeline.id, + idempotency_key=idempotency_key, + request_hash=request_hash, + ) + if existing is not None: + return existing, True + run = _new_pipeline_run( + tenant_id=tenant_id, + actor_id=actor_id, + pipeline=pipeline, + revision=revision, + request=request, + idempotency_key=idempotency_key, + request_hash=request_hash, + ) + session.add(run) + session.flush() + _execute_pipeline_run( + session, + run=run, + pipeline=pipeline, + revision=revision, + request=request, + principal=principal, + registry=registry, + ) + session.flush() + return run, False + + +def _run_definition( + session: Session, + *, + tenant_id: str, + principal: ApiPrincipal, + registry: object | None, + request: DataflowRunRequest, +) -> tuple[DataflowPipeline, DataflowPipelineRevision]: pipeline_id = _strip_ref(request.pipeline_ref, "pipeline:") if not pipeline_id: raise DataflowNotFoundError("Dataflow pipeline not found") @@ -784,6 +834,10 @@ def start_pipeline_run( pipeline=pipeline, revision=request.revision, ) + return pipeline, revision + + +def _validated_run_identity(request: DataflowRunRequest) -> tuple[str, str]: idempotency_key = request.idempotency_key.strip() if not idempotency_key or len(idempotency_key) > 255: raise DataflowConflictError( @@ -793,11 +847,21 @@ def start_pipeline_run( raise DataflowConflictError( "The bounded Dataflow runner supports between 1 and 500 output rows." ) - request_hash = _run_request_hash(request) + return idempotency_key, _run_request_hash(request) + + +def _existing_pipeline_run( + session: Session, + *, + tenant_id: str, + pipeline_id: str, + idempotency_key: str, + request_hash: str, +) -> DataflowRun | None: existing = session.scalar( select(DataflowRun).where( DataflowRun.tenant_id == tenant_id, - DataflowRun.pipeline_id == pipeline.id, + DataflowRun.pipeline_id == pipeline_id, DataflowRun.idempotency_key == idempotency_key, ) ) @@ -807,11 +871,20 @@ def start_pipeline_run( "The Dataflow run idempotency key was already used with " "different parameters." ) - return existing, True + return existing - graph = PipelineGraph.model_validate(revision.graph) - started_at = utcnow() - run = DataflowRun( + +def _new_pipeline_run( + *, + tenant_id: str, + actor_id: str | None, + pipeline: DataflowPipeline, + revision: DataflowPipelineRevision, + request: DataflowRunRequest, + idempotency_key: str, + request_hash: str, +) -> DataflowRun: + return DataflowRun( tenant_id=tenant_id, pipeline_id=pipeline.id, pipeline_revision_id=revision.id, @@ -838,15 +911,24 @@ def start_pipeline_run( diagnostics=[], input_row_count=0, output_row_count=0, - started_at=started_at, + started_at=utcnow(), created_by=actor_id, ) - session.add(run) - session.flush() + +def _execute_pipeline_run( + session: Session, + *, + run: DataflowRun, + pipeline: DataflowPipeline, + revision: DataflowPipelineRevision, + request: DataflowRunRequest, + principal: ApiPrincipal, + registry: object | None, +) -> None: try: result = execute_preview( - graph, + PipelineGraph.model_validate(revision.graph), row_limit=request.row_limit, source_resolver=_datasource_source_resolver( session=session, @@ -854,91 +936,128 @@ def start_pipeline_run( registry=registry, ), ) - if request.publication and ( - result.truncated - or any( - bool(item.get("truncated")) - for item in result.source_fingerprints - ) - ): - raise PipelineExecutionError( - "The bounded runner cannot publish a truncated result or a " - "result calculated from truncated source data." - ) - - run.source_fingerprints = result.source_fingerprints - run.result_schema = [ - item.model_dump(mode="json") for item in result.columns - ] - run.diagnostics = [ - item.model_dump(mode="json") for item in result.diagnostics - ] - run.input_row_count = result.input_row_count - run.output_row_count = result.total_rows - + _apply_pipeline_result(run, result) if request.publication: - publisher = datasource_publication(registry) - if publisher is None: - raise PipelineExecutionError( - "Publishing Dataflow output requires the Datasources " - "publication capability." - ) - target = request.publication - publication = publisher.publish_rows( + _ensure_publishable(result) + _publish_pipeline_result( session, - principal, - request=DatasourcePublicationRequest( - producer_module="dataflow", - producer_run_ref=f"dataflow-run:{run.id}", - idempotency_key=f"{pipeline.id}:{idempotency_key}", - rows=tuple(dict(row) for row in result.rows), - target_datasource_ref=target.target_datasource_ref, - name=target.name or f"{pipeline.name} output", - source_name=target.source_name, - description=target.description, - freeze=target.freeze, - frozen_label=target.frozen_label, - set_current=target.set_current, - provenance={ - "pipeline_ref": f"pipeline:{pipeline.id}", - "pipeline_revision": revision.revision, - "definition_hash": revision.content_hash, - "source_fingerprints": result.source_fingerprints, - }, - metadata={ - **dict(target.metadata), - "dataflow_run_ref": f"dataflow-run:{run.id}", - }, - ), + run=run, + pipeline=pipeline, + revision=revision, + request=request, + result=result, + principal=principal, + registry=registry, ) - run.output_publication_ref = publication.ref - run.output_datasource_ref = publication.datasource.ref - run.output_materialization_ref = publication.materialization.ref run.status = "succeeded" run.finished_at = utcnow() run.error = None except (DatasourceError, PipelineExecutionError) as exc: - run.status = "failed" - run.finished_at = utcnow() - run.error = str(exc) - diagnostics = list(getattr(exc, "diagnostics", ())) - diagnostics.append( - DataflowDiagnostic( - severity="error", - code="run.execution", - message=str(exc), - node_id=getattr(exc, "node_id", None), - ) + _mark_pipeline_run_failed(run, exc) + + +def _apply_pipeline_result( + run: DataflowRun, + result: PipelineExecutionResult, +) -> None: + run.source_fingerprints = result.source_fingerprints + run.result_schema = [ + item.model_dump(mode="json") for item in result.columns + ] + run.diagnostics = [ + item.model_dump(mode="json") for item in result.diagnostics + ] + run.input_row_count = result.input_row_count + run.output_row_count = result.total_rows + + +def _ensure_publishable(result: PipelineExecutionResult) -> None: + source_truncated = any( + bool(item.get("truncated")) + for item in result.source_fingerprints + ) + if result.truncated or source_truncated: + raise PipelineExecutionError( + "The bounded runner cannot publish a truncated result or a " + "result calculated from truncated source data." ) - run.diagnostics = [ - item.model_dump(mode="json") for item in diagnostics - ] - run.source_fingerprints = list( - getattr(exc, "source_fingerprints", ()) + + +def _publish_pipeline_result( + session: Session, + *, + run: DataflowRun, + pipeline: DataflowPipeline, + revision: DataflowPipelineRevision, + request: DataflowRunRequest, + result: PipelineExecutionResult, + principal: ApiPrincipal, + registry: object | None, +) -> None: + publisher = datasource_publication(registry) + if publisher is None: + raise PipelineExecutionError( + "Publishing Dataflow output requires the Datasources " + "publication capability." ) - run.input_row_count = int(getattr(exc, "input_row_count", 0)) - session.flush() - return run, False + target = request.publication + if target is None: + return + publication = publisher.publish_rows( + session, + principal, + request=DatasourcePublicationRequest( + producer_module="dataflow", + producer_run_ref=f"dataflow-run:{run.id}", + idempotency_key=f"{pipeline.id}:{request.idempotency_key.strip()}", + rows=tuple(dict(row) for row in result.rows), + target_datasource_ref=target.target_datasource_ref, + name=target.name or f"{pipeline.name} output", + source_name=target.source_name, + description=target.description, + freeze=target.freeze, + frozen_label=target.frozen_label, + set_current=target.set_current, + provenance={ + "pipeline_ref": f"pipeline:{pipeline.id}", + "pipeline_revision": revision.revision, + "definition_hash": revision.content_hash, + "source_fingerprints": result.source_fingerprints, + }, + metadata={ + **dict(target.metadata), + "dataflow_run_ref": f"dataflow-run:{run.id}", + }, + ), + ) + run.output_publication_ref = publication.ref + run.output_datasource_ref = publication.datasource.ref + run.output_materialization_ref = publication.materialization.ref + + +def _mark_pipeline_run_failed( + run: DataflowRun, + exc: DatasourceError | PipelineExecutionError, +) -> None: + run.status = "failed" + run.finished_at = utcnow() + run.error = str(exc) + diagnostics = list(getattr(exc, "diagnostics", ())) + diagnostics.append( + DataflowDiagnostic( + severity="error", + code="run.execution", + message=str(exc), + node_id=getattr(exc, "node_id", None), + ) + ) + run.diagnostics = [ + item.model_dump(mode="json") for item in diagnostics + ] + run.source_fingerprints = list( + getattr(exc, "source_fingerprints", ()) + ) + run.input_row_count = int(getattr(exc, "input_row_count", 0)) def cancel_pipeline_run( diff --git a/src/govoplan_dataflow/backend/sql_compiler.py b/src/govoplan_dataflow/backend/sql_compiler.py index 3b46c86..ada79e3 100644 --- a/src/govoplan_dataflow/backend/sql_compiler.py +++ b/src/govoplan_dataflow/backend/sql_compiler.py @@ -44,25 +44,112 @@ class _SqlRenderState: distinct: bool = False +@dataclass +class _SqlCompileState: + source_nodes: list[GraphNode] + nodes: list[GraphNode] = field(default_factory=list) + edges: list[GraphEdge] = field(default_factory=list) + qualifier_prefixes: dict[str, str] | None = None + previous_node_id: str = "" + next_transform_layer: int = 0 + edge_sequence: int = 0 + + def next_edge_id(self) -> str: + self.edge_sequence += 1 + return f"edge-{self.edge_sequence}" + + def next_transform_position(self) -> GraphPosition: + position = _position(self.next_transform_layer) + self.next_transform_layer += 1 + return position + + def append_transform(self, node: GraphNode) -> None: + self.nodes.append(node) + self.edges.append( + GraphEdge( + id=self.next_edge_id(), + source=self.previous_node_id, + target=node.id, + ) + ) + self.previous_node_id = node.id + + +@dataclass(frozen=True) +class _ProjectionPlan: + group_by: list[str] + aggregates: list[dict[str, Any]] + fields: list[dict[str, str]] + saw_star: bool + saw_aggregate: bool + + +@dataclass(frozen=True) +class _SqlRenderSources: + left_source: GraphNode | None = None + right_source: GraphNode | None = None + join_node: GraphNode | None = None + union_node: GraphNode | None = None + union_expression: exp.Expression | None = None + left_source_name: str | None = None + right_source_name: str | None = None + right_prefix: str | None = None + right_alias: str | None = None + + def compile_sql( sql_text: str, *, source_nodes: Iterable[GraphNode] = (), ) -> tuple[PipelineGraph, str, list[DataflowDiagnostic]]: + query, union_expression = _parse_sql_query(sql_text) + from_clause, joins = _validated_source_clause( + query, + union_expression=union_expression, + ) + state = _initialize_compile_sources( + from_clause, + joins=joins, + union_expression=union_expression, + source_nodes=list(source_nodes), + ) + _append_filter_nodes(state, query) + projection = _projection_plan( + query, + qualifier_prefixes=state.qualifier_prefixes, + ) + _append_projection_node(state, query, projection) + _append_distinct_node(state, query, projection) + _append_sort_node(state, query) + _append_limit_node(state, query) + state.append_transform( + GraphNode( + id="output", + type="output", + label="Preview output", + position=state.next_transform_position(), + config={}, + ) + ) + graph = PipelineGraph(nodes=state.nodes, edges=state.edges) + diagnostics = validate_graph(graph) + if any(item.severity == "error" for item in diagnostics): + raise SqlCompilationError(diagnostics) + return graph, query.sql(dialect="duckdb", pretty=True), diagnostics + + +def _parse_sql_query(sql_text: str) -> tuple[exp.Select, exp.Union | None]: source_text = sql_text.strip() if not source_text: raise SqlCompilationError([_sql_error("sql.empty", "Enter a SELECT query.")]) try: statements = sqlglot.parse(source_text, read="duckdb") except ParseError as exc: - detail = exc.errors[0] if exc.errors else {} - line = detail.get("line") - column = detail.get("col") - location = f" at line {line}, column {column}" if line and column else "" - raise SqlCompilationError( - [_sql_error("sql.parse", f"SQL could not be parsed{location}: {detail.get('description') or exc}")] - ) from exc - if len(statements) != 1 or not isinstance(statements[0], (exp.Select, exp.Union)): + raise _parse_error(exc) from exc + if len(statements) != 1 or not isinstance( + statements[0], + (exp.Select, exp.Union), + ): raise SqlCompilationError( [ _sql_error( @@ -71,16 +158,32 @@ def compile_sql( ) ] ) - query, union_expression = _select_query(statements[0]) + return _select_query(statements[0]) - from_clause = query.args.get("from_") - union_subquery = ( - from_clause.this - if from_clause is not None - and isinstance(from_clause.this, exp.Subquery) - and isinstance(from_clause.this.this, exp.Union) - else None + +def _parse_error(exc: ParseError) -> SqlCompilationError: + detail = exc.errors[0] if exc.errors else {} + line = detail.get("line") + column = detail.get("col") + location = f" at line {line}, column {column}" if line and column else "" + description = detail.get("description") or exc + return SqlCompilationError( + [ + _sql_error( + "sql.parse", + f"SQL could not be parsed{location}: {description}", + ) + ] ) + + +def _validated_source_clause( + query: exp.Select, + *, + union_expression: exp.Union | None, +) -> tuple[exp.From, list[exp.Join]]: + from_clause = query.args.get("from_") + union_subquery = _union_subquery(from_clause) _reject_unsupported_query_shape(query, allowed_subquery=union_subquery) if from_clause is None or ( union_expression is None @@ -94,335 +197,471 @@ def compile_sql( raise SqlCompilationError( [_sql_error("sql.join_count", "Dataflow SQL currently supports one two-source join.")] ) - source_node_list = list(source_nodes) - nodes: list[GraphNode] = [] - edges: list[GraphEdge] = [] - qualifier_prefixes: dict[str, str] | None = None - previous_node_id: str - edge_sequence = 0 + return from_clause, joins - def next_edge_id() -> str: - nonlocal edge_sequence - edge_sequence += 1 - return f"edge-{edge_sequence}" +def _union_subquery(from_clause: exp.From | None) -> exp.Subquery | None: + if ( + from_clause is not None + and isinstance(from_clause.this, exp.Subquery) + and isinstance(from_clause.this.this, exp.Union) + ): + return from_clause.this + return None + + +def _initialize_compile_sources( + from_clause: exp.From, + *, + joins: list[exp.Join], + union_expression: exp.Union | None, + source_nodes: list[GraphNode], +) -> _SqlCompileState: + state = _SqlCompileState(source_nodes=source_nodes) if union_expression is not None: - if joins: - raise SqlCompilationError( - [_sql_error("sql.union_join", "JOIN outside UNION BY NAME is not supported.")] - ) - union_tables, union_mode = _union_tables(union_expression) - branch_positions = _branch_positions(len(union_tables)) - for index, table in enumerate(union_tables, start=1): - source = _source_node( - table, - source_node_list, - fallback_id=f"source-{index}", - position=branch_positions[index - 1], - ) - if any(existing.id == source.id for existing in nodes): - raise SqlCompilationError( - [ - _sql_error( - "sql.source_identity", - "UNION inputs must use different graph nodes.", - ) - ] - ) - nodes.append(source) - union_node = GraphNode( - id="union", - type="combine.union", - label="Append rows", - position=_position( - 1, - y=sum(position.y for position in branch_positions) - / len(branch_positions), - ), - config={"mode": union_mode}, - ) - nodes.append(union_node) - edges.extend( - GraphEdge( - id=next_edge_id(), - source=source.id, - target=union_node.id, - ) - for source in nodes - if source.type.startswith("source.") - ) - previous_node_id = union_node.id - next_transform_layer = 2 + _initialize_union_sources(state, union_expression, joins=joins) else: - left_table = from_clause.this - if not isinstance(left_table, exp.Table): - raise SqlCompilationError( - [_sql_error("sql.source_required", "SELECT requires a logical tabular source.")] - ) - right_table = joins[0].this if joins else None - if right_table is not None and not isinstance(right_table, exp.Table): - raise SqlCompilationError( - [_sql_error("sql.join_source", "JOIN requires a logical tabular source.")] - ) - tables = [left_table, *([right_table] if isinstance(right_table, exp.Table) else [])] - for table in tables: - if table.catalog or table.db: - raise SqlCompilationError( - [_sql_error("sql.qualified_source", "Use logical source names without a catalog or schema.")] + _initialize_tabular_sources(state, from_clause, joins=joins) + return state + + +def _initialize_union_sources( + state: _SqlCompileState, + union_expression: exp.Union, + *, + joins: list[exp.Join], +) -> None: + if joins: + raise SqlCompilationError( + [_sql_error("sql.union_join", "JOIN outside UNION BY NAME is not supported.")] + ) + union_tables, union_mode = _union_tables(union_expression) + branch_positions = _branch_positions(len(union_tables)) + for index, table in enumerate(union_tables, start=1): + source = _source_node( + table, + state.source_nodes, + fallback_id=f"source-{index}", + position=branch_positions[index - 1], + ) + _append_unique_union_source(state, source) + union_node = GraphNode( + id="union", + type="combine.union", + label="Append rows", + position=_position( + 1, + y=sum(position.y for position in branch_positions) + / len(branch_positions), + ), + config={"mode": union_mode}, + ) + state.nodes.append(union_node) + state.edges.extend( + GraphEdge( + id=state.next_edge_id(), + source=source.id, + target=union_node.id, + ) + for source in state.nodes + if source.type.startswith("source.") + ) + state.previous_node_id = union_node.id + state.next_transform_layer = 2 + + +def _append_unique_union_source( + state: _SqlCompileState, + source: GraphNode, +) -> None: + if any(existing.id == source.id for existing in state.nodes): + raise SqlCompilationError( + [ + _sql_error( + "sql.source_identity", + "UNION inputs must use different graph nodes.", ) - left_source = _source_node( - left_table, - source_node_list, - fallback_id="source-left" if right_table is not None else "source", - position=( - _branch_positions(2)[0] - if right_table is not None - else _position(0) + ] + ) + state.nodes.append(source) + + +def _initialize_tabular_sources( + state: _SqlCompileState, + from_clause: exp.From, + *, + joins: list[exp.Join], +) -> None: + left_table = from_clause.this + if not isinstance(left_table, exp.Table): + raise SqlCompilationError( + [_sql_error("sql.source_required", "SELECT requires a logical tabular source.")] + ) + right_table = joins[0].this if joins else None + if right_table is not None and not isinstance(right_table, exp.Table): + raise SqlCompilationError( + [_sql_error("sql.join_source", "JOIN requires a logical tabular source.")] + ) + _validate_logical_tables( + [left_table, *([right_table] if isinstance(right_table, exp.Table) else [])] + ) + left_source = _source_node( + left_table, + state.source_nodes, + fallback_id="source-left" if right_table is not None else "source", + position=_branch_positions(2)[0] if right_table is not None else _position(0), + ) + state.nodes.append(left_source) + state.previous_node_id = left_source.id + state.next_transform_layer = 1 + if isinstance(right_table, exp.Table): + _initialize_join_source( + state, + joins[0], + left_table=left_table, + right_table=right_table, + left_source=left_source, + ) + + +def _validate_logical_tables(tables: list[exp.Table]) -> None: + if any(table.catalog or table.db for table in tables): + raise SqlCompilationError( + [ + _sql_error( + "sql.qualified_source", + "Use logical source names without a catalog or schema.", + ) + ] + ) + + +def _initialize_join_source( + state: _SqlCompileState, + join: exp.Join, + *, + left_table: exp.Table, + right_table: exp.Table, + left_source: GraphNode, +) -> None: + right_source = _source_node( + right_table, + state.source_nodes, + fallback_id="source-right", + position=_branch_positions(2)[1], + ) + if right_source.id == left_source.id: + raise SqlCompilationError( + [_sql_error("sql.source_identity", "Joined sources must use different graph nodes.")] + ) + right_prefix = f"{right_table.alias_or_name}_" + join_config = _join_config( + join, + left_table=left_table, + right_table=right_table, + ) + join_config["right_prefix"] = right_prefix + join_node = GraphNode( + id="join", + type="combine.join", + label=f"Join {left_table.name} and {right_table.name}", + position=_position(1), + config=join_config, + ) + state.nodes.extend((right_source, join_node)) + state.edges.extend( + ( + GraphEdge( + id=state.next_edge_id(), + source=left_source.id, + target=join_node.id, + target_port="left", + ), + GraphEdge( + id=state.next_edge_id(), + source=right_source.id, + target=join_node.id, + target_port="right", ), ) - nodes.append(left_source) - previous_node_id = left_source.id - next_transform_layer = 1 - if isinstance(right_table, exp.Table): - join_positions = _branch_positions(2) - right_source = _source_node( - right_table, - source_node_list, - fallback_id="source-right", - position=join_positions[1], - ) - if right_source.id == left_source.id: - raise SqlCompilationError( - [_sql_error("sql.source_identity", "Joined sources must use different graph nodes.")] - ) - right_prefix = f"{right_table.alias_or_name}_" - join_config = _join_config(joins[0], left_table=left_table, right_table=right_table) - join_config["right_prefix"] = right_prefix - join_node = GraphNode( - id="join", - type="combine.join", - label=f"Join {left_table.name} and {right_table.name}", - position=_position(1), - config=join_config, - ) - nodes.extend((right_source, join_node)) - edges.extend( - ( - GraphEdge( - id=next_edge_id(), - source=left_source.id, - target=join_node.id, - target_port="left", - ), - GraphEdge( - id=next_edge_id(), - source=right_source.id, - target=join_node.id, - target_port="right", - ), - ) - ) - previous_node_id = join_node.id - qualifier_prefixes = _join_qualifier_prefixes( - left_table, - right_table, - right_prefix=right_prefix, - ) - next_transform_layer = 2 + ) + state.previous_node_id = join_node.id + state.qualifier_prefixes = _join_qualifier_prefixes( + left_table, + right_table, + right_prefix=right_prefix, + ) + state.next_transform_layer = 2 - def next_transform_position() -> GraphPosition: - nonlocal next_transform_layer - position = _position(next_transform_layer) - next_transform_layer += 1 - return position - def append_transform(node: GraphNode) -> None: - nonlocal previous_node_id - nodes.append(node) - edges.append( - GraphEdge( - id=next_edge_id(), - source=previous_node_id, - target=node.id, - ) - ) - previous_node_id = node.id - - conditions = _flatten_and(query.args.get("where").this) if query.args.get("where") else [] +def _append_filter_nodes(state: _SqlCompileState, query: exp.Select) -> None: + where = query.args.get("where") + conditions = _flatten_and(where.this) if where is not None else [] for index, condition in enumerate(conditions, start=1): - append_transform( + state.append_transform( GraphNode( id=f"filter-{index}", type="filter", label=f"Filter {index}", - position=next_transform_position(), + position=state.next_transform_position(), config=_condition_config( condition, - qualifier_prefixes=qualifier_prefixes, + qualifier_prefixes=state.qualifier_prefixes, ), ) ) - group = query.args.get("group") - group_by = [ - _column_name( - item, - context="GROUP BY", - qualifier_prefixes=qualifier_prefixes, - ) - for item in group.expressions - ] if group else [] - aggregate_specs: list[dict[str, Any]] = [] - projection_fields: list[dict[str, str]] = [] - saw_star = False - saw_aggregate = False - for item in query.expressions: - inner = item.this if isinstance(item, exp.Alias) else item - alias = item.alias if isinstance(item, exp.Alias) else "" - aggregate = _aggregate_config( - inner, - alias=alias, - qualifier_prefixes=qualifier_prefixes, - ) - if aggregate is not None: - saw_aggregate = True - aggregate_specs.append(aggregate) - continue - if isinstance(inner, exp.Star): - saw_star = True - continue - if not isinstance(inner, exp.Column): - raise SqlCompilationError( - [_sql_error("sql.select_expression", "SELECT supports columns and COUNT/SUM/AVG/MIN/MAX only.")] - ) - column = _column_name( - inner, - context="SELECT", - qualifier_prefixes=qualifier_prefixes, - ) - projection_fields.append({"column": column, "alias": alias or column}) - if saw_star and len(query.expressions) != 1: +def _projection_plan( + query: exp.Select, + *, + qualifier_prefixes: dict[str, str] | None, +) -> _ProjectionPlan: + group = query.args.get("group") + group_by = ( + [ + _column_name( + item, + context="GROUP BY", + qualifier_prefixes=qualifier_prefixes, + ) + for item in group.expressions + ] + if group + else [] + ) + aggregates: list[dict[str, Any]] = [] + fields: list[dict[str, str]] = [] + saw_star = False + for item in query.expressions: + kind, value = _projection_item( + item, + qualifier_prefixes=qualifier_prefixes, + ) + if kind == "aggregate": + aggregates.append(value) + elif kind == "field": + fields.append(value) + else: + saw_star = True + plan = _ProjectionPlan( + group_by=group_by, + aggregates=aggregates, + fields=fields, + saw_star=saw_star, + saw_aggregate=bool(aggregates), + ) + _validate_projection_plan(query, plan) + return plan + + +def _projection_item( + item: exp.Expression, + *, + qualifier_prefixes: dict[str, str] | None, +) -> tuple[str, Any]: + inner = item.this if isinstance(item, exp.Alias) else item + alias = item.alias if isinstance(item, exp.Alias) else "" + aggregate = _aggregate_config( + inner, + alias=alias, + qualifier_prefixes=qualifier_prefixes, + ) + if aggregate is not None: + return "aggregate", aggregate + if isinstance(inner, exp.Star): + return "star", None + if not isinstance(inner, exp.Column): + raise SqlCompilationError( + [ + _sql_error( + "sql.select_expression", + "SELECT supports columns and COUNT/SUM/AVG/MIN/MAX only.", + ) + ] + ) + column = _column_name( + inner, + context="SELECT", + qualifier_prefixes=qualifier_prefixes, + ) + return "field", {"column": column, "alias": alias or column} + + +def _validate_projection_plan( + query: exp.Select, + plan: _ProjectionPlan, +) -> None: + if plan.saw_star and len(query.expressions) != 1: raise SqlCompilationError( [_sql_error("sql.star_mix", "SELECT * cannot be mixed with other expressions in this dialect.")] ) - if saw_aggregate or group_by: - if saw_star: - raise SqlCompilationError([_sql_error("sql.aggregate_star", "SELECT * cannot be grouped.")]) - plain_columns = [field["column"] for field in projection_fields] - if any(column not in group_by for column in plain_columns): - raise SqlCompilationError( - [_sql_error("sql.grouping", "Every non-aggregate SELECT column must appear in GROUP BY.")] - ) - if any(field["alias"] != field["column"] for field in projection_fields): - raise SqlCompilationError( - [_sql_error("sql.group_alias", "Aliases for GROUP BY columns are not supported yet.")] - ) - if not aggregate_specs: - raise SqlCompilationError( - [_sql_error("sql.aggregate_required", "GROUP BY requires at least one aggregate in this dialect.")] - ) - append_transform( + if not (plan.saw_aggregate or plan.group_by): + return + if plan.saw_star: + raise SqlCompilationError( + [_sql_error("sql.aggregate_star", "SELECT * cannot be grouped.")] + ) + if any(field["column"] not in plan.group_by for field in plan.fields): + raise SqlCompilationError( + [_sql_error("sql.grouping", "Every non-aggregate SELECT column must appear in GROUP BY.")] + ) + if any(field["alias"] != field["column"] for field in plan.fields): + raise SqlCompilationError( + [_sql_error("sql.group_alias", "Aliases for GROUP BY columns are not supported yet.")] + ) + if not plan.aggregates: + raise SqlCompilationError( + [_sql_error("sql.aggregate_required", "GROUP BY requires at least one aggregate in this dialect.")] + ) + + +def _append_projection_node( + state: _SqlCompileState, + query: exp.Select, + plan: _ProjectionPlan, +) -> None: + if plan.saw_aggregate or plan.group_by: + state.append_transform( GraphNode( id="aggregate", type="aggregate", label="Aggregate", - position=next_transform_position(), - config={"group_by": group_by, "aggregates": aggregate_specs}, + position=state.next_transform_position(), + config={ + "group_by": plan.group_by, + "aggregates": plan.aggregates, + }, ) ) - elif not saw_star: - append_transform( + elif not plan.saw_star: + state.append_transform( GraphNode( id="select", type="select", label="Select columns", - position=next_transform_position(), - config={"fields": projection_fields}, + position=state.next_transform_position(), + config={"fields": plan.fields}, ) ) - if query.args.get("distinct"): - if saw_aggregate or group_by: - raise SqlCompilationError( - [_sql_error("sql.distinct_group", "DISTINCT cannot be combined with aggregation yet.")] - ) - append_transform( - GraphNode( - id="distinct", - type="distinct", - label="Remove duplicates", - position=next_transform_position(), - config={"columns": []}, - ) - ) - order = query.args.get("order") - if order: - fields: list[dict[str, str]] = [] - for item in order.expressions: - if not isinstance(item, exp.Ordered): - raise SqlCompilationError([_sql_error("sql.order", "Unsupported ORDER BY expression.")]) - fields.append( - { - "column": _column_name( - item.this, - context="ORDER BY", - qualifier_prefixes=qualifier_prefixes, - ), - "direction": "desc" if item.args.get("desc") else "asc", - } - ) - append_transform( - GraphNode( - id="sort", - type="sort", - label="Sort", - position=next_transform_position(), - config={"fields": fields}, - ) +def _append_distinct_node( + state: _SqlCompileState, + query: exp.Select, + plan: _ProjectionPlan, +) -> None: + if not query.args.get("distinct"): + return + if plan.saw_aggregate or plan.group_by: + raise SqlCompilationError( + [_sql_error("sql.distinct_group", "DISTINCT cannot be combined with aggregation yet.")] ) - - limit = query.args.get("limit") - if limit: - expression = limit.args.get("expression") - if not isinstance(expression, exp.Literal) or expression.is_string: - raise SqlCompilationError([_sql_error("sql.limit", "LIMIT must be a positive integer.")]) - try: - count = int(expression.this) - except (TypeError, ValueError) as exc: - raise SqlCompilationError([_sql_error("sql.limit", "LIMIT must be a positive integer.")]) from exc - if not 1 <= count <= 100_000: - raise SqlCompilationError( - [_sql_error("sql.limit_range", "LIMIT must be between 1 and 100,000.")] - ) - append_transform( - GraphNode( - id="limit", - type="limit", - label="Limit", - position=next_transform_position(), - config={"count": count}, - ) - ) - - append_transform( + state.append_transform( GraphNode( - id="output", - type="output", - label="Preview output", - position=next_transform_position(), - config={}, + id="distinct", + type="distinct", + label="Remove duplicates", + position=state.next_transform_position(), + config={"columns": []}, ) ) - graph = PipelineGraph(nodes=nodes, edges=edges) - diagnostics = validate_graph(graph) - if any(item.severity == "error" for item in diagnostics): - raise SqlCompilationError(diagnostics) - return graph, query.sql(dialect="duckdb", pretty=True), diagnostics + + +def _append_sort_node(state: _SqlCompileState, query: exp.Select) -> None: + order = query.args.get("order") + if not order: + return + fields = [ + _order_field( + item, + qualifier_prefixes=state.qualifier_prefixes, + ) + for item in order.expressions + ] + state.append_transform( + GraphNode( + id="sort", + type="sort", + label="Sort", + position=state.next_transform_position(), + config={"fields": fields}, + ) + ) + + +def _order_field( + item: exp.Expression, + *, + qualifier_prefixes: dict[str, str] | None, +) -> dict[str, str]: + if not isinstance(item, exp.Ordered): + raise SqlCompilationError( + [_sql_error("sql.order", "Unsupported ORDER BY expression.")] + ) + return { + "column": _column_name( + item.this, + context="ORDER BY", + qualifier_prefixes=qualifier_prefixes, + ), + "direction": "desc" if item.args.get("desc") else "asc", + } + + +def _append_limit_node(state: _SqlCompileState, query: exp.Select) -> None: + limit = query.args.get("limit") + if not limit: + return + count = _limit_count(limit.args.get("expression")) + state.append_transform( + GraphNode( + id="limit", + type="limit", + label="Limit", + position=state.next_transform_position(), + config={"count": count}, + ) + ) + + +def _limit_count(expression: exp.Expression | None) -> int: + if not isinstance(expression, exp.Literal) or expression.is_string: + raise SqlCompilationError( + [_sql_error("sql.limit", "LIMIT must be a positive integer.")] + ) + try: + count = int(expression.this) + except (TypeError, ValueError) as exc: + raise SqlCompilationError( + [_sql_error("sql.limit", "LIMIT must be a positive integer.")] + ) from exc + if not 1 <= count <= 100_000: + raise SqlCompilationError( + [_sql_error("sql.limit_range", "LIMIT must be between 1 and 100,000.")] + ) + return count def render_sql(graph: PipelineGraph) -> tuple[str, list[DataflowDiagnostic]]: + diagnostics, ordered, node_by_id = _validated_render_graph(graph) + sources = _render_sources(graph, node_by_id=node_by_id) + state = _SqlRenderState() + _render_operator_nodes( + ordered, + node_by_id=node_by_id, + sources=sources, + state=state, + ) + query = _render_base_query(state, sources=sources) + query = _render_join_clause(query, sources=sources) + query = _apply_render_state(query, state=state) + return query.sql(dialect="duckdb", pretty=True), diagnostics + + +def _validated_render_graph( + graph: PipelineGraph, +) -> tuple[list[DataflowDiagnostic], list[str], dict[str, GraphNode]]: diagnostics = validate_graph(graph) if any(item.severity == "error" for item in diagnostics): raise SqlCompilationError(diagnostics) @@ -430,9 +669,42 @@ def render_sql(graph: PipelineGraph) -> tuple[str, list[DataflowDiagnostic]]: if cyclic: raise SqlCompilationError([_sql_error("graph.cycle", "A cyclic graph cannot be rendered as SQL.")]) node_by_id = {node.id: node for node in graph.nodes} + return diagnostics, ordered, node_by_id + + +def _render_sources( + graph: PipelineGraph, + *, + node_by_id: dict[str, GraphNode], +) -> _SqlRenderSources: source_nodes = [node for node in graph.nodes if node.type.startswith("source.")] join_nodes = [node for node in graph.nodes if node.type == "combine.join"] union_nodes = [node for node in graph.nodes if node.type == "combine.union"] + _validate_render_combines(join_nodes, union_nodes) + if union_nodes: + return _render_union_sources( + graph, + node_by_id=node_by_id, + source_nodes=source_nodes, + union_node=union_nodes[0], + ) + if join_nodes: + return _render_join_sources( + graph, + node_by_id=node_by_id, + join_node=join_nodes[0], + ) + if len(source_nodes) == 1: + return _render_single_source(source_nodes[0]) + raise SqlCompilationError( + [_sql_error("sql.source_count", "SQL rendering needs one source or one two-source join.")] + ) + + +def _validate_render_combines( + join_nodes: list[GraphNode], + union_nodes: list[GraphNode], +) -> None: if len(join_nodes) > 1: raise SqlCompilationError( [_sql_error("sql.join_count", "Only one two-source join can be rendered as SQL.")] @@ -446,99 +718,147 @@ def render_sql(graph: PipelineGraph) -> tuple[str, list[DataflowDiagnostic]]: [_sql_error("sql.combine_count", "JOIN and UNION cannot yet be combined in Dataflow SQL.")] ) - left_source: GraphNode | None = None - right_source: GraphNode | None = None - join_node = join_nodes[0] if join_nodes else None - union_node = union_nodes[0] if union_nodes else None - union_expression: exp.Expression | None = None - left_source_name: str | None = None - right_prefix: str | None = None - right_alias: str | None = None - right_source_name: str | None = None - if union_node is not None: - inputs = graph_inputs_by_port(graph).get(union_node.id, {}) - union_sources = [ - node_by_id[source_id] - for source_id in inputs.get("input", []) - if source_id in node_by_id - ] - if ( - len(union_sources) != len(source_nodes) - or any(not node.type.startswith("source.") for node in union_sources) - or {node.id for node in union_sources} != {node.id for node in source_nodes} - ): - raise SqlCompilationError( - [ - _node_sql_error( - union_node.id, - "sql.union_shape", - "SQL rendering currently requires each append input to connect directly to a source.", - ) - ] - ) - source_names = [ - str(node.config.get("source_name", "")).strip() - for node in union_sources - ] - if any(not source_name for source_name in source_names): - raise SqlCompilationError( - [_sql_error("source.name_required", "Every source needs a logical SQL name.")] - ) - union_expression = exp.select("*").from_(exp.to_table(source_names[0])) - for source_name in source_names[1:]: - union_expression = exp.Union( - this=union_expression, - expression=exp.select("*").from_(exp.to_table(source_name)), - distinct=union_node.config.get("mode") == "distinct", - by_name=True, - ) - elif join_node is not None: - inputs = graph_inputs_by_port(graph).get(join_node.id, {}) - left_source = node_by_id[inputs["left"][0]] - right_source = node_by_id[inputs["right"][0]] - right_prefix = str(join_node.config["right_prefix"]) - right_alias = right_prefix[:-1] - elif len(source_nodes) == 1: - left_source = source_nodes[0] - else: + +def _render_union_sources( + graph: PipelineGraph, + *, + node_by_id: dict[str, GraphNode], + source_nodes: list[GraphNode], + union_node: GraphNode, +) -> _SqlRenderSources: + inputs = graph_inputs_by_port(graph).get(union_node.id, {}) + union_sources = [ + node_by_id[source_id] + for source_id in inputs.get("input", []) + if source_id in node_by_id + ] + _validate_union_source_shape( + union_node, + union_sources=union_sources, + source_nodes=source_nodes, + ) + source_names = [ + str(node.config.get("source_name", "")).strip() + for node in union_sources + ] + _require_source_names(source_names) + union_expression: exp.Expression = exp.select("*").from_( + exp.to_table(source_names[0]) + ) + for source_name in source_names[1:]: + union_expression = exp.Union( + this=union_expression, + expression=exp.select("*").from_(exp.to_table(source_name)), + distinct=union_node.config.get("mode") == "distinct", + by_name=True, + ) + return _SqlRenderSources( + union_node=union_node, + union_expression=union_expression, + ) + + +def _validate_union_source_shape( + union_node: GraphNode, + *, + union_sources: list[GraphNode], + source_nodes: list[GraphNode], +) -> None: + valid = ( + len(union_sources) == len(source_nodes) + and all(node.type.startswith("source.") for node in union_sources) + and {node.id for node in union_sources} + == {node.id for node in source_nodes} + ) + if not valid: raise SqlCompilationError( - [_sql_error("sql.source_count", "SQL rendering needs one source or one two-source join.")] - ) - - if union_node is None: - if left_source is None: - raise SqlCompilationError( - [_sql_error("sql.source_required", "SQL rendering needs a source.")] - ) - left_source_name = str(left_source.config.get("source_name", "")).strip() - right_source_name = ( - str(right_source.config.get("source_name", "")).strip() - if right_source is not None - else None - ) - if not left_source_name or (right_source is not None and not right_source_name): - raise SqlCompilationError( - [_sql_error("source.name_required", "Every source needs a logical SQL name.")] - ) - if right_alias and right_alias.casefold() == left_source_name.casefold(): - raise SqlCompilationError( - [_sql_error("sql.join_alias", "The right-column prefix conflicts with the left source name.")] - ) - - def column_expression(name: str) -> exp.Column: - if right_prefix and right_alias and name.startswith(right_prefix): - right_name = name[len(right_prefix) :] - if not right_name: - raise SqlCompilationError( - [_sql_error("sql.column", "A right-side column name is missing.")] + [ + _node_sql_error( + union_node.id, + "sql.union_shape", + "SQL rendering currently requires each append input to connect directly to a source.", ) - return exp.column(right_name, table=right_alias) - if right_source is not None and left_source_name: - return exp.column(name, table=left_source_name) - return exp.column(name) + ] + ) - state = _SqlRenderState() +def _require_source_names(source_names: list[str]) -> None: + if any(not source_name for source_name in source_names): + raise SqlCompilationError( + [_sql_error("source.name_required", "Every source needs a logical SQL name.")] + ) + + +def _render_join_sources( + graph: PipelineGraph, + *, + node_by_id: dict[str, GraphNode], + join_node: GraphNode, +) -> _SqlRenderSources: + inputs = graph_inputs_by_port(graph).get(join_node.id, {}) + left_source = node_by_id[inputs["left"][0]] + right_source = node_by_id[inputs["right"][0]] + right_prefix = str(join_node.config["right_prefix"]) + right_alias = right_prefix[:-1] + left_source_name = _source_name(left_source) + right_source_name = _source_name(right_source) + _require_source_names([left_source_name, right_source_name]) + if right_alias.casefold() == left_source_name.casefold(): + raise SqlCompilationError( + [_sql_error("sql.join_alias", "The right-column prefix conflicts with the left source name.")] + ) + return _SqlRenderSources( + left_source=left_source, + right_source=right_source, + join_node=join_node, + left_source_name=left_source_name, + right_source_name=right_source_name, + right_prefix=right_prefix, + right_alias=right_alias, + ) + + +def _render_single_source(source: GraphNode) -> _SqlRenderSources: + source_name = _source_name(source) + _require_source_names([source_name]) + return _SqlRenderSources( + left_source=source, + left_source_name=source_name, + ) + + +def _source_name(source: GraphNode) -> str: + return str(source.config.get("source_name", "")).strip() + + +def _column_expression( + name: str, + *, + sources: _SqlRenderSources, +) -> exp.Column: + if ( + sources.right_prefix + and sources.right_alias + and name.startswith(sources.right_prefix) + ): + right_name = name[len(sources.right_prefix) :] + if not right_name: + raise SqlCompilationError( + [_sql_error("sql.column", "A right-side column name is missing.")] + ) + return exp.column(right_name, table=sources.right_alias) + if sources.right_source is not None and sources.left_source_name: + return exp.column(name, table=sources.left_source_name) + return exp.column(name) + + +def _render_operator_nodes( + ordered: list[str], + *, + node_by_id: dict[str, GraphNode], + sources: _SqlRenderSources, + state: _SqlRenderState, +) -> None: for node_id in ordered: node = node_by_id[node_id] if node.type.startswith("source.") or node.type in {"combine.join", "combine.union"}: @@ -554,23 +874,46 @@ def render_sql(graph: PipelineGraph) -> tuple[str, list[DataflowDiagnostic]]: ) ] ) - renderer(node, state, column_expression) + renderer( + node, + state, + lambda name: _column_expression(name, sources=sources), + ) - if union_expression is not None: - query = exp.select(*state.select_expressions).from_( - union_expression.subquery("_unioned") + +def _render_base_query( + state: _SqlRenderState, + *, + sources: _SqlRenderSources, +) -> exp.Select: + if sources.union_expression is not None: + return exp.select(*state.select_expressions).from_( + sources.union_expression.subquery("_unioned") ) - elif left_source_name: - query = exp.select(*state.select_expressions).from_(exp.to_table(left_source_name)) - else: - raise SqlCompilationError( - [_sql_error("sql.source_required", "SQL rendering needs a source.")] + if sources.left_source_name: + return exp.select(*state.select_expressions).from_( + exp.to_table(sources.left_source_name) ) - if join_node is not None and right_source_name and right_alias: + raise SqlCompilationError( + [_sql_error("sql.source_required", "SQL rendering needs a source.")] + ) + + +def _render_join_clause( + query: exp.Select, + *, + sources: _SqlRenderSources, +) -> exp.Select: + join_node = sources.join_node + if ( + join_node is not None + and sources.right_source_name + and sources.right_alias + ): join_conditions = [ exp.EQ( - this=exp.column(str(left_key), table=left_source_name), - expression=exp.column(str(right_key), table=right_alias), + this=exp.column(str(left_key), table=sources.left_source_name), + expression=exp.column(str(right_key), table=sources.right_alias), ) for left_key, right_key in zip( join_node.config["left_keys"], @@ -578,11 +921,19 @@ def render_sql(graph: PipelineGraph) -> tuple[str, list[DataflowDiagnostic]]: strict=True, ) ] - query = query.join( - exp.to_table(right_source_name).as_(right_alias), + return query.join( + exp.to_table(sources.right_source_name).as_(sources.right_alias), on=_combine_and(join_conditions), join_type=str(join_node.config.get("join_type", "inner")), ) + return query + + +def _apply_render_state( + query: exp.Select, + *, + state: _SqlRenderState, +) -> exp.Select: if state.where_conditions: query = query.where(_combine_and(state.where_conditions)) if state.group_by: @@ -593,7 +944,7 @@ def render_sql(graph: PipelineGraph) -> tuple[str, list[DataflowDiagnostic]]: query = query.distinct() if state.limit is not None: query = query.limit(state.limit) - return query.sql(dialect="duckdb", pretty=True), diagnostics + return query def _select_query( diff --git a/tests/test_operators.py b/tests/test_operators.py index f23561e..07863cd 100644 --- a/tests/test_operators.py +++ b/tests/test_operators.py @@ -276,9 +276,58 @@ class DataflowOperatorTests(unittest.TestCase): self.assertEqual(set(NODE_TYPES), set(coverage)) for node_type, facets in coverage.items(): self.assertIn("config_validator", facets, node_type) + self.assertIn("schema_propagator", facets, node_type) self.assertIn("executor", facets, node_type) self.assertIn("sql_renderer", facets, node_type) + def test_maximum_size_linear_graph_validates_and_executes(self) -> None: + nodes = [ + node( + "source", + "source.inline", + { + "source_name": "bounded_input", + "rows": [{"value": 1}], + }, + x=0, + ) + ] + edges: list[GraphEdge] = [] + previous_id = "source" + for index in range(1, 99): + node_id = f"limit-{index}" + nodes.append( + node( + node_id, + "limit", + {"count": 1}, + x=index * 10, + ) + ) + edges.append( + GraphEdge( + id=f"edge-{index}", + source=previous_id, + target=node_id, + ) + ) + previous_id = node_id + nodes.append(node("output", "output", {}, x=990)) + edges.append( + GraphEdge( + id="edge-99", + source=previous_id, + target="output", + ) + ) + graph = PipelineGraph(nodes=nodes, edges=edges) + + self.assertEqual([], validate_graph(graph)) + result = execute_preview(graph, row_limit=1) + + self.assertEqual([{"value": 1}], result.rows) + self.assertEqual(100, len(result.node_diagnostics)) + def test_expression_node_renders_through_registered_sql_compiler(self) -> None: graph = PipelineGraph( nodes=[