Initialize governed Dataflow module
This commit is contained in:
@@ -0,0 +1,483 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Iterable
|
||||
|
||||
import sqlglot
|
||||
from sqlglot import exp
|
||||
from sqlglot.errors import ParseError
|
||||
|
||||
from govoplan_dataflow.backend.graph import topological_order, validate_graph
|
||||
from govoplan_dataflow.backend.schemas import (
|
||||
DataflowDiagnostic,
|
||||
GraphEdge,
|
||||
GraphNode,
|
||||
GraphPosition,
|
||||
PipelineGraph,
|
||||
)
|
||||
|
||||
|
||||
class SqlCompilationError(ValueError):
|
||||
def __init__(self, diagnostics: list[DataflowDiagnostic]) -> None:
|
||||
super().__init__(diagnostics[0].message if diagnostics else "SQL compilation failed")
|
||||
self.diagnostics = diagnostics
|
||||
|
||||
|
||||
def compile_sql(
|
||||
sql_text: str,
|
||||
*,
|
||||
source_nodes: Iterable[GraphNode] = (),
|
||||
) -> tuple[PipelineGraph, str, list[DataflowDiagnostic]]:
|
||||
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):
|
||||
raise SqlCompilationError(
|
||||
[_sql_error("sql.select_only", "Dataflow SQL accepts exactly one SELECT statement.")]
|
||||
)
|
||||
query = statements[0]
|
||||
_reject_unsupported_query_shape(query)
|
||||
|
||||
tables = list(query.find_all(exp.Table))
|
||||
if len(tables) != 1:
|
||||
raise SqlCompilationError(
|
||||
[_sql_error("sql.source_count", "The first release supports exactly one logical source.")]
|
||||
)
|
||||
table = tables[0]
|
||||
if table.catalog or table.db:
|
||||
raise SqlCompilationError(
|
||||
[_sql_error("sql.qualified_source", "Use the logical source name without a catalog or schema.")]
|
||||
)
|
||||
source_name = table.name
|
||||
preserved_source = next(
|
||||
(
|
||||
node.model_copy(deep=True)
|
||||
for node in source_nodes
|
||||
if node.type.startswith("source.")
|
||||
and str(node.config.get("source_name", "")).casefold() == source_name.casefold()
|
||||
),
|
||||
None,
|
||||
)
|
||||
source_node = preserved_source or GraphNode(
|
||||
id="source",
|
||||
type="source.reference",
|
||||
label=source_name,
|
||||
position=GraphPosition(x=80, y=180),
|
||||
config={"source_name": source_name},
|
||||
)
|
||||
|
||||
nodes = [source_node]
|
||||
conditions = _flatten_and(query.args.get("where").this) if query.args.get("where") else []
|
||||
for index, condition in enumerate(conditions, start=1):
|
||||
nodes.append(
|
||||
GraphNode(
|
||||
id=f"filter-{index}",
|
||||
type="filter",
|
||||
label=f"Filter {index}",
|
||||
position=_position(len(nodes)),
|
||||
config=_condition_config(condition),
|
||||
)
|
||||
)
|
||||
|
||||
group = query.args.get("group")
|
||||
group_by = [_column_name(item, context="GROUP BY") 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)
|
||||
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")
|
||||
projection_fields.append({"column": column, "alias": alias or column})
|
||||
|
||||
if 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.")]
|
||||
)
|
||||
nodes.append(
|
||||
GraphNode(
|
||||
id="aggregate",
|
||||
type="aggregate",
|
||||
label="Aggregate",
|
||||
position=_position(len(nodes)),
|
||||
config={"group_by": group_by, "aggregates": aggregate_specs},
|
||||
)
|
||||
)
|
||||
elif not saw_star:
|
||||
nodes.append(
|
||||
GraphNode(
|
||||
id="select",
|
||||
type="select",
|
||||
label="Select columns",
|
||||
position=_position(len(nodes)),
|
||||
config={"fields": projection_fields},
|
||||
)
|
||||
)
|
||||
|
||||
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"),
|
||||
"direction": "desc" if item.args.get("desc") else "asc",
|
||||
}
|
||||
)
|
||||
nodes.append(
|
||||
GraphNode(
|
||||
id="sort",
|
||||
type="sort",
|
||||
label="Sort",
|
||||
position=_position(len(nodes)),
|
||||
config={"fields": fields},
|
||||
)
|
||||
)
|
||||
|
||||
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.")]
|
||||
)
|
||||
nodes.append(
|
||||
GraphNode(
|
||||
id="limit",
|
||||
type="limit",
|
||||
label="Limit",
|
||||
position=_position(len(nodes)),
|
||||
config={"count": count},
|
||||
)
|
||||
)
|
||||
|
||||
nodes.append(
|
||||
GraphNode(
|
||||
id="output",
|
||||
type="output",
|
||||
label="Preview output",
|
||||
position=_position(len(nodes)),
|
||||
config={},
|
||||
)
|
||||
)
|
||||
edges = [
|
||||
GraphEdge(
|
||||
id=f"edge-{source.id}-{target.id}",
|
||||
source=source.id,
|
||||
target=target.id,
|
||||
)
|
||||
for source, target in zip(nodes, nodes[1:])
|
||||
]
|
||||
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 render_sql(graph: PipelineGraph) -> tuple[str, list[DataflowDiagnostic]]:
|
||||
diagnostics = validate_graph(graph)
|
||||
if any(item.severity == "error" for item in diagnostics):
|
||||
raise SqlCompilationError(diagnostics)
|
||||
ordered, cyclic = topological_order(graph)
|
||||
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}
|
||||
source = node_by_id[ordered[0]]
|
||||
source_name = str(source.config.get("source_name", "")).strip()
|
||||
if not source_name:
|
||||
raise SqlCompilationError([_sql_error("source.name_required", "The source needs a logical SQL name.")])
|
||||
|
||||
where_conditions: list[exp.Expression] = []
|
||||
select_expressions: list[exp.Expression] = [exp.Star()]
|
||||
group_by: list[exp.Expression] = []
|
||||
order_by: list[exp.Expression] = []
|
||||
limit: int | None = None
|
||||
selected = False
|
||||
|
||||
for node_id in ordered[1:]:
|
||||
node = node_by_id[node_id]
|
||||
if node.type == "filter":
|
||||
if selected:
|
||||
raise SqlCompilationError(
|
||||
[_node_sql_error(node.id, "sql.filter_order", "Filters after projection or aggregation are not representable yet.")]
|
||||
)
|
||||
where_conditions.append(_condition_expression(node.config))
|
||||
elif node.type == "select":
|
||||
if selected:
|
||||
raise SqlCompilationError(
|
||||
[_node_sql_error(node.id, "sql.multiple_select", "Only one select or aggregate transform is supported.")]
|
||||
)
|
||||
select_expressions = []
|
||||
for field in node.config["fields"]:
|
||||
column = field if isinstance(field, str) else str(field["column"])
|
||||
alias = column if isinstance(field, str) else str(field.get("alias") or column)
|
||||
expression: exp.Expression = exp.column(column)
|
||||
if alias != column:
|
||||
expression = expression.as_(alias)
|
||||
select_expressions.append(expression)
|
||||
selected = True
|
||||
elif node.type == "aggregate":
|
||||
if selected:
|
||||
raise SqlCompilationError(
|
||||
[_node_sql_error(node.id, "sql.multiple_select", "Only one select or aggregate transform is supported.")]
|
||||
)
|
||||
select_expressions = [exp.column(column) for column in node.config.get("group_by", [])]
|
||||
group_by = [exp.column(column) for column in node.config.get("group_by", [])]
|
||||
for aggregate in node.config["aggregates"]:
|
||||
function = str(aggregate["function"])
|
||||
column = aggregate.get("column")
|
||||
argument: exp.Expression = exp.Star() if function == "count" and column in (None, "", "*") else exp.column(str(column))
|
||||
aggregate_expression = _aggregate_expression(function, argument)
|
||||
select_expressions.append(aggregate_expression.as_(str(aggregate["alias"])))
|
||||
selected = True
|
||||
elif node.type == "sort":
|
||||
order_by = [
|
||||
exp.Ordered(
|
||||
this=exp.column(str(field["column"])),
|
||||
desc=field.get("direction", "asc") == "desc",
|
||||
nulls_first=False,
|
||||
)
|
||||
for field in node.config["fields"]
|
||||
]
|
||||
elif node.type == "limit":
|
||||
limit = int(node.config["count"])
|
||||
elif node.type != "output":
|
||||
raise SqlCompilationError(
|
||||
[_node_sql_error(node.id, "sql.node_not_representable", f"{node.type!r} cannot be rendered as SQL.")]
|
||||
)
|
||||
|
||||
query = exp.select(*select_expressions).from_(exp.to_table(source_name))
|
||||
if where_conditions:
|
||||
query = query.where(_combine_and(where_conditions))
|
||||
if group_by:
|
||||
query = query.group_by(*group_by)
|
||||
if order_by:
|
||||
query = query.order_by(*order_by)
|
||||
if limit is not None:
|
||||
query = query.limit(limit)
|
||||
return query.sql(dialect="duckdb", pretty=True), diagnostics
|
||||
|
||||
|
||||
def _reject_unsupported_query_shape(query: exp.Select) -> None:
|
||||
unsupported_args = {
|
||||
"with_": "WITH queries",
|
||||
"distinct": "DISTINCT",
|
||||
"having": "HAVING",
|
||||
"qualify": "QUALIFY",
|
||||
"offset": "OFFSET",
|
||||
"windows": "window definitions",
|
||||
"locks": "locking clauses",
|
||||
}
|
||||
for key, label in unsupported_args.items():
|
||||
if query.args.get(key):
|
||||
raise SqlCompilationError(
|
||||
[_sql_error("sql.unsupported_clause", f"{label} are not supported by the first Dataflow dialect.")]
|
||||
)
|
||||
if any(True for _ in query.find_all(exp.Join)):
|
||||
raise SqlCompilationError([_sql_error("sql.join", "JOIN support belongs to the comparison/reconciliation slice.")])
|
||||
if any(True for _ in query.find_all(exp.Subquery)):
|
||||
raise SqlCompilationError([_sql_error("sql.subquery", "Subqueries are not supported by the first Dataflow dialect.")])
|
||||
|
||||
|
||||
def _flatten_and(expression: exp.Expression) -> list[exp.Expression]:
|
||||
if isinstance(expression, exp.And):
|
||||
return [*_flatten_and(expression.this), *_flatten_and(expression.expression)]
|
||||
if isinstance(expression, exp.Or):
|
||||
raise SqlCompilationError([_sql_error("sql.or", "OR conditions are not supported yet; use separate pipelines.")])
|
||||
return [expression]
|
||||
|
||||
|
||||
def _condition_config(expression: exp.Expression) -> dict[str, Any]:
|
||||
if isinstance(expression, exp.Not) and isinstance(expression.this, exp.Is):
|
||||
inner = expression.this
|
||||
if isinstance(inner.this, exp.Column) and isinstance(inner.expression, exp.Null):
|
||||
return {"column": _column_name(inner.this, context="WHERE"), "operator": "not_null"}
|
||||
if isinstance(expression, exp.Is):
|
||||
if isinstance(expression.this, exp.Column) and isinstance(expression.expression, exp.Null):
|
||||
return {"column": _column_name(expression.this, context="WHERE"), "operator": "is_null"}
|
||||
mapping: tuple[tuple[type[exp.Expression], str], ...] = (
|
||||
(exp.EQ, "eq"),
|
||||
(exp.NEQ, "ne"),
|
||||
(exp.GT, "gt"),
|
||||
(exp.GTE, "gte"),
|
||||
(exp.LT, "lt"),
|
||||
(exp.LTE, "lte"),
|
||||
)
|
||||
for expression_type, operator in mapping:
|
||||
if isinstance(expression, expression_type):
|
||||
if not isinstance(expression.this, exp.Column):
|
||||
break
|
||||
return {
|
||||
"column": _column_name(expression.this, context="WHERE"),
|
||||
"operator": operator,
|
||||
"value": _literal_value(expression.expression),
|
||||
}
|
||||
if isinstance(expression, exp.Like) and isinstance(expression.this, exp.Column):
|
||||
value = _literal_value(expression.expression)
|
||||
if not isinstance(value, str) or not (value.startswith("%") and value.endswith("%")):
|
||||
raise SqlCompilationError(
|
||||
[_sql_error("sql.like", "LIKE is supported only as a contains pattern: LIKE '%value%'.")]
|
||||
)
|
||||
return {
|
||||
"column": _column_name(expression.this, context="WHERE"),
|
||||
"operator": "contains",
|
||||
"value": value[1:-1],
|
||||
}
|
||||
raise SqlCompilationError(
|
||||
[_sql_error("sql.condition", "WHERE supports simple column comparisons joined with AND.")]
|
||||
)
|
||||
|
||||
|
||||
def _condition_expression(config: dict[str, Any]) -> exp.Expression:
|
||||
column = exp.column(str(config["column"]))
|
||||
operator = str(config["operator"])
|
||||
if operator == "is_null":
|
||||
return exp.Is(this=column, expression=exp.Null())
|
||||
if operator == "not_null":
|
||||
return exp.Not(this=exp.Is(this=column, expression=exp.Null()))
|
||||
value = exp.convert(config.get("value"))
|
||||
mapping: dict[str, type[exp.Expression]] = {
|
||||
"eq": exp.EQ,
|
||||
"ne": exp.NEQ,
|
||||
"gt": exp.GT,
|
||||
"gte": exp.GTE,
|
||||
"lt": exp.LT,
|
||||
"lte": exp.LTE,
|
||||
}
|
||||
if operator in mapping:
|
||||
return mapping[operator](this=column, expression=value)
|
||||
if operator == "contains":
|
||||
return exp.Like(this=column, expression=exp.Literal.string(f"%{config.get('value', '')}%"))
|
||||
raise SqlCompilationError([_sql_error("filter.operator", f"Unsupported filter operator {operator!r}.")])
|
||||
|
||||
|
||||
def _literal_value(expression: exp.Expression) -> Any:
|
||||
if isinstance(expression, exp.Null):
|
||||
return None
|
||||
if isinstance(expression, exp.Boolean):
|
||||
return bool(expression.this)
|
||||
if not isinstance(expression, exp.Literal):
|
||||
raise SqlCompilationError([_sql_error("sql.literal", "Comparisons require a literal value.")])
|
||||
if expression.is_string:
|
||||
return str(expression.this)
|
||||
text = str(expression.this)
|
||||
try:
|
||||
return int(text)
|
||||
except ValueError:
|
||||
try:
|
||||
return float(text)
|
||||
except ValueError as exc:
|
||||
raise SqlCompilationError([_sql_error("sql.literal", f"Unsupported literal {text!r}.")]) from exc
|
||||
|
||||
|
||||
def _aggregate_config(expression: exp.Expression, *, alias: str) -> dict[str, Any] | None:
|
||||
mapping: tuple[tuple[type[exp.Expression], str], ...] = (
|
||||
(exp.Count, "count"),
|
||||
(exp.Sum, "sum"),
|
||||
(exp.Avg, "avg"),
|
||||
(exp.Min, "min"),
|
||||
(exp.Max, "max"),
|
||||
)
|
||||
for expression_type, function in mapping:
|
||||
if not isinstance(expression, expression_type):
|
||||
continue
|
||||
argument = expression.this
|
||||
if isinstance(argument, exp.Star):
|
||||
column = "*"
|
||||
elif isinstance(argument, exp.Column):
|
||||
column = _column_name(argument, context=function.upper())
|
||||
else:
|
||||
raise SqlCompilationError(
|
||||
[_sql_error("sql.aggregate_argument", f"{function.upper()} requires a column or * argument.")]
|
||||
)
|
||||
output_alias = alias or (f"{function}_all" if column == "*" else f"{function}_{column}")
|
||||
return {"function": function, "column": column, "alias": output_alias}
|
||||
return None
|
||||
|
||||
|
||||
def _aggregate_expression(function: str, argument: exp.Expression) -> exp.Expression:
|
||||
mapping: dict[str, type[exp.Expression]] = {
|
||||
"count": exp.Count,
|
||||
"sum": exp.Sum,
|
||||
"avg": exp.Avg,
|
||||
"min": exp.Min,
|
||||
"max": exp.Max,
|
||||
}
|
||||
return mapping[function](this=argument)
|
||||
|
||||
|
||||
def _column_name(expression: exp.Expression, *, context: str) -> str:
|
||||
if not isinstance(expression, exp.Column) or expression.table:
|
||||
raise SqlCompilationError(
|
||||
[_sql_error("sql.column", f"{context} accepts unqualified column names only.")]
|
||||
)
|
||||
return expression.name
|
||||
|
||||
|
||||
def _position(index: int) -> GraphPosition:
|
||||
return GraphPosition(x=80 + index * 220, y=180)
|
||||
|
||||
|
||||
def _combine_and(expressions: list[exp.Expression]) -> exp.Expression:
|
||||
combined = expressions[0]
|
||||
for expression in expressions[1:]:
|
||||
combined = exp.And(this=combined, expression=expression)
|
||||
return combined
|
||||
|
||||
|
||||
def _sql_error(code: str, message: str) -> DataflowDiagnostic:
|
||||
return DataflowDiagnostic(severity="error", code=code, message=message, field="sql_text")
|
||||
|
||||
|
||||
def _node_sql_error(node_id: str, code: str, message: str) -> DataflowDiagnostic:
|
||||
return DataflowDiagnostic(severity="error", code=code, message=message, node_id=node_id)
|
||||
|
||||
|
||||
__all__ = ["SqlCompilationError", "compile_sql", "render_sql"]
|
||||
Reference in New Issue
Block a user