diff --git a/crates/flowscope-core/src/analyzer/query.rs b/crates/flowscope-core/src/analyzer/query.rs index de8e84fa..c9000270 100644 --- a/crates/flowscope-core/src/analyzer/query.rs +++ b/crates/flowscope-core/src/analyzer/query.rs @@ -7,6 +7,7 @@ use super::context::{ ColumnRef, OutputColumn, PendingWildcard, StatementContext, DBT_MODEL_SINK_METADATA_KEY, }; +use super::expression::ExpressionAnalyzer; use super::helpers::{ generate_column_node_id, generate_edge_id, generate_node_id, normalize_schema_type, }; @@ -557,9 +558,11 @@ impl<'a> Analyzer<'a> { target_node: Option<&str>, options: &ast::WildcardAdditionalOptions, ) { - // What `* EXCLUDE (...)` or `* EXCEPT (...)` leaves out, and what - // `* RENAME (old AS new)` returns under another name, compared as the - // expanded columns' own names are. REPLACE and ILIKE are not read. + // What `* EXCLUDE (...)` or `* EXCEPT (...)` leaves out, what + // `* RENAME (old AS new)` returns under another name, and what + // `* REPLACE (expr AS col)` computes in a column's place, compared as + // the expanded columns' own names are. `* ILIKE 'pattern'` keeps only + // the columns whose name the pattern matches, without regard to case. let mut excluded: HashSet = HashSet::new(); match &options.opt_exclude { Some(ast::ExcludeSelectItem::Single(ident)) => { @@ -595,6 +598,21 @@ impl<'a> Analyzer<'a> { ) }) .collect(); + let replaced: HashMap = options + .opt_replace + .iter() + .flat_map(|replace| &replace.items) + .map(|item| { + ( + self.normalize_identifier(&item.column_name.to_string()), + &item.expr, + ) + }) + .collect(); + let ilike = options + .opt_ilike + .as_ref() + .map(|ilike| ilike.pattern.as_str()); // Resolve wildcard sources as (canonical, qualifier) pairs so repeated // relation instances in self-joins are expanded independently. @@ -659,10 +677,39 @@ impl<'a> Analyzer<'a> { if excluded.contains(&normalized) { continue; } + if ilike + .is_some_and(|pattern| !like_matches_ignoring_case(pattern, &col_info.name)) + { + continue; + } let name = renamed .get(&normalized) .cloned() .unwrap_or_else(|| col_info.name.clone()); + // The replacing expression stands where the column was, and + // reads what it reads, the column itself or any other. + if let Some(expr) = replaced.get(&normalized) { + let (sources, aggregation) = { + let mut ea = ExpressionAnalyzer::new(self, ctx); + ( + ea.extract_column_refs_with_warning(expr), + ea.detect_aggregation(expr), + ) + }; + self.add_output_column_with_aggregation( + ctx, + OutputColumnParams { + name, + sources, + expression: Some(expr.to_string()), + data_type: None, + target_node: target_node.map(str::to_string), + approximate: false, + aggregation, + }, + ); + continue; + } let sources = vec![ColumnRef { table: Some(source_qualifier.clone()), column: col_info.name.clone(), @@ -1753,3 +1800,38 @@ impl<'a> Analyzer<'a> { Some(col_node_id) } } + +/// SQL's `LIKE`, without regard to case: `%` matches any run of characters and +/// `_` any one. What `* ILIKE 'pattern'` keeps of a star. +fn like_matches_ignoring_case(pattern: &str, name: &str) -> bool { + let pattern: Vec = pattern.to_lowercase().chars().collect(); + let name: Vec = name.to_lowercase().chars().collect(); + // Positions in the pattern reachable after each character of the name. + let mut at = vec![false; pattern.len() + 1]; + at[0] = true; + let close = |at: &mut Vec| { + for i in 0..pattern.len() { + if at[i] && pattern[i] == '%' { + at[i + 1] = true; + } + } + }; + close(&mut at); + for c in name { + let mut next = vec![false; pattern.len() + 1]; + for i in 0..pattern.len() { + if !at[i] { + continue; + } + match pattern[i] { + '%' => next[i] = true, + '_' => next[i + 1] = true, + p if p == c => next[i + 1] = true, + _ => {} + } + } + close(&mut next); + at = next; + } + at[pattern.len()] +} diff --git a/crates/flowscope-core/tests/lineage_engine.rs b/crates/flowscope-core/tests/lineage_engine.rs index 690abb07..99605eff 100644 --- a/crates/flowscope-core/tests/lineage_engine.rs +++ b/crates/flowscope-core/tests/lineage_engine.rs @@ -13550,3 +13550,64 @@ fn a_derived_table_inside_a_later_branch_keeps_its_own_names() { expected_sources(&["t1.k", "t2.k"]) ); } + +#[test] +fn star_ilike_keeps_the_columns_its_pattern_matches() { + for (sql, want) in [ + ("SELECT * ILIKE 'a%' FROM p", vec!["a"]), + ("SELECT * ILIKE '%B' FROM p", vec!["b"]), + ("SELECT * ILIKE '_' FROM p", vec!["a", "b", "c"]), + ("SELECT p.* ILIKE '%c%' FROM p", vec!["c"]), + ( + "WITH c AS (SELECT * ILIKE 'C' FROM p) SELECT * FROM c", + vec!["c"], + ), + ( + "WITH c AS (SELECT * ILIKE '%' FROM p) SELECT * FROM c", + vec!["a", "b", "c"], + ), + ] { + let result = run_analysis(sql, Dialect::Snowflake, Some(star_options_schema())); + let stmt = first_statement(&result); + assert_eq!(returned_columns(&stmt), want, "{sql}"); + } +} + +#[test] +fn star_replace_computes_the_column_it_replaces() { + // The replacing expression stands where the column was, keeps its place + // and name, and reads what it reads. + for (sql, sources) in [ + ("SELECT * REPLACE (UPPER(b) AS b) FROM p", vec!["p.b"]), + ( + "SELECT p.* REPLACE (b || c AS b) FROM p", + vec!["p.b", "p.c"], + ), + ( + "WITH r AS (SELECT * REPLACE (a + 1 AS a) FROM p) SELECT * FROM r", + vec!["p.a"], + ), + ] { + let result = run_analysis(sql, Dialect::Snowflake, Some(star_options_schema())); + let stmt = first_statement(&result); + assert_eq!(returned_columns(&stmt), vec!["a", "b", "c"], "{sql}"); + let replaced = if sql.contains("a + 1") { "a" } else { "b" }; + assert_eq!( + base_sources_of(&stmt, replaced), + expected_sources(&sources), + "{sql}" + ); + let expressions: Vec = stmt + .nodes + .iter() + .filter(|node| { + node.node_type == NodeType::Column && node.label.eq_ignore_ascii_case(replaced) + }) + .map(|node| node.expression.is_some()) + .collect(); + assert!( + expressions.contains(&true), + "{sql}: the replaced column carries its expression" + ); + } +}