From e5d6ce7ebd7f39c9db17c3760ff61f712c875f8a Mon Sep 17 00:00:00 2001 From: William Horel <80535383+IL-William@users.noreply.github.com> Date: Sat, 10 Oct 2026 18:23:57 +0200 Subject: [PATCH] fix(core): give each CTE its own node when a nested WITH reuses its name A CTE's node was keyed by statement and name, and its output columns were kept in one map per statement keyed by name. Two CTEs of one statement sharing a name in different WITH blocks, as dbt writes when it inlines ephemeral models that each name a CTE `final`, therefore shared one node and one column list, and each reader received both lineages: WITH a AS ( WITH final AS (SELECT k, ts FROM sat_a) SELECT k, ts FROM final ), b AS ( WITH final AS (SELECT k, ts FROM sat_b) SELECT k, ts FROM final ) SELECT a.ts AS a_ts, b.ts AS b_ts FROM a JOIN b ON a.k = b.k a_ts came out derived from sat_a.ts and sat_b.ts both, and so did b_ts, with no issue raised. Count the CTEs of each name per statement, as derived tables already are. The first keeps the key it always had, so nothing changes for a statement that does not repeat a name, and each later one is keyed `name#n`, with a node and column nodes of its own and the name as its label. A WITH that defines a name an earlier CTE has hides that CTE's definition and columns until its query ends, then puts them back, so a name read after a nested WITH is the one in scope there. Backward inference through `SELECT *` grouped wildcards by the target's name and chained them by the source's. It now keys them by node, so two CTEs of one name with no schema infer their own columns; a derived table reusing an alias, which the earlier fix left to inference by name, is covered too. Tests: the query above; the same inside two derived tables; an outer CTE shadowed in a nested WITH and read after it; the same with no schema; a filter on one of two CTEs of one name; an unqualified column of an outer CTE beside an inner namesake. All six fail without the change. Two more pass either way and guard it: stars expanded per CTE, and dbt's shape, an outer `final` over an inlined model holding its own `final`. Co-Authored-By: Claude Opus 5.5 --- crates/flowscope-core/src/analyzer/context.rs | 12 +- crates/flowscope-core/src/analyzer/query.rs | 77 +++--- crates/flowscope-core/src/analyzer/visitor.rs | 49 +++- crates/flowscope-core/tests/lineage_engine.rs | 221 ++++++++++++++++++ 4 files changed, 304 insertions(+), 55 deletions(-) diff --git a/crates/flowscope-core/src/analyzer/context.rs b/crates/flowscope-core/src/analyzer/context.rs index c1b89b3a..d0320635 100644 --- a/crates/flowscope-core/src/analyzer/context.rs +++ b/crates/flowscope-core/src/analyzer/context.rs @@ -14,10 +14,14 @@ pub(crate) const DBT_MODEL_SINK_METADATA_KEY: &str = "dbtModelSink"; pub(crate) struct PendingWildcard { /// The source canonical name (table or CTE) with unknown columns pub(crate) source_canonical: String, - /// The target entity (CTE name or derived table alias) receiving the wildcard - pub(crate) target_name: String, /// Node ID of the source pub(crate) source_node_id: Arc, + /// Node ID of the target. Two CTEs of one name in nested `WITH` blocks + /// share `target_name` and nothing else. + pub(crate) target_node_id: Arc, + /// The node whose own wildcards feed the source: the CTE definition the + /// source names, else the source itself. + pub(crate) upstream_node_id: Arc, } /// A single table/view reference in a FROM or JOIN clause. @@ -144,6 +148,9 @@ pub(crate) struct StatementContext { /// a derived table reusing the alias of an earlier one in another scope /// gets a node of its own. pub(crate) derived_occurrences: HashMap, + /// How many CTEs of each name the statement has defined, so that a nested + /// `WITH` defining a name an earlier CTE has gets a node of its own. + pub(crate) cte_occurrences: HashMap, /// How many subqueries a predicate reads the statement has opened, each /// analysed against a node of its own. pub(crate) predicate_subqueries: usize, @@ -243,6 +250,7 @@ impl StatementContext { table_aliases: HashMap::new(), subquery_aliases: HashSet::new(), derived_occurrences: HashMap::new(), + cte_occurrences: HashMap::new(), predicate_subqueries: 0, last_operation: None, current_join_info: JoinInfo::default(), diff --git a/crates/flowscope-core/src/analyzer/query.rs b/crates/flowscope-core/src/analyzer/query.rs index 9a233405..d27eed84 100644 --- a/crates/flowscope-core/src/analyzer/query.rs +++ b/crates/flowscope-core/src/analyzer/query.rs @@ -717,17 +717,19 @@ impl<'a> Analyzer<'a> { }); } - // Find the CTE/alias name from the node ID for backward inference - // Use the cte_node_to_name reverse mapping for efficient lookup - let target_alias_name = - ctx.cte_node_to_name.get(&Arc::from(target)).cloned(); - - // Record pending wildcard for backward column inference - if let Some(alias_name) = target_alias_name { + // Record pending wildcard for backward column inference, when + // the target is a CTE or a derived table + if ctx.cte_node_to_name.contains_key(&Arc::from(target)) { + let upstream_node_id = ctx + .cte_definitions + .get(&table_canonical) + .cloned() + .unwrap_or_else(|| source_node_id.clone()); ctx.pending_wildcards.push(PendingWildcard { source_canonical: table_canonical.clone(), - target_name: alias_name, source_node_id, + target_node_id: Arc::from(target), + upstream_node_id, }); } } @@ -1494,16 +1496,18 @@ impl<'a> Analyzer<'a> { /// /// # Algorithm Overview /// - /// 1. **Group wildcards by target**: Wildcards are grouped by their `target_name` - /// (the CTE or derived table alias that receives the `SELECT *` columns). + /// 1. **Group wildcards by target**: Wildcards are grouped by their + /// `target_node_id` (the CTE or derived table node that receives the + /// `SELECT *` columns), not by name: two CTEs of one name in nested `WITH` + /// blocks are two targets. /// /// 2. **Build node index**: Creates an O(1) lookup map from node ID to node /// to avoid repeated linear scans when collecting owned columns. /// /// 3. **Propagate columns**: For each target, finds its owned columns and - /// creates corresponding columns on the source tables. The `source_canonical` - /// field in `PendingWildcard` matches the `target_name` of upstream wildcards, - /// enabling recursive chain propagation. + /// creates corresponding columns on the source tables. The `upstream_node_id` + /// field in `PendingWildcard` matches the `target_node_id` of upstream + /// wildcards, enabling recursive chain propagation. /// /// 4. **Cycle detection**: Uses `visited_pairs` to track (target, source) pairs /// and prevent infinite recursion on cyclic references. @@ -1512,11 +1516,13 @@ impl<'a> Analyzer<'a> { return; } - // Build map: target_name -> Vec - let mut wildcards_by_target: HashMap> = HashMap::new(); + // Build map: target node -> Vec. Keyed by node rather + // than by name, so that two CTEs of one name in nested `WITH` blocks + // infer their own columns. + let mut wildcards_by_target: HashMap, Vec> = HashMap::new(); for pw in ctx.pending_wildcards.drain(..) { wildcards_by_target - .entry(pw.target_name.clone()) + .entry(pw.target_node_id.clone()) .or_default() .push(pw); } @@ -1531,14 +1537,10 @@ impl<'a> Analyzer<'a> { .collect(); // Track visited target/source pairs to prevent cycles - let mut visited_pairs: HashSet<(String, String)> = HashSet::new(); - - for (target_name, wildcards) in &wildcards_by_target { - let Some(cte_node_id) = self.lookup_inference_target_node(ctx, target_name) else { - continue; - }; + let mut visited_pairs: HashSet<(Arc, Arc)> = HashSet::new(); - let owned_columns = self.collect_owned_columns(ctx, &cte_node_id, &node_index); + for (cte_node_id, wildcards) in &wildcards_by_target { + let owned_columns = self.collect_owned_columns(ctx, cte_node_id, &node_index); if owned_columns.is_empty() { continue; } @@ -1546,7 +1548,7 @@ impl<'a> Analyzer<'a> { for wildcard in wildcards { self.propagate_wildcard_columns( ctx, - target_name, + cte_node_id, wildcard, &owned_columns, &wildcards_by_target, @@ -1557,21 +1559,6 @@ impl<'a> Analyzer<'a> { } } - /// Locates the node ID to use as the inference target for a wildcard. - fn lookup_inference_target_node( - &self, - ctx: &StatementContext, - target_name: &str, - ) -> Option> { - if let Some(node_id) = ctx.cte_definitions.get(target_name) { - return Some(node_id.clone()); - } - - ctx.cte_node_to_name - .iter() - .find_map(|(node_id, name)| (name == target_name).then(|| node_id.clone())) - } - /// Collects column information owned by a CTE node. /// /// Uses the provided `node_index` for O(1) node lookups instead of linear scans. @@ -1607,11 +1594,11 @@ impl<'a> Analyzer<'a> { fn propagate_wildcard_columns( &mut self, ctx: &mut StatementContext, - target_name: &str, + target_node_id: &Arc, wildcard: &PendingWildcard, columns: &[(String, Option, Arc)], - wildcards_by_target: &HashMap>, - visited_pairs: &mut HashSet<(String, String)>, + wildcards_by_target: &HashMap, Vec>, + visited_pairs: &mut HashSet<(Arc, Arc)>, depth: usize, ) { // Enforce recursion depth limit to prevent stack overflow @@ -1619,7 +1606,7 @@ impl<'a> Analyzer<'a> { return; } - let pair = (target_name.to_string(), wildcard.source_canonical.clone()); + let pair = (target_node_id.clone(), wildcard.upstream_node_id.clone()); if !visited_pairs.insert(pair.clone()) { return; } @@ -1639,11 +1626,11 @@ impl<'a> Analyzer<'a> { .collect(); // Recursively propagate if this source is itself a target of another wildcard - if let Some(upstream_wildcards) = wildcards_by_target.get(&wildcard.source_canonical) { + if let Some(upstream_wildcards) = wildcards_by_target.get(&wildcard.upstream_node_id) { for upstream_wildcard in upstream_wildcards { self.propagate_wildcard_columns( ctx, - &wildcard.source_canonical, + &wildcard.upstream_node_id, upstream_wildcard, &source_columns, wildcards_by_target, diff --git a/crates/flowscope-core/src/analyzer/visitor.rs b/crates/flowscope-core/src/analyzer/visitor.rs index 89e1a2b9..7607429c 100644 --- a/crates/flowscope-core/src/analyzer/visitor.rs +++ b/crates/flowscope-core/src/analyzer/visitor.rs @@ -4,7 +4,7 @@ //! and building lineage graphs. It separates traversal logic (the `Visitor` trait) //! from analysis logic (the `LineageVisitor` implementation). -use super::context::StatementContext; +use super::context::{OutputColumn, StatementContext}; use super::expression::ExpressionAnalyzer; use super::helpers::{ alias_visibility_warning, find_cte_body_span, find_cte_definition_span, @@ -622,10 +622,28 @@ impl<'a, 'b> LineageVisitor<'a, 'b> { impl<'a, 'b> Visitor for LineageVisitor<'a, 'b> { fn visit_query(&mut self, query: &Query) { + // A name this WITH defines hides, until the query ends, the CTE an + // enclosing or earlier WITH gave it: its node and its columns, put back + // on the way out. Keyed by the name alone, the second definition would + // take the first's node and overwrite its columns, and every reader of + // either would get both lineages. + let mut hidden: Vec<(String, Arc, Option>)> = Vec::new(); if let Some(with) = &query.with { let mut cte_ids: Vec<(String, Arc)> = Vec::new(); for cte in &with.cte_tables { let cte_name = cte.alias.name.to_string(); + let occurrence = self + .ctx + .cte_occurrences + .entry(cte_name.clone()) + .or_insert(0); + *occurrence += 1; + // The first keeps the key it always had, so a statement that + // never repeats a name is analysed as before. + let key = match *occurrence { + 1 => cte_name.clone(), + n => format!("{cte_name}#{n}"), + }; let cte_span = self.locate_cte_definition_span(&cte_name); let body_span = cte_span.and_then(|span| { let sql = if let Some(source) = &self.analyzer.current_statement_source { @@ -636,11 +654,7 @@ impl<'a, 'b> Visitor for LineageVisitor<'a, 'b> { find_cte_body_span(sql, span) }); let cte_id = self.ctx.add_node(Node { - id: generate_statement_scoped_node_id( - "cte", - self.ctx.statement_index, - &cte_name, - ), + id: generate_statement_scoped_node_id("cte", self.ctx.statement_index, &key), node_type: NodeType::Cte, label: cte_name.clone().into(), qualified_name: Some(cte_name.clone().into()), @@ -654,9 +668,16 @@ impl<'a, 'b> Visitor for LineageVisitor<'a, 'b> { *cursor = (*cursor).max(span.end); } - self.ctx + if let Some(previous) = self + .ctx .cte_definitions - .insert(cte_name.clone(), cte_id.clone()); + .insert(cte_name.clone(), cte_id.clone()) + { + if previous != cte_id { + let columns = self.ctx.aliased_subquery_columns.remove(&cte_name); + hidden.push((cte_name.clone(), previous, columns)); + } + } self.ctx .cte_node_to_name .insert(cte_id.clone(), cte_name.clone()); @@ -680,6 +701,18 @@ impl<'a, 'b> Visitor for LineageVisitor<'a, 'b> { if let Some(order_by) = &query.order_by { self.analyze_order_by(order_by); } + + for (name, previous, columns) in hidden.into_iter().rev() { + self.ctx.cte_definitions.insert(name.clone(), previous); + match columns { + Some(columns) => { + self.ctx.aliased_subquery_columns.insert(name, columns); + } + None => { + self.ctx.aliased_subquery_columns.remove(&name); + } + } + } } fn visit_set_expr(&mut self, set_expr: &SetExpr) { diff --git a/crates/flowscope-core/tests/lineage_engine.rs b/crates/flowscope-core/tests/lineage_engine.rs index 27cea89a..de68e63e 100644 --- a/crates/flowscope-core/tests/lineage_engine.rs +++ b/crates/flowscope-core/tests/lineage_engine.rs @@ -13176,3 +13176,224 @@ fn star_rename_returns_the_column_under_its_new_name() { ); } } + +#[test] +fn ctes_sharing_a_name_in_two_nested_withs_keep_their_own_columns() { + // dbt inlines each ephemeral model as a CTE holding its own WITH, and two + // of them often name an inner CTE alike. + let sql = r#" + WITH a AS ( + WITH final AS (SELECT k, ts FROM sat_a) SELECT k, ts FROM final + ), + b AS ( + WITH final AS (SELECT k, ts FROM sat_b) SELECT k, ts FROM final + ) + SELECT a.k, a.ts AS a_ts, b.ts AS b_ts + FROM a + JOIN b ON a.k = b.k + "#; + let schema = schema_of(vec![ + schema_table(None, None, "sat_a", &["k", "ts"]), + schema_table(None, None, "sat_b", &["k", "ts"]), + ]); + + let result = run_analysis(sql, Dialect::Snowflake, Some(schema)); + let stmt = first_statement(&result); + assert_eq!( + base_sources_of(&stmt, "a_ts"), + expected_sources(&["sat_a.ts"]) + ); + assert_eq!( + base_sources_of(&stmt, "b_ts"), + expected_sources(&["sat_b.ts"]) + ); +} + +#[test] +fn ctes_sharing_a_name_in_two_derived_tables_keep_their_own_columns() { + let sql = r#" + SELECT p.ts AS paid_ts, q.ts AS open_ts + FROM (WITH z AS (SELECT ts FROM payments) SELECT ts FROM z) AS p + CROSS JOIN (WITH z AS (SELECT ts FROM invoices) SELECT ts FROM z) AS q + "#; + let schema = schema_of(vec![ + schema_table(None, None, "payments", &["ts"]), + schema_table(None, None, "invoices", &["ts"]), + ]); + + let result = run_analysis(sql, Dialect::Snowflake, Some(schema)); + let stmt = first_statement(&result); + assert_eq!( + base_sources_of(&stmt, "paid_ts"), + expected_sources(&["payments.ts"]) + ); + assert_eq!( + base_sources_of(&stmt, "open_ts"), + expected_sources(&["invoices.ts"]) + ); +} + +#[test] +fn a_cte_shadowed_in_a_nested_with_is_still_itself_outside_it() { + let sql = r#" + WITH final AS (SELECT ts FROM sat_a), + inner_q AS ( + WITH final AS (SELECT ts FROM sat_b) SELECT ts FROM final + ) + SELECT final.ts AS outer_ts, inner_q.ts AS inner_ts + FROM final + CROSS JOIN inner_q + "#; + let schema = schema_of(vec![ + schema_table(None, None, "sat_a", &["ts"]), + schema_table(None, None, "sat_b", &["ts"]), + ]); + + let result = run_analysis(sql, Dialect::Snowflake, Some(schema)); + let stmt = first_statement(&result); + assert_eq!( + base_sources_of(&stmt, "outer_ts"), + expected_sources(&["sat_a.ts"]) + ); + assert_eq!( + base_sources_of(&stmt, "inner_ts"), + expected_sources(&["sat_b.ts"]) + ); +} + +#[test] +fn ctes_sharing_a_name_expand_their_own_stars() { + let sql = r#" + WITH a AS ( + WITH final AS (SELECT k, ts FROM sat_a) SELECT * FROM final + ), + b AS ( + WITH final AS (SELECT k, amount FROM sat_b) SELECT * FROM final + ) + SELECT a.ts AS a_ts, b.amount AS b_amount + FROM a + JOIN b ON a.k = b.k + "#; + let schema = schema_of(vec![ + schema_table(None, None, "sat_a", &["k", "ts"]), + schema_table(None, None, "sat_b", &["k", "amount"]), + ]); + + let result = run_analysis(sql, Dialect::Snowflake, Some(schema)); + let stmt = first_statement(&result); + assert_eq!( + base_sources_of(&stmt, "a_ts"), + expected_sources(&["sat_a.ts"]) + ); + assert_eq!( + base_sources_of(&stmt, "b_amount"), + expected_sources(&["sat_b.amount"]) + ); +} + +#[test] +fn ctes_sharing_a_name_infer_their_own_columns_without_a_schema() { + let sql = r#" + WITH a AS ( + WITH final AS (SELECT * FROM sat_a) SELECT * FROM final + ), + b AS ( + WITH final AS (SELECT * FROM sat_b) SELECT * FROM final + ) + SELECT a.ts AS a_ts, b.amount AS b_amount + FROM a + CROSS JOIN b + "#; + + let result = run_analysis(sql, Dialect::Snowflake, None); + let stmt = first_statement(&result); + assert_eq!( + base_sources_of(&stmt, "a_ts"), + expected_sources(&["sat_a.ts"]) + ); + assert_eq!( + base_sources_of(&stmt, "b_amount"), + expected_sources(&["sat_b.amount"]) + ); +} + +#[test] +fn a_filter_on_one_of_two_ctes_sharing_a_name_stays_on_it() { + let sql = r#" + WITH a AS ( + WITH final AS (SELECT k, ts FROM sat_a) SELECT k, ts FROM final + ), + b AS ( + WITH final AS (SELECT k, ts FROM sat_b) SELECT k, ts FROM final WHERE final.k = 1 + ) + SELECT a.ts AS a_ts, b.ts AS b_ts + FROM a + JOIN b ON a.k = b.k + "#; + let schema = schema_of(vec![ + schema_table(None, None, "sat_a", &["k", "ts"]), + schema_table(None, None, "sat_b", &["k", "ts"]), + ]); + + let result = run_analysis(sql, Dialect::Snowflake, Some(schema)); + let stmt = first_statement(&result); + let finals: Vec<&Node> = stmt + .nodes + .iter() + .copied() + .filter(|node| node.node_type == NodeType::Cte && &*node.label == "final") + .collect(); + assert_eq!(finals.len(), 2, "one node per CTE"); + let filtered: Vec = finals.iter().map(|node| node.filters.len()).collect(); + assert_eq!( + filtered.iter().filter(|n| **n > 0).count(), + 1, + "only the filtered one carries the predicate: {filtered:?}" + ); +} + +#[test] +fn an_unqualified_column_of_an_outer_cte_is_not_read_from_an_inner_namesake() { + let sql = r#" + WITH final AS (SELECT ts FROM sat_a) + SELECT ts AS outer_ts, d.amount AS inner_amount + FROM final + CROSS JOIN (WITH final AS (SELECT amount FROM sat_b) SELECT amount FROM final) AS d + "#; + let schema = schema_of(vec![ + schema_table(None, None, "sat_a", &["ts"]), + schema_table(None, None, "sat_b", &["amount"]), + ]); + + let result = run_analysis(sql, Dialect::Snowflake, Some(schema)); + let stmt = first_statement(&result); + assert_eq!( + base_sources_of(&stmt, "outer_ts"), + expected_sources(&["sat_a.ts"]) + ); + assert_eq!( + base_sources_of(&stmt, "inner_amount"), + expected_sources(&["sat_b.amount"]) + ); +} + +#[test] +fn an_outer_cte_named_like_an_inlined_model_s_inner_one_reads_through_it() { + // dbt inlines an ephemeral model holding its own `final`, and the model + // reading it names its last CTE `final` too. + let sql = r#" + WITH __dbt__cte__a AS ( + WITH final AS (SELECT k, ts FROM t) SELECT * FROM final + ), + final AS (SELECT * FROM __dbt__cte__a) + SELECT ts AS out_ts FROM final + "#; + let schema = schema_of(vec![schema_table(None, None, "t", &["k", "ts"])]); + + let result = run_analysis(sql, Dialect::Snowflake, Some(schema)); + let stmt = first_statement(&result); + assert_eq!( + base_sources_of(&stmt, "out_ts"), + expected_sources(&["t.ts"]) + ); +}