Fix union SQL and selection previews
This commit is contained in:
@@ -45,7 +45,9 @@ class GraphNode(BaseModel):
|
||||
|
||||
|
||||
class GraphEdge(BaseModel):
|
||||
id: str = Field(min_length=1, max_length=120, pattern=r"^[A-Za-z0-9_.:-]+$")
|
||||
# Older WebUI releases composed edge IDs from both node IDs. Keep those
|
||||
# definitions readable while new clients use short opaque edge IDs.
|
||||
id: str = Field(min_length=1, max_length=255, pattern=r"^[A-Za-z0-9_.:-]+$")
|
||||
source: str = Field(min_length=1, max_length=100)
|
||||
target: str = Field(min_length=1, max_length=100)
|
||||
source_port: str = Field(default="output", min_length=1, max_length=80)
|
||||
|
||||
@@ -40,15 +40,30 @@ def compile_sql(
|
||||
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):
|
||||
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 statement.")]
|
||||
[
|
||||
_sql_error(
|
||||
"sql.select_only",
|
||||
"Dataflow SQL accepts exactly one SELECT or UNION BY NAME statement.",
|
||||
)
|
||||
]
|
||||
)
|
||||
query = statements[0]
|
||||
_reject_unsupported_query_shape(query)
|
||||
query, union_expression = _select_query(statements[0])
|
||||
|
||||
from_clause = query.args.get("from_")
|
||||
if from_clause is None or not isinstance(from_clause.this, exp.Table):
|
||||
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
|
||||
)
|
||||
_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.")]
|
||||
)
|
||||
@@ -57,81 +72,135 @@ def compile_sql(
|
||||
raise SqlCompilationError(
|
||||
[_sql_error("sql.join_count", "Dataflow SQL currently supports one two-source join.")]
|
||||
)
|
||||
left_table = from_clause.this
|
||||
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.")]
|
||||
)
|
||||
|
||||
source_node_list = list(source_nodes)
|
||||
left_source = _source_node(
|
||||
left_table,
|
||||
source_node_list,
|
||||
fallback_id="source-left" if right_table is not None else "source",
|
||||
position=GraphPosition(x=60, y=120 if right_table is not None else 180),
|
||||
)
|
||||
nodes = [left_source]
|
||||
nodes: list[GraphNode] = []
|
||||
edges: list[GraphEdge] = []
|
||||
previous_node_id = left_source.id
|
||||
qualifier_prefixes: dict[str, str] | None = None
|
||||
if isinstance(right_table, exp.Table):
|
||||
right_source = _source_node(
|
||||
right_table,
|
||||
source_node_list,
|
||||
fallback_id="source-right",
|
||||
position=GraphPosition(x=60, y=280),
|
||||
)
|
||||
if right_source.id == left_source.id:
|
||||
previous_node_id: str
|
||||
edge_sequence = 0
|
||||
|
||||
def next_edge_id() -> str:
|
||||
nonlocal edge_sequence
|
||||
edge_sequence += 1
|
||||
return f"edge-{edge_sequence}"
|
||||
|
||||
if union_expression is not None:
|
||||
if joins:
|
||||
raise SqlCompilationError(
|
||||
[_sql_error("sql.source_identity", "Joined sources must use different graph nodes.")]
|
||||
[_sql_error("sql.union_join", "JOIN outside UNION BY NAME is not supported.")]
|
||||
)
|
||||
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=GraphPosition(x=300, y=200),
|
||||
config=join_config,
|
||||
union_tables, union_mode = _union_tables(union_expression)
|
||||
for index, table in enumerate(union_tables, start=1):
|
||||
source = _source_node(
|
||||
table,
|
||||
source_node_list,
|
||||
fallback_id=f"source-{index}",
|
||||
position=GraphPosition(x=60, y=80 + ((index - 1) * 140)),
|
||||
)
|
||||
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=GraphPosition(x=300, y=80 + ((len(nodes) - 1) * 70)),
|
||||
config={"mode": union_mode},
|
||||
)
|
||||
nodes.extend((right_source, join_node))
|
||||
nodes.append(union_node)
|
||||
edges.extend(
|
||||
(
|
||||
GraphEdge(
|
||||
id=f"edge-{left_source.id}-{join_node.id}-left",
|
||||
source=left_source.id,
|
||||
target=join_node.id,
|
||||
target_port="left",
|
||||
),
|
||||
GraphEdge(
|
||||
id=f"edge-{right_source.id}-{join_node.id}-right",
|
||||
source=right_source.id,
|
||||
target=join_node.id,
|
||||
target_port="right",
|
||||
),
|
||||
GraphEdge(
|
||||
id=next_edge_id(),
|
||||
source=source.id,
|
||||
target=union_node.id,
|
||||
)
|
||||
for source in nodes
|
||||
if source.type.startswith("source.")
|
||||
)
|
||||
previous_node_id = join_node.id
|
||||
qualifier_prefixes = _join_qualifier_prefixes(
|
||||
previous_node_id = union_node.id
|
||||
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.")]
|
||||
)
|
||||
left_source = _source_node(
|
||||
left_table,
|
||||
right_table,
|
||||
right_prefix=right_prefix,
|
||||
source_node_list,
|
||||
fallback_id="source-left" if right_table is not None else "source",
|
||||
position=GraphPosition(x=60, y=120 if right_table is not None else 180),
|
||||
)
|
||||
nodes.append(left_source)
|
||||
previous_node_id = left_source.id
|
||||
if isinstance(right_table, exp.Table):
|
||||
right_source = _source_node(
|
||||
right_table,
|
||||
source_node_list,
|
||||
fallback_id="source-right",
|
||||
position=GraphPosition(x=60, y=280),
|
||||
)
|
||||
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=GraphPosition(x=300, y=200),
|
||||
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,
|
||||
)
|
||||
|
||||
def append_transform(node: GraphNode) -> None:
|
||||
nonlocal previous_node_id
|
||||
nodes.append(node)
|
||||
edges.append(
|
||||
GraphEdge(
|
||||
id=f"edge-{previous_node_id}-{node.id}",
|
||||
id=next_edge_id(),
|
||||
source=previous_node_id,
|
||||
target=node.id,
|
||||
)
|
||||
@@ -322,17 +391,67 @@ def render_sql(graph: PipelineGraph) -> tuple[str, list[DataflowDiagnostic]]:
|
||||
node_by_id = {node.id: node for node in graph.nodes}
|
||||
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"]
|
||||
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.")]
|
||||
)
|
||||
|
||||
left_source: GraphNode
|
||||
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
|
||||
if join_node is not 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]]
|
||||
@@ -345,20 +464,25 @@ def render_sql(graph: PipelineGraph) -> tuple[str, list[DataflowDiagnostic]]:
|
||||
[_sql_error("sql.source_count", "SQL rendering needs one source or one two-source join.")]
|
||||
)
|
||||
|
||||
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.")]
|
||||
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):
|
||||
@@ -368,7 +492,7 @@ def render_sql(graph: PipelineGraph) -> tuple[str, list[DataflowDiagnostic]]:
|
||||
[_sql_error("sql.column", "A right-side column name is missing.")]
|
||||
)
|
||||
return exp.column(right_name, table=right_alias)
|
||||
if right_source is not None:
|
||||
if right_source is not None and left_source_name:
|
||||
return exp.column(name, table=left_source_name)
|
||||
return exp.column(name)
|
||||
|
||||
@@ -382,7 +506,7 @@ def render_sql(graph: PipelineGraph) -> tuple[str, list[DataflowDiagnostic]]:
|
||||
|
||||
for node_id in ordered:
|
||||
node = node_by_id[node_id]
|
||||
if node.type.startswith("source.") or node.type == "combine.join":
|
||||
if node.type.startswith("source.") or node.type in {"combine.join", "combine.union"}:
|
||||
continue
|
||||
if node.type == "filter":
|
||||
if selected:
|
||||
@@ -482,7 +606,16 @@ def render_sql(graph: PipelineGraph) -> tuple[str, list[DataflowDiagnostic]]:
|
||||
[_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(left_source_name))
|
||||
if union_expression is not None:
|
||||
query = exp.select(*select_expressions).from_(
|
||||
union_expression.subquery("_unioned")
|
||||
)
|
||||
elif left_source_name:
|
||||
query = exp.select(*select_expressions).from_(exp.to_table(left_source_name))
|
||||
else:
|
||||
raise SqlCompilationError(
|
||||
[_sql_error("sql.source_required", "SQL rendering needs a source.")]
|
||||
)
|
||||
if join_node is not None and right_source_name and right_alias:
|
||||
join_conditions = [
|
||||
exp.EQ(
|
||||
@@ -513,7 +646,124 @@ def render_sql(graph: PipelineGraph) -> tuple[str, list[DataflowDiagnostic]]:
|
||||
return query.sql(dialect="duckdb", pretty=True), diagnostics
|
||||
|
||||
|
||||
def _reject_unsupported_query_shape(query: exp.Select) -> None:
|
||||
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.",
|
||||
)
|
||||
]
|
||||
)
|
||||
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.",
|
||||
)
|
||||
]
|
||||
)
|
||||
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.")]
|
||||
)
|
||||
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.",
|
||||
)
|
||||
]
|
||||
)
|
||||
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.",
|
||||
)
|
||||
]
|
||||
)
|
||||
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.")]
|
||||
)
|
||||
if table.catalog or table.db:
|
||||
raise SqlCompilationError(
|
||||
[_sql_error("sql.qualified_source", "Use logical source names without a catalog or schema.")]
|
||||
)
|
||||
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",
|
||||
@@ -527,7 +777,10 @@ def _reject_unsupported_query_shape(query: exp.Select) -> None:
|
||||
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.Subquery)):
|
||||
if any(
|
||||
subquery is not allowed_subquery
|
||||
for subquery in query.find_all(exp.Subquery)
|
||||
):
|
||||
raise SqlCompilationError([_sql_error("sql.subquery", "Subqueries are not supported by the first Dataflow dialect.")])
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user