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",