Files
govoplan-dataflow/src/govoplan_dataflow/backend/sql_compiler.py
T

2225 lines
68 KiB
Python

from __future__ import annotations
from contextvars import ContextVar
from dataclasses import dataclass, field
from typing import Any, Callable, Iterable
import sqlglot
from sqlglot import exp
from sqlglot.errors import ParseError
from govoplan_dataflow.backend.graph import graph_inputs_by_port, topological_order, validate_graph
from govoplan_dataflow.backend.expressions import parse_expression
from govoplan_dataflow.backend.operator_registry import OPERATOR_REGISTRY
from govoplan_dataflow.backend.schemas import (
DataflowDiagnostic,
DiagnosticSourceLocation,
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
LAYOUT_ORIGIN_X = 60
LAYOUT_CENTER_Y = 180
LAYOUT_LAYER_GAP = 240
LAYOUT_BRANCH_GAP = 120
@dataclass(frozen=True)
class _SqlDiagnosticContext:
source_text: str
root_expression: exp.Expression | None = None
_SQL_DIAGNOSTIC_CONTEXT: ContextVar[_SqlDiagnosticContext | None] = ContextVar(
"dataflow_sql_diagnostic_context",
default=None,
)
@dataclass
class _SqlRenderState:
where_conditions: list[exp.Expression] = field(default_factory=list)
select_expressions: list[exp.Expression] = field(
default_factory=lambda: [exp.Star()]
)
group_by: list[exp.Expression] = field(default_factory=list)
order_by: list[exp.Expression] = field(default_factory=list)
limit: int | None = None
selected: bool = False
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]]:
source_text = sql_text
context_token = _SQL_DIAGNOSTIC_CONTEXT.set(
_SqlDiagnosticContext(source_text=source_text)
)
try:
query, union_expression = _parse_sql_query(source_text)
_SQL_DIAGNOSTIC_CONTEXT.set(
_SqlDiagnosticContext(
source_text=source_text,
root_expression=union_expression or query,
)
)
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
finally:
_SQL_DIAGNOSTIC_CONTEXT.reset(context_token)
def _parse_sql_query(sql_text: str) -> tuple[exp.Select, exp.Union | None]:
if not sql_text.strip():
raise SqlCompilationError([_sql_error("sql.empty", "Enter a SELECT query.")])
try:
statements = sqlglot.parse(sql_text, read="duckdb")
except ParseError as exc:
raise _parse_error(exc) from exc
if len(statements) != 1 or not isinstance(
statements[0],
(exp.Select, exp.Union),
):
raise SqlCompilationError(
[
_sql_error(
"sql.select_only",
"Dataflow SQL accepts exactly one SELECT or UNION BY NAME statement.",
expression=statements[0] if statements else None,
)
]
)
return _select_query(statements[0])
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}",
source_location=_parse_error_location(detail),
)
]
)
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
and not isinstance(from_clause.this, exp.Table)
):
raise SqlCompilationError(
[_sql_error("sql.source_required", "SELECT requires a logical tabular source.")]
)
joins = list(query.args.get("joins") or [])
if len(joins) > 1:
raise SqlCompilationError(
[_sql_error("sql.join_count", "Dataflow SQL currently supports one two-source join.")]
)
return from_clause, joins
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:
_initialize_union_sources(state, union_expression, joins=joins)
else:
_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.",
)
]
)
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:
qualified_table = next(
(table for table in tables if table.catalog or table.db),
None,
)
if qualified_table is not None:
raise SqlCompilationError(
[
_sql_error(
"sql.qualified_source",
"Use logical source names without a catalog or schema.",
expression=qualified_table,
)
]
)
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",
),
)
)
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 _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):
state.append_transform(
GraphNode(
id=f"filter-{index}",
type="filter",
label=f"Filter {index}",
position=state.next_transform_position(),
config=_condition_config(
condition,
qualifier_prefixes=state.qualifier_prefixes,
),
)
)
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.",
expression=item,
)
]
)
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.",
expression=query.expressions[0],
)
]
)
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.",
expression=query.expressions[0],
)
]
)
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.",
expression=query.args.get("group") or query,
)
]
)
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.",
expression=query.args.get("group") or query,
)
]
)
if not plan.aggregates:
raise SqlCompilationError(
[
_sql_error(
"sql.aggregate_required",
"GROUP BY requires at least one aggregate in this dialect.",
expression=query.args.get("group") or query,
)
]
)
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=state.next_transform_position(),
config={
"group_by": plan.group_by,
"aggregates": plan.aggregates,
},
)
)
elif not plan.saw_star:
state.append_transform(
GraphNode(
id="select",
type="select",
label="Select columns",
position=state.next_transform_position(),
config={"fields": plan.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.",
expression=query.args.get("distinct") or query,
)
]
)
state.append_transform(
GraphNode(
id="distinct",
type="distinct",
label="Remove duplicates",
position=state.next_transform_position(),
config={"columns": []},
)
)
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.",
expression=item,
)
]
)
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.",
expression=expression,
)
]
)
try:
count = int(expression.this)
except (TypeError, ValueError) as exc:
raise SqlCompilationError(
[
_sql_error(
"sql.limit",
"LIMIT must be a positive integer.",
expression=expression,
)
]
) from exc
if not 1 <= count <= 100_000:
raise SqlCompilationError(
[
_sql_error(
"sql.limit_range",
"LIMIT must be between 1 and 100,000.",
expression=expression,
)
]
)
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)
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}
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.")]
)
if len(union_nodes) > 1:
raise SqlCompilationError(
[_sql_error("sql.union_count", "Only one append node can be rendered as SQL.")]
)
if join_nodes and union_nodes:
raise SqlCompilationError(
[_sql_error("sql.combine_count", "JOIN and UNION cannot yet be combined in Dataflow SQL.")]
)
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(
[
_node_sql_error(
union_node.id,
"sql.union_shape",
"SQL rendering currently requires each append input to connect directly to a source.",
)
]
)
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"}:
continue
renderer = OPERATOR_REGISTRY.sql_renderer(node.type)
if renderer is None:
raise SqlCompilationError(
[
_node_sql_error(
node.id,
"sql.renderer_missing",
f"Node type {node.type!r} has no registered SQL renderer.",
)
]
)
renderer(
node,
state,
lambda name: _column_expression(name, sources=sources),
)
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")
)
if sources.left_source_name:
return exp.select(*state.select_expressions).from_(
exp.to_table(sources.left_source_name)
)
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=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"],
join_node.config["right_keys"],
strict=True,
)
]
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:
query = query.group_by(*state.group_by)
if state.order_by:
query = query.order_by(*state.order_by)
if state.distinct:
query = query.distinct()
if state.limit is not None:
query = query.limit(state.limit)
return query
def _select_query(
statement: exp.Select | exp.Union,
) -> tuple[exp.Select, exp.Union | None]:
if isinstance(statement, exp.Select):
from_clause = statement.args.get("from_")
if (
from_clause is not None
and isinstance(from_clause.this, exp.Subquery)
and isinstance(from_clause.this.this, exp.Union)
):
return statement, from_clause.this.this
return statement, None
union_expression = statement.copy()
order = union_expression.args.get("order")
limit = union_expression.args.get("limit")
union_expression.set("order", None)
union_expression.set("limit", None)
query = exp.select("*").from_(union_expression.subquery("_unioned"))
if order is not None:
query.set("order", order.copy())
if limit is not None:
query.set("limit", limit.copy())
return query, union_expression
def _union_tables(union_expression: exp.Union) -> tuple[list[exp.Table], str]:
branches: list[exp.Select] = []
modes: set[bool] = set()
def collect(expression: exp.Expression) -> None:
if isinstance(expression, exp.Union):
if expression.args.get("by_name") is not True:
raise SqlCompilationError(
[
_sql_error(
"sql.union_by_name",
"Use UNION BY NAME so SQL matches Dataflow column-name append semantics.",
expression=expression,
)
]
)
if expression.args.get("order") or expression.args.get("limit"):
raise SqlCompilationError(
[
_sql_error(
"sql.union_branch_order",
"Apply ORDER BY or LIMIT after the complete UNION BY NAME.",
expression=expression.args.get("order")
or expression.args.get("limit")
or expression,
)
]
)
modes.add(bool(expression.args.get("distinct")))
collect(expression.this)
collect(expression.expression)
return
if not isinstance(expression, exp.Select):
raise SqlCompilationError(
[
_sql_error(
"sql.union_branch",
"UNION BY NAME branches must be SELECT statements.",
expression=expression,
)
]
)
branches.append(expression)
collect(union_expression)
if len(modes) != 1:
raise SqlCompilationError(
[
_sql_error(
"sql.union_mode",
"A Dataflow append node cannot mix UNION and UNION ALL.",
)
]
)
tables: list[exp.Table] = []
for branch in branches:
_reject_unsupported_query_shape(branch)
if (
branch.args.get("joins")
or branch.args.get("where")
or branch.args.get("group")
or branch.args.get("order")
or branch.args.get("limit")
or branch.args.get("distinct")
):
raise SqlCompilationError(
[
_sql_error(
"sql.union_branch_shape",
"UNION BY NAME inputs must be direct SELECT * source reads.",
expression=branch,
)
]
)
if len(branch.expressions) != 1 or not isinstance(branch.expressions[0], exp.Star):
raise SqlCompilationError(
[
_sql_error(
"sql.union_branch_projection",
"UNION BY NAME inputs must select all source columns.",
expression=branch.expressions[0] if branch.expressions else branch,
)
]
)
from_clause = branch.args.get("from_")
table = from_clause.this if from_clause is not None else None
if not isinstance(table, exp.Table):
raise SqlCompilationError(
[
_sql_error(
"sql.union_source",
"Each UNION BY NAME input needs a logical source.",
expression=branch,
)
]
)
if table.catalog or table.db:
raise SqlCompilationError(
[
_sql_error(
"sql.qualified_source",
"Use logical source names without a catalog or schema.",
expression=table,
)
]
)
tables.append(table)
return tables, "distinct" if modes == {True} else "all"
def _reject_unsupported_query_shape(
query: exp.Select,
*,
allowed_subquery: exp.Subquery | None = None,
) -> None:
unsupported_args = {
"with_": "WITH queries",
"having": "HAVING",
"qualify": "QUALIFY",
"offset": "OFFSET",
"windows": "window definitions",
"locks": "locking clauses",
}
for key, label in unsupported_args.items():
unsupported = query.args.get(key)
if unsupported:
raise SqlCompilationError(
[
_sql_error(
"sql.unsupported_clause",
f"{label} are not supported by the first Dataflow dialect.",
expression=unsupported,
)
]
)
unsupported_subquery = next(
(
subquery
for subquery in query.find_all(exp.Subquery)
if subquery is not allowed_subquery
),
None,
)
if unsupported_subquery is not None:
raise SqlCompilationError(
[
_sql_error(
"sql.subquery",
"Subqueries are not supported by the first Dataflow dialect.",
expression=unsupported_subquery,
)
]
)
def _source_node(
table: exp.Table,
source_nodes: list[GraphNode],
*,
fallback_id: str,
position: GraphPosition,
) -> GraphNode:
source_name = table.name
preserved = 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,
)
if preserved is not None:
return preserved.model_copy(
update={"position": position},
deep=True,
)
return GraphNode(
id=fallback_id,
type="source.reference",
label=source_name,
position=position,
config={
"source_ref": "",
"source_name": source_name,
"expected_fingerprint": "",
},
)
def _join_config(
join: exp.Join,
*,
left_table: exp.Table,
right_table: exp.Table,
) -> dict[str, Any]:
kind = str(join.args.get("kind") or "").casefold()
side = str(join.args.get("side") or "").casefold()
if kind not in {"", "inner", "outer", "semi", "anti"} or side not in {
"",
"left",
"right",
"full",
}:
raise SqlCompilationError(
[
_sql_error(
"sql.join_type",
(
"JOIN supports INNER, LEFT, RIGHT, FULL, SEMI, "
"or ANTI joins only."
),
expression=join,
)
]
)
if kind in {"semi", "anti"} and side:
raise SqlCompilationError(
[
_sql_error(
"sql.join_type",
"SEMI and ANTI joins cannot use a side qualifier.",
expression=join,
)
]
)
join_type = (
kind
if kind in {"semi", "anti"}
else side or ("inner" if kind in {"", "inner"} else "")
)
if not join_type:
raise SqlCompilationError(
[
_sql_error(
"sql.join_type",
"OUTER JOIN requires LEFT, RIGHT, or FULL.",
expression=join,
)
]
)
on_expression = join.args.get("on")
if on_expression is None:
raise SqlCompilationError(
[
_sql_error(
"sql.join_condition",
"JOIN requires an ON key comparison.",
expression=join,
)
]
)
left_qualifiers = _table_qualifiers(left_table)
right_qualifiers = _table_qualifiers(right_table)
if left_qualifiers & right_qualifiers:
raise SqlCompilationError(
[
_sql_error(
"sql.join_alias",
"Joined sources need distinct names or aliases.",
expression=join,
)
]
)
left_keys: list[str] = []
right_keys: list[str] = []
for condition in _flatten_and(on_expression):
if (
not isinstance(condition, exp.EQ)
or not isinstance(condition.this, exp.Column)
or not isinstance(condition.expression, exp.Column)
):
raise SqlCompilationError(
[
_sql_error(
"sql.join_condition",
"JOIN ON accepts equality comparisons between source columns.",
expression=condition,
)
]
)
first = condition.this
second = condition.expression
first_qualifier = first.table.casefold()
second_qualifier = second.table.casefold()
if first_qualifier in left_qualifiers and second_qualifier in right_qualifiers:
left_keys.append(first.name)
right_keys.append(second.name)
elif first_qualifier in right_qualifiers and second_qualifier in left_qualifiers:
left_keys.append(second.name)
right_keys.append(first.name)
else:
raise SqlCompilationError(
[
_sql_error(
"sql.join_qualification",
"Every JOIN key must qualify one left and one right source column.",
expression=condition,
)
]
)
return {
"join_type": join_type,
"left_keys": left_keys,
"right_keys": right_keys,
}
def _join_qualifier_prefixes(
left_table: exp.Table,
right_table: exp.Table,
*,
right_prefix: str,
) -> dict[str, str]:
return {
**{qualifier: "" for qualifier in _table_qualifiers(left_table)},
**{qualifier: right_prefix for qualifier in _table_qualifiers(right_table)},
}
def _table_qualifiers(table: exp.Table) -> set[str]:
return {
qualifier.casefold()
for qualifier in (table.name, table.alias_or_name)
if qualifier
}
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.",
expression=expression,
)
]
)
return [expression]
def _condition_config(
expression: exp.Expression,
*,
qualifier_prefixes: dict[str, str] | None = None,
) -> 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",
qualifier_prefixes=qualifier_prefixes,
),
"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",
qualifier_prefixes=qualifier_prefixes,
),
"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",
qualifier_prefixes=qualifier_prefixes,
),
"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%'.",
expression=expression,
)
]
)
return {
"column": _column_name(
expression.this,
context="WHERE",
qualifier_prefixes=qualifier_prefixes,
),
"operator": "contains",
"value": value[1:-1],
}
raise SqlCompilationError(
[
_sql_error(
"sql.condition",
"WHERE supports simple column comparisons joined with AND.",
expression=expression,
)
]
)
def _condition_expression(
config: dict[str, Any],
*,
column_expression: Callable[[str], exp.Column] = exp.column,
) -> exp.Expression:
column = column_expression(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 _qualified_expression(
source: str,
*,
column_expression: Callable[[str], exp.Column],
) -> exp.Expression:
parsed = parse_expression(source)
expression = parsed.expression.copy()
for column in list(expression.find_all(exp.Column)):
column.replace(column_expression(column.name))
return expression
def _sql_data_type(data_type: str) -> exp.DataType:
mapping = {
"string": exp.DataType.Type.TEXT,
"integer": exp.DataType.Type.BIGINT,
"number": exp.DataType.Type.DECIMAL,
"boolean": exp.DataType.Type.BOOLEAN,
"date": exp.DataType.Type.DATE,
"datetime": exp.DataType.Type.TIMESTAMPTZ,
}
target = mapping.get(data_type)
if target is None:
raise SqlCompilationError(
[_sql_error("convert.target_type", f"Unsupported conversion type {data_type!r}.")]
)
return exp.DataType.build(target)
def _render_filter(
node: GraphNode,
state: _SqlRenderState,
column_expression: Callable[[str], exp.Column],
) -> None:
if state.selected:
raise SqlCompilationError(
[
_node_sql_error(
node.id,
"sql.filter_order",
"Filters after projection or aggregation are not representable yet.",
)
]
)
state.where_conditions.append(
_condition_expression(node.config, column_expression=column_expression)
)
def _render_expression_filter(
node: GraphNode,
state: _SqlRenderState,
column_expression: Callable[[str], exp.Column],
) -> None:
if state.selected:
raise SqlCompilationError(
[
_node_sql_error(
node.id,
"sql.filter_order",
"Expression filters after projection are not representable yet.",
)
]
)
state.where_conditions.append(
_qualified_expression(
str(node.config["expression"]),
column_expression=column_expression,
)
)
def _render_select(
node: GraphNode,
state: _SqlRenderState,
column_expression: Callable[[str], exp.Column],
) -> None:
_require_projection_slot(node, state)
if state.distinct:
raise SqlCompilationError(
[
_node_sql_error(
node.id,
"sql.distinct_order",
"Projection after deduplication is not representable as SELECT DISTINCT.",
)
]
)
state.select_expressions = []
for field_config in node.config["fields"]:
column = (
field_config
if isinstance(field_config, str)
else str(field_config["column"])
)
alias = (
column
if isinstance(field_config, str)
else str(field_config.get("alias") or column)
)
expression: exp.Expression = column_expression(column)
if alias != expression.alias_or_name:
expression = expression.as_(alias)
state.select_expressions.append(expression)
state.selected = True
def _render_aggregate(
node: GraphNode,
state: _SqlRenderState,
column_expression: Callable[[str], exp.Column],
) -> None:
_require_projection_slot(node, state)
if state.distinct:
raise SqlCompilationError(
[
_node_sql_error(
node.id,
"sql.distinct_order",
"Aggregation after deduplication is not representable in the constrained dialect.",
)
]
)
state.select_expressions = [
column_expression(str(column))
for column in node.config.get("group_by", [])
]
state.group_by = [
column_expression(str(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 column_expression(str(column))
)
aggregate_expression = _aggregate_expression(function, argument)
state.select_expressions.append(
aggregate_expression.as_(str(aggregate["alias"]))
)
state.selected = True
def _render_expression(
node: GraphNode,
state: _SqlRenderState,
column_expression: Callable[[str], exp.Column],
) -> None:
_require_projection_slot(node, state)
expression = _qualified_expression(
str(node.config["expression"]),
column_expression=column_expression,
)
state.select_expressions = [
exp.Star(),
expression.as_(str(node.config["target_column"])),
]
state.selected = True
def _render_calculate(
node: GraphNode,
state: _SqlRenderState,
column_expression: Callable[[str], exp.Column],
) -> None:
_require_projection_slot(node, state)
calculated_targets: set[str] = set()
expressions: list[exp.Expression] = [exp.Star()]
for item in node.config["calculations"]:
parsed = parse_expression(str(item["expression"]))
dependencies = set(parsed.columns) & calculated_targets
if dependencies:
raise SqlCompilationError(
[
_node_sql_error(
node.id,
"sql.sequential_calculation",
(
"SQL view cannot render a calculated column that "
"depends on an earlier calculation in the same block."
),
)
]
)
expression = _qualified_expression(
str(item["expression"]),
column_expression=column_expression,
)
target = str(item["target_column"])
expressions.append(expression.as_(target))
calculated_targets.add(target)
state.select_expressions = expressions
state.selected = True
def _render_derive(
node: GraphNode,
state: _SqlRenderState,
column_expression: Callable[[str], exp.Column],
) -> None:
_require_projection_slot(node, state)
columns = [
column_expression(str(column))
for column in node.config["source_columns"]
]
operation = str(node.config["operation"])
if operation == "copy":
expression: exp.Expression = columns[0]
elif operation == "upper":
expression = exp.Upper(this=columns[0])
elif operation == "lower":
expression = exp.Lower(this=columns[0])
elif operation == "trim":
expression = exp.Trim(this=columns[0])
elif operation == "concat":
separator = exp.Literal.string(str(node.config.get("separator", " ")))
expressions: list[exp.Expression] = []
for index, column in enumerate(columns):
if index:
expressions.append(separator.copy())
expressions.append(column)
expression = exp.Concat(expressions=expressions)
elif operation == "coalesce":
expression = exp.Coalesce(
this=columns[0],
expressions=columns[1:],
)
elif operation in {"add", "subtract", "multiply", "divide"}:
operation_type = {
"add": exp.Add,
"subtract": exp.Sub,
"multiply": exp.Mul,
"divide": exp.Div,
}[operation]
expression = operation_type(this=columns[0], expression=columns[1])
else:
raise SqlCompilationError(
[
_node_sql_error(
node.id,
"sql.derive_operation",
f"Derive operation {operation!r} cannot be rendered as SQL.",
)
]
)
state.select_expressions = [
exp.Star(),
expression.as_(str(node.config["target_column"])),
]
state.selected = True
def _render_convert(
node: GraphNode,
state: _SqlRenderState,
column_expression: Callable[[str], exp.Column],
) -> None:
_require_projection_slot(node, state)
on_error = str(node.config.get("on_error", "fail"))
if on_error == "keep":
raise SqlCompilationError(
[
_node_sql_error(
node.id,
"sql.convert_keep",
"Keeping the original value after a conversion error is not representable in SQL view.",
)
]
)
cast_class = exp.TryCast if on_error == "null" else exp.Cast
expression = cast_class(
this=column_expression(str(node.config["source_column"])),
to=_sql_data_type(str(node.config["target_type"])),
)
state.select_expressions = [
exp.Star(),
expression.as_(str(node.config["target_column"])),
]
state.selected = True
def _render_distinct(
node: GraphNode,
state: _SqlRenderState,
_column_expression: Callable[[str], exp.Column],
) -> None:
if node.config.get("columns"):
raise SqlCompilationError(
[
_node_sql_error(
node.id,
"sql.distinct_keys",
"Key-based deduplication is not representable as SELECT DISTINCT.",
)
]
)
if state.distinct:
raise SqlCompilationError(
[
_node_sql_error(
node.id,
"sql.multiple_distinct",
"Only one DISTINCT transform is supported.",
)
]
)
state.distinct = True
def _render_sort(
node: GraphNode,
state: _SqlRenderState,
column_expression: Callable[[str], exp.Column],
) -> None:
state.order_by = [
exp.Ordered(
this=column_expression(str(item["column"])),
desc=item.get("direction", "asc") == "desc",
nulls_first=False,
)
for item in node.config["fields"]
]
def _render_rank(
node: GraphNode,
state: _SqlRenderState,
column_expression: Callable[[str], exp.Column],
) -> None:
_require_projection_slot(node, state)
function = {
"row_number": exp.RowNumber,
"rank": exp.Rank,
"dense_rank": exp.DenseRank,
}[str(node.config.get("method", "row_number"))]()
window = exp.Window(
this=function,
partition_by=[
column_expression(str(column))
for column in node.config.get("partition_by", [])
],
order=exp.Order(
expressions=[
exp.Ordered(
this=column_expression(str(item["column"])),
desc=item.get("direction", "asc") == "desc",
nulls_first=False,
)
for item in node.config["order_by"]
]
),
)
state.select_expressions = [
exp.Star(),
window.as_(str(node.config["target_column"])),
]
state.selected = True
def _render_limit(
node: GraphNode,
state: _SqlRenderState,
_column_expression: Callable[[str], exp.Column],
) -> None:
state.limit = int(node.config["count"])
def _render_noop(
_node: GraphNode,
_state: _SqlRenderState,
_column_expression: Callable[[str], exp.Column],
) -> None:
return None
def _render_unsupported(
node: GraphNode,
_state: _SqlRenderState,
_column_expression: Callable[[str], exp.Column],
) -> None:
raise SqlCompilationError(
[
_node_sql_error(
node.id,
"sql.node_not_representable",
f"{node.type!r} cannot be rendered as SQL.",
)
]
)
def _require_projection_slot(
node: GraphNode,
state: _SqlRenderState,
) -> None:
if state.selected:
raise SqlCompilationError(
[
_node_sql_error(
node.id,
"sql.multiple_select",
(
"Only one calculation, ranking, conversion, derive, "
"select, or aggregate transform is supported in SQL view."
),
)
]
)
def _register_sql_renderers() -> None:
renderers = {
"source.inline": _render_noop,
"source.reference": _render_noop,
"combine.union": _render_noop,
"combine.join": _render_noop,
"filter": _render_filter,
"filter.expression": _render_expression_filter,
"distinct": _render_distinct,
"select": _render_select,
"derive": _render_derive,
"expression": _render_expression,
"calculate": _render_calculate,
"convert": _render_convert,
"replace": _render_unsupported,
"aggregate": _render_aggregate,
"sort": _render_sort,
"window.rank": _render_rank,
"limit": _render_limit,
"quality.rules": _render_unsupported,
"reconcile.compare": _render_unsupported,
"reconcile.decisions": _render_unsupported,
"subflow": _render_unsupported,
"output": _render_noop,
}
for node_type, renderer in renderers.items():
if OPERATOR_REGISTRY.sql_renderer(node_type) is None:
OPERATOR_REGISTRY.register_sql_renderer(node_type, renderer)
_register_sql_renderers()
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.",
expression=expression,
)
]
)
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}.",
expression=expression,
)
]
) from exc
def _aggregate_config(
expression: exp.Expression,
*,
alias: str,
qualifier_prefixes: dict[str, str] | None = None,
) -> 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(),
qualifier_prefixes=qualifier_prefixes,
)
else:
raise SqlCompilationError(
[
_sql_error(
"sql.aggregate_argument",
f"{function.upper()} requires a column or * argument.",
expression=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,
qualifier_prefixes: dict[str, str] | None = None,
) -> str:
if not isinstance(expression, exp.Column):
raise SqlCompilationError(
[
_sql_error(
"sql.column",
f"{context} accepts column names only.",
expression=expression,
)
]
)
if not expression.table and qualifier_prefixes is not None:
raise SqlCompilationError(
[
_sql_error(
"sql.join_column_qualification",
f"{context} columns must be source-qualified when a JOIN is present.",
expression=expression,
)
]
)
if not expression.table:
return expression.name
if qualifier_prefixes is None:
raise SqlCompilationError(
[
_sql_error(
"sql.column",
f"{context} accepts unqualified column names only.",
expression=expression,
)
]
)
prefix = qualifier_prefixes.get(expression.table.casefold())
if prefix is None:
raise SqlCompilationError(
[
_sql_error(
"sql.column_source",
f"{context} references an unknown source qualifier.",
expression=expression,
)
]
)
return f"{prefix}{expression.name}"
def _position(layer: int, *, y: float = LAYOUT_CENTER_Y) -> GraphPosition:
return GraphPosition(
x=LAYOUT_ORIGIN_X + layer * LAYOUT_LAYER_GAP,
y=y,
)
def _branch_positions(count: int) -> list[GraphPosition]:
top = max(
40,
LAYOUT_CENTER_Y - ((count - 1) * LAYOUT_BRANCH_GAP / 2),
)
return [
_position(0, y=top + index * LAYOUT_BRANCH_GAP)
for index in range(count)
]
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,
*,
expression: exp.Expression | None = None,
source_location: DiagnosticSourceLocation | None = None,
) -> DataflowDiagnostic:
context = _SQL_DIAGNOSTIC_CONTEXT.get()
location_expression = expression or (context.root_expression if context else None)
return DataflowDiagnostic(
severity="error",
code=code,
message=message,
field="sql_text",
source_location=source_location or _expression_location(location_expression),
)
def _expression_location(
expression: exp.Expression | None,
) -> DiagnosticSourceLocation | None:
context = _SQL_DIAGNOSTIC_CONTEXT.get()
if context is None or expression is None:
return None
offsets = [
(int(node.meta["start"]), int(node.meta["end"]))
for node in expression.walk()
if node.meta.get("start") is not None and node.meta.get("end") is not None
]
if not offsets:
return None
start_offset = min(item[0] for item in offsets)
end_offset = max(item[1] for item in offsets)
start_line, start_column = _line_and_column(context.source_text, start_offset)
end_line, end_column = _line_and_column(context.source_text, end_offset)
return DiagnosticSourceLocation(
start_line=start_line,
start_column=start_column,
end_line=end_line,
end_column=end_column,
start_offset=start_offset,
end_offset=end_offset,
)
def _parse_error_location(detail: dict[str, Any]) -> DiagnosticSourceLocation | None:
context = _SQL_DIAGNOSTIC_CONTEXT.get()
line = detail.get("line")
column = detail.get("col")
if context is None or not isinstance(line, int) or not isinstance(column, int):
return None
highlight = str(detail.get("highlight") or "")
start_column = max(1, column - max(1, len(highlight)) + 1)
start_offset = _offset_for_line_and_column(
context.source_text,
line,
start_column,
)
end_offset = min(
len(context.source_text) - 1,
start_offset + max(0, len(highlight) - 1),
)
return DiagnosticSourceLocation(
start_line=line,
start_column=start_column,
end_line=line,
end_column=max(start_column, column),
start_offset=max(0, start_offset),
end_offset=max(0, end_offset),
)
def _line_and_column(source_text: str, offset: int) -> tuple[int, int]:
bounded_offset = min(max(0, offset), max(0, len(source_text) - 1))
line = source_text.count("\n", 0, bounded_offset) + 1
last_newline = source_text.rfind("\n", 0, bounded_offset)
return line, bounded_offset - last_newline
def _offset_for_line_and_column(source_text: str, line: int, column: int) -> int:
lines = source_text.splitlines(keepends=True)
if not lines:
return 0
bounded_line = min(max(1, line), len(lines))
line_start = sum(len(item) for item in lines[: bounded_line - 1])
line_content = lines[bounded_line - 1].rstrip("\r\n")
return line_start + min(max(0, column - 1), max(0, len(line_content) - 1))
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"]