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"]) + ); +}