Initialize governed Dataflow module
This commit is contained in:
@@ -0,0 +1,312 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass
|
||||
from decimal import Decimal
|
||||
from typing import Any
|
||||
|
||||
from govoplan_dataflow.backend.graph import graph_input_map, topological_order, validate_graph
|
||||
from govoplan_dataflow.backend.schemas import (
|
||||
DataflowDiagnostic,
|
||||
NodePreviewDiagnostic,
|
||||
PipelineGraph,
|
||||
PreviewColumn,
|
||||
)
|
||||
|
||||
|
||||
EXECUTOR_VERSION = "dataflow-preview-v1"
|
||||
MAX_EXECUTION_SECONDS = 2.0
|
||||
MAX_RESULT_BYTES = 1_000_000
|
||||
|
||||
|
||||
class PipelineExecutionError(RuntimeError):
|
||||
def __init__(self, message: str, *, node_id: str | None = None) -> None:
|
||||
super().__init__(message)
|
||||
self.node_id = node_id
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PipelineExecutionResult:
|
||||
rows: list[dict[str, Any]]
|
||||
total_rows: int
|
||||
truncated: bool
|
||||
columns: list[PreviewColumn]
|
||||
diagnostics: list[DataflowDiagnostic]
|
||||
node_diagnostics: list[NodePreviewDiagnostic]
|
||||
source_fingerprints: list[dict[str, Any]]
|
||||
input_row_count: int
|
||||
|
||||
|
||||
def execute_preview(graph: PipelineGraph, *, row_limit: int) -> PipelineExecutionResult:
|
||||
validation = validate_graph(graph)
|
||||
errors = [item for item in validation if item.severity == "error"]
|
||||
if errors:
|
||||
raise PipelineExecutionError(errors[0].message, node_id=errors[0].node_id)
|
||||
|
||||
node_by_id = {node.id: node for node in graph.nodes}
|
||||
inputs = graph_input_map(graph)
|
||||
ordered, cyclic = topological_order(graph)
|
||||
if cyclic:
|
||||
raise PipelineExecutionError("Pipeline graph contains a cycle")
|
||||
|
||||
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:
|
||||
if time.monotonic() - started > MAX_EXECUTION_SECONDS:
|
||||
raise PipelineExecutionError("Preview exceeded the two-second execution limit", node_id=node_id)
|
||||
node = node_by_id[node_id]
|
||||
node_started = time.monotonic()
|
||||
input_rows = [] if node.type.startswith("source.") else outputs[inputs[node_id]]
|
||||
try:
|
||||
if node.type == "source.inline":
|
||||
output_rows = [dict(row) for row in node.config.get("rows", [])]
|
||||
input_row_count += len(output_rows)
|
||||
source_fingerprints.append(
|
||||
{
|
||||
"node_id": node.id,
|
||||
"source_name": node.config.get("source_name"),
|
||||
"kind": "inline",
|
||||
"fingerprint": _rows_fingerprint(output_rows),
|
||||
"row_count": len(output_rows),
|
||||
}
|
||||
)
|
||||
elif node.type == "source.reference":
|
||||
raise PipelineExecutionError(
|
||||
"This connector-backed source is not available to the local preview executor.",
|
||||
node_id=node.id,
|
||||
)
|
||||
elif node.type == "filter":
|
||||
output_rows = _filter_rows(input_rows, node.config, node_id=node.id)
|
||||
elif node.type == "select":
|
||||
output_rows = _select_rows(input_rows, node.config)
|
||||
elif node.type == "aggregate":
|
||||
output_rows = _aggregate_rows(input_rows, node.config, node_id=node.id)
|
||||
elif node.type == "sort":
|
||||
output_rows = _sort_rows(input_rows, node.config)
|
||||
elif node.type == "limit":
|
||||
output_rows = input_rows[: int(node.config["count"])]
|
||||
elif node.type == "output":
|
||||
output_rows = [dict(row) for row in input_rows]
|
||||
else:
|
||||
raise PipelineExecutionError(f"Unsupported node type: {node.type}", node_id=node.id)
|
||||
except PipelineExecutionError:
|
||||
raise
|
||||
except (ArithmeticError, TypeError, ValueError) as exc:
|
||||
raise PipelineExecutionError(str(exc), node_id=node.id) from exc
|
||||
|
||||
if len(json.dumps(output_rows, default=str).encode("utf-8")) > MAX_RESULT_BYTES:
|
||||
raise PipelineExecutionError(
|
||||
"A preview node exceeded the one-megabyte result limit.",
|
||||
node_id=node.id,
|
||||
)
|
||||
outputs[node.id] = output_rows
|
||||
node_diagnostics.append(
|
||||
NodePreviewDiagnostic(
|
||||
node_id=node.id,
|
||||
status="succeeded",
|
||||
input_rows=len(input_rows),
|
||||
output_rows=len(output_rows),
|
||||
duration_ms=round((time.monotonic() - node_started) * 1000, 3),
|
||||
columns=infer_columns(output_rows),
|
||||
)
|
||||
)
|
||||
|
||||
output_node = next(node for node in graph.nodes if node.type == "output")
|
||||
all_rows = outputs[output_node.id]
|
||||
rows = all_rows[:row_limit]
|
||||
return PipelineExecutionResult(
|
||||
rows=rows,
|
||||
total_rows=len(all_rows),
|
||||
truncated=len(rows) < len(all_rows),
|
||||
columns=infer_columns(all_rows),
|
||||
diagnostics=[],
|
||||
node_diagnostics=node_diagnostics,
|
||||
source_fingerprints=source_fingerprints,
|
||||
input_row_count=input_row_count,
|
||||
)
|
||||
|
||||
|
||||
def infer_columns(rows: list[dict[str, Any]]) -> list[PreviewColumn]:
|
||||
names: list[str] = []
|
||||
for row in rows:
|
||||
for name in row:
|
||||
if name not in names:
|
||||
names.append(name)
|
||||
columns: list[PreviewColumn] = []
|
||||
for name in names:
|
||||
values = [row.get(name) for row in rows]
|
||||
concrete = [value for value in values if value is not None]
|
||||
inferred = _type_name(concrete[0]) if concrete else "unknown"
|
||||
if any(_type_name(value) != inferred for value in concrete[1:]):
|
||||
inferred = "mixed"
|
||||
columns.append(
|
||||
PreviewColumn(
|
||||
name=name,
|
||||
type=inferred,
|
||||
nullable=len(concrete) != len(values),
|
||||
)
|
||||
)
|
||||
return columns
|
||||
|
||||
|
||||
def _filter_rows(
|
||||
rows: list[dict[str, Any]],
|
||||
config: dict[str, Any],
|
||||
*,
|
||||
node_id: str,
|
||||
) -> list[dict[str, Any]]:
|
||||
column = str(config["column"])
|
||||
operator = str(config["operator"])
|
||||
expected = config.get("value")
|
||||
result: list[dict[str, Any]] = []
|
||||
for row in rows:
|
||||
actual = row.get(column)
|
||||
try:
|
||||
matches = _compare(actual, operator, expected)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise PipelineExecutionError(
|
||||
f"Cannot apply {operator!r} to column {column!r}: {exc}",
|
||||
node_id=node_id,
|
||||
) from exc
|
||||
if matches:
|
||||
result.append(dict(row))
|
||||
return result
|
||||
|
||||
|
||||
def _compare(actual: Any, operator: str, expected: Any) -> bool:
|
||||
if operator == "is_null":
|
||||
return actual is None
|
||||
if operator == "not_null":
|
||||
return actual is not None
|
||||
if operator == "eq":
|
||||
return actual == expected
|
||||
if operator == "ne":
|
||||
return actual != expected
|
||||
if operator == "contains":
|
||||
return expected is not None and str(expected).casefold() in str(actual or "").casefold()
|
||||
if actual is None or expected is None:
|
||||
return False
|
||||
if operator == "gt":
|
||||
return actual > expected
|
||||
if operator == "gte":
|
||||
return actual >= expected
|
||||
if operator == "lt":
|
||||
return actual < expected
|
||||
if operator == "lte":
|
||||
return actual <= expected
|
||||
raise ValueError(f"unknown operator {operator!r}")
|
||||
|
||||
|
||||
def _select_rows(rows: list[dict[str, Any]], config: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
fields = config["fields"]
|
||||
normalized = [
|
||||
(
|
||||
field if isinstance(field, str) else str(field["column"]),
|
||||
field if isinstance(field, str) else str(field.get("alias") or field["column"]),
|
||||
)
|
||||
for field in fields
|
||||
]
|
||||
return [
|
||||
{alias: row.get(column) for column, alias in normalized}
|
||||
for row in rows
|
||||
]
|
||||
|
||||
|
||||
def _aggregate_rows(
|
||||
rows: list[dict[str, Any]],
|
||||
config: dict[str, Any],
|
||||
*,
|
||||
node_id: str,
|
||||
) -> list[dict[str, Any]]:
|
||||
group_by = [str(item) for item in config.get("group_by", [])]
|
||||
aggregates = list(config["aggregates"])
|
||||
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[()] = []
|
||||
|
||||
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 _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"
|
||||
result.sort(
|
||||
key=lambda row: (row.get(column) is None, _sortable_value(row.get(column))),
|
||||
reverse=reverse,
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def _sortable_value(value: Any) -> tuple[str, Any]:
|
||||
if isinstance(value, (int, float, Decimal, str)):
|
||||
return type(value).__name__, value
|
||||
return type(value).__name__, str(value)
|
||||
|
||||
|
||||
def _rows_fingerprint(rows: list[dict[str, Any]]) -> str:
|
||||
encoded = json.dumps(rows, sort_keys=True, separators=(",", ":"), default=str)
|
||||
return hashlib.sha256(encoded.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def _type_name(value: Any) -> str:
|
||||
if value is None:
|
||||
return "unknown"
|
||||
if isinstance(value, bool):
|
||||
return "boolean"
|
||||
if isinstance(value, int):
|
||||
return "integer"
|
||||
if isinstance(value, (float, Decimal)):
|
||||
return "number"
|
||||
if isinstance(value, str):
|
||||
return "string"
|
||||
if isinstance(value, list):
|
||||
return "array"
|
||||
if isinstance(value, dict):
|
||||
return "object"
|
||||
return type(value).__name__.lower()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"EXECUTOR_VERSION",
|
||||
"PipelineExecutionError",
|
||||
"PipelineExecutionResult",
|
||||
"execute_preview",
|
||||
"infer_columns",
|
||||
]
|
||||
Reference in New Issue
Block a user