From 8eddd1dced81aa43e394aae850d4b85c0c1bd0c0 Mon Sep 17 00:00:00 2001 From: William Horel <80535383+IL-William@users.noreply.github.com> Date: Sun, 11 Oct 2026 07:19:07 +0200 Subject: [PATCH] fix(core): pair the branches of a set operation by position SQL pairs the columns of a UNION, an INTERSECT or an EXCEPT by position: the output takes the first branch's names, and each later branch's column feeds the output column at its place. The analyzer added each branch's columns to the target under their own names instead, so a later branch naming a column otherwise lost it to the column of that name, or added one of its own. SELECT k, a AS x FROM t1 UNION ALL SELECT m.k, m.b FROM t2 AS m came out with three outputs, k, x and b, x derived from t1.a alone. A later branch's columns now take the first branch's names, position by position, stars expanded. Only the columns it gives the set operation's target are renamed: a derived table or a subquery inside the branch keeps its own. UNION BY NAME still pairs by name. A nested set operation hands its first branch's names up, so a third branch pairs with the first. The recursive CTE snapshot gains the edge its second branch writes: `SELECT e.employee_id, e.employee_id, e.name` feeds manager_id from employee_id. The test of CTE aliases reused across UNION branches now finds one output column, read from each branch's own CTE. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_017VTEw2dNNL9Pn57Hkeeirn --- crates/flowscope-core/src/analyzer/context.rs | 20 +++ crates/flowscope-core/src/analyzer/query.rs | 12 +- crates/flowscope-core/src/analyzer/visitor.rs | 45 ++++- crates/flowscope-core/tests/lineage_engine.rs | 164 ++++++++++++++++-- ...hots__run_snapshot_test@recursive_cte.snap | 10 ++ 5 files changed, 226 insertions(+), 25 deletions(-) diff --git a/crates/flowscope-core/src/analyzer/context.rs b/crates/flowscope-core/src/analyzer/context.rs index d0320635..fe03f9dd 100644 --- a/crates/flowscope-core/src/analyzer/context.rs +++ b/crates/flowscope-core/src/analyzer/context.rs @@ -154,6 +154,13 @@ pub(crate) struct StatementContext { /// How many subqueries a predicate reads the statement has opened, each /// analysed against a node of its own. pub(crate) predicate_subqueries: usize, + /// While a later branch of a set operation is read, the names its columns + /// take. SQL pairs the branches of a UNION, an INTERSECT or an EXCEPT by + /// position, not by name, unless it says BY NAME. + pub(crate) branch_names: Option, + /// The first branch's names of the set operation read last, for one that + /// holds it as its own first branch. + pub(crate) last_branch_names: Option<(Option, Vec)>, /// Last join/operation type for edge labeling pub(crate) last_operation: Option, /// Current join information (type + condition) for edge labeling @@ -214,6 +221,17 @@ pub(crate) struct StatementContext { pub(crate) instance_limit_reached: bool, } +/// The names a later branch of a set operation gives its columns: the first +/// branch's, by position. Only columns of `target` take them, so that a +/// derived table or a subquery inside the branch keeps its own. +#[derive(Debug, Clone)] +pub(crate) struct BranchNames { + pub(crate) target: Option, + pub(crate) names: Vec, + /// The position of the next column the branch gives `target`. + pub(crate) next: usize, +} + /// Represents an output column in the SELECT list #[derive(Debug, Clone)] pub(crate) struct OutputColumn { @@ -252,6 +270,8 @@ impl StatementContext { derived_occurrences: HashMap::new(), cte_occurrences: HashMap::new(), predicate_subqueries: 0, + branch_names: None, + last_branch_names: None, last_operation: None, current_join_info: JoinInfo::default(), table_node_ids: HashMap::new(), diff --git a/crates/flowscope-core/src/analyzer/query.rs b/crates/flowscope-core/src/analyzer/query.rs index d27eed84..de8e84fa 100644 --- a/crates/flowscope-core/src/analyzer/query.rs +++ b/crates/flowscope-core/src/analyzer/query.rs @@ -1001,7 +1001,17 @@ impl<'a> Analyzer<'a> { ctx: &mut StatementContext, params: OutputColumnParams, ) { - let normalized_name = self.normalize_identifier(¶ms.name); + // A later branch of a set operation names no column: each takes the + // first branch's name at its position. + let name = match ctx.branch_names.as_mut() { + Some(branch) if branch.target == params.target_node => { + let at = branch.next; + branch.next += 1; + branch.names.get(at).cloned() + } + _ => None, + }; + let normalized_name = self.normalize_identifier(name.as_deref().unwrap_or(¶ms.name)); let node_id = generate_column_node_id(params.target_node.as_deref(), &normalized_name); let ownership_edge_id = params .target_node diff --git a/crates/flowscope-core/src/analyzer/visitor.rs b/crates/flowscope-core/src/analyzer/visitor.rs index 3e0e0e75..7e060a64 100644 --- a/crates/flowscope-core/src/analyzer/visitor.rs +++ b/crates/flowscope-core/src/analyzer/visitor.rs @@ -4,11 +4,11 @@ //! and building lineage graphs. It separates traversal logic (the `Visitor` trait) //! from analysis logic (the `LineageVisitor` implementation). -use super::context::{OutputColumn, StatementContext}; +use super::context::{BranchNames, OutputColumn, StatementContext}; use super::expression::ExpressionAnalyzer; use super::helpers::{ alias_visibility_warning, find_cte_body_span, find_cte_definition_span, - find_derived_table_alias_span, generate_statement_scoped_node_id, + find_derived_table_alias_span, generate_column_node_id, generate_statement_scoped_node_id, }; use super::query::OutputColumnParams; use super::select_analyzer::SelectAnalyzer; @@ -16,8 +16,8 @@ use super::Analyzer; use crate::generated::is_value_table_function; use crate::types::{issue_codes, Issue, Node, NodeType, Span}; use sqlparser::ast::{ - self, CreateView, Cte, Expr, Ident, Join, Query, Select, SetExpr, SetOperator, Statement, - TableAlias, TableFactor, TableWithJoins, Values, + self, CreateView, Cte, Expr, Ident, Join, Query, Select, SetExpr, SetOperator, SetQuantifier, + Statement, TableAlias, TableFactor, TableWithJoins, Values, }; use std::sync::Arc; @@ -723,7 +723,10 @@ impl<'a, 'b> Visitor for LineageVisitor<'a, 'b> { SetExpr::Select(select) => self.visit_select(select), SetExpr::Query(query) => self.visit_query(query), SetExpr::SetOperation { - op, left, right, .. + op, + set_quantifier, + left, + right, } => { let op_name = match op { SetOperator::Union => "UNION", @@ -731,8 +734,40 @@ impl<'a, 'b> Visitor for LineageVisitor<'a, 'b> { SetOperator::Except => "EXCEPT", SetOperator::Minus => "MINUS", }; + let sink = self.ctx.sink_node_id().map(|id| id.to_string()); + let target = self.target_node.clone().or(sink); + let checkpoint = self.ctx.output_columns.len(); + self.ctx.last_branch_names = None; self.visit_set_expr(left); + // The first branch's names: a nested set operation's own first + // branch's, or the columns the branch just gave the target. + let names = match self.ctx.last_branch_names.take() { + Some((of, names)) if of == target => names, + _ => self.ctx.output_columns[checkpoint..] + .iter() + .filter(|c| { + generate_column_node_id(target.as_deref(), &c.name) == c.node_id + }) + .map(|c| c.name.clone()) + .collect(), + }; + let by_name = matches!( + set_quantifier, + SetQuantifier::ByName + | SetQuantifier::AllByName + | SetQuantifier::DistinctByName + ); + let outer = std::mem::replace( + &mut self.ctx.branch_names, + (!by_name).then(|| BranchNames { + target: target.clone(), + names: names.clone(), + next: 0, + }), + ); self.visit_set_expr(right); + self.ctx.branch_names = outer; + self.ctx.last_branch_names = Some((target, names)); if self.target_node.is_some() { self.ctx.last_operation = Some(op_name.to_string()); } diff --git a/crates/flowscope-core/tests/lineage_engine.rs b/crates/flowscope-core/tests/lineage_engine.rs index 8d187c91..690abb07 100644 --- a/crates/flowscope-core/tests/lineage_engine.rs +++ b/crates/flowscope-core/tests/lineage_engine.rs @@ -2430,46 +2430,48 @@ fn cte_self_join_alias_columns_do_not_leak_across_union_scopes() { .iter() .find(|node| node.node_type == NodeType::Output) .expect("Output node should exist"); - let dept_output = stmt + // The branches pair by position: the output has the first branch's name + // only, and the second branch's `b.name` flows into it. + let output_columns: Vec<_> = stmt .nodes .iter() - .find(|node| { + .filter(|node| { node.node_type == NodeType::Column - && &*node.label == "dept_name" && stmt.edges.iter().any(|edge| { edge.edge_type == EdgeType::Ownership && edge.from == output_node.id && edge.to == node.id }) }) - .expect("dept_name output column should exist"); + .collect(); + assert_eq!( + output_columns + .iter() + .map(|node| node.label.to_string()) + .collect::>(), + vec!["emp_name"], + "a UNION's output takes its first branch's names" + ); - let dept_source_ids: HashSet<_> = stmt + let source_ids: HashSet<_> = stmt .edges .iter() - .filter(|edge| edge.edge_type == EdgeType::DataFlow && edge.to == dept_output.id) + .filter(|edge| edge.edge_type == EdgeType::DataFlow && edge.to == output_columns[0].id) .map(|edge| edge.from.clone()) .collect(); - assert!( - !dept_source_ids.is_empty(), - "dept_name should receive a direct data-flow edge from the dept branch source column" - ); - let source_owner_labels: HashSet<_> = stmt .edges .iter() - .filter(|edge| edge.edge_type == EdgeType::Ownership && dept_source_ids.contains(&edge.to)) + .filter(|edge| edge.edge_type == EdgeType::Ownership && source_ids.contains(&edge.to)) .filter_map(|edge| stmt.nodes.iter().find(|node| node.id == edge.from)) .map(|node| node.label.to_string()) .collect(); - assert!( - source_owner_labels.contains("dept"), - "dept_name should be sourced from a dept CTE instance, saw owners {source_owner_labels:?}" - ); - assert!( - !source_owner_labels.contains("emp"), - "dept_name should not reuse emp columns when aliases are reused across UNION branches, saw owners {source_owner_labels:?}" + // Each branch's `b` is its own CTE: the dept branch must not reuse emp's. + assert_eq!( + source_owner_labels, + HashSet::from(["emp".to_string(), "dept".to_string()]), + "each branch should read its own CTE instance, saw owners {source_owner_labels:?}" ); } @@ -13424,3 +13426,127 @@ fn an_outer_cte_named_like_an_inlined_model_s_inner_one_reads_through_it() { expected_sources(&["t.ts"]) ); } + +#[test] +fn a_union_pairs_its_branches_by_position() { + // The second branch names its columns otherwise: SQL pairs them with the + // first branch's by position, and the output has the first branch's names. + let sql = r#" + SELECT k, a AS x FROM t1 + UNION ALL + SELECT m.k, m.b FROM t2 AS m + "#; + let schema = schema_of(vec![ + schema_table(None, None, "t1", &["k", "a"]), + schema_table(None, None, "t2", &["k", "b"]), + ]); + + let result = run_analysis(sql, Dialect::Snowflake, Some(schema)); + let stmt = first_statement(&result); + // The output has the first branch's names, none of the second's. + let output = stmt + .nodes + .iter() + .find(|node| node.node_type == NodeType::Output) + .expect("Output node should exist"); + let mut outputs: Vec = stmt + .edges + .iter() + .filter(|edge| edge.edge_type == EdgeType::Ownership && edge.from == output.id) + .filter_map(|edge| stmt.nodes.iter().find(|node| node.id == edge.to)) + .map(|node| node.label.to_lowercase()) + .collect(); + outputs.sort(); + assert_eq!(outputs, vec!["k", "x"]); + assert_eq!( + base_sources_of(&stmt, "x"), + expected_sources(&["t1.a", "t2.b"]) + ); + assert_eq!( + base_sources_of(&stmt, "k"), + expected_sources(&["t1.k", "t2.k"]) + ); +} + +#[test] +fn a_union_in_a_cte_pairs_by_position_through_its_stars_and_nesting() { + // Three branches, the last a star, inside a CTE the statement reads. + let sql = r#" + WITH u AS ( + SELECT k, 'one' AS origin, a AS x FROM t1 + UNION ALL + SELECT k, 'two', b FROM t2 + UNION ALL + SELECT * FROM t3 + ) + SELECT u.k, u.origin, u.x FROM u + "#; + let schema = schema_of(vec![ + schema_table(None, None, "t1", &["k", "a"]), + schema_table(None, None, "t2", &["k", "b"]), + schema_table(None, None, "t3", &["kk", "label", "c"]), + ]); + + let result = run_analysis(sql, Dialect::Snowflake, Some(schema)); + let stmt = first_statement(&result); + assert_eq!( + base_sources_of(&stmt, "x"), + expected_sources(&["t1.a", "t2.b", "t3.c"]) + ); + assert_eq!( + base_sources_of(&stmt, "k"), + expected_sources(&["t1.k", "t2.k", "t3.kk"]) + ); + assert_eq!( + base_sources_of(&stmt, "origin"), + expected_sources(&["t3.label"]) + ); +} + +#[test] +fn a_union_by_name_still_pairs_by_name() { + let sql = r#" + SELECT k, a AS x FROM t1 + UNION ALL BY NAME + SELECT b AS x, k FROM t2 + "#; + let schema = schema_of(vec![ + schema_table(None, None, "t1", &["k", "a"]), + schema_table(None, None, "t2", &["k", "b"]), + ]); + + let result = run_analysis(sql, Dialect::Snowflake, Some(schema)); + let stmt = first_statement(&result); + assert_eq!( + base_sources_of(&stmt, "x"), + expected_sources(&["t1.a", "t2.b"]) + ); + assert_eq!( + base_sources_of(&stmt, "k"), + expected_sources(&["t1.k", "t2.k"]) + ); +} + +#[test] +fn a_derived_table_inside_a_later_branch_keeps_its_own_names() { + let sql = r#" + SELECT k, a AS x FROM t1 + UNION ALL + SELECT d.k, d.y FROM (SELECT k, b AS y FROM t2) AS d + "#; + let schema = schema_of(vec![ + schema_table(None, None, "t1", &["k", "a"]), + schema_table(None, None, "t2", &["k", "b"]), + ]); + + let result = run_analysis(sql, Dialect::Snowflake, Some(schema)); + let stmt = first_statement(&result); + assert_eq!( + base_sources_of(&stmt, "x"), + expected_sources(&["t1.a", "t2.b"]) + ); + assert_eq!( + base_sources_of(&stmt, "k"), + expected_sources(&["t1.k", "t2.k"]) + ); +} diff --git a/crates/flowscope-core/tests/snapshots/snapshots__run_snapshot_test@recursive_cte.snap b/crates/flowscope-core/tests/snapshots/snapshots__run_snapshot_test@recursive_cte.snap index a66b4b43..35b94dfc 100644 --- a/crates/flowscope-core/tests/snapshots/snapshots__run_snapshot_test@recursive_cte.snap +++ b/crates/flowscope-core/tests/snapshots/snapshots__run_snapshot_test@recursive_cte.snap @@ -1,5 +1,6 @@ --- source: crates/flowscope-core/tests/snapshots.rs +assertion_line: 27 expression: clean_result --- { @@ -273,6 +274,15 @@ expression: clean_result 0 ] }, + { + "id": "edge_8e6e75d17a8dea9d", + "from": "column_78a6f51170482ff9", + "to": "column_a5ca696854a3dcdd", + "type": "data_flow", + "statementIds": [ + 0 + ] + }, { "id": "edge_9837a0968ff8ff3f", "from": "column_f0521d6d6b3f8f6b",