diff --git a/sqllineage-python/sqllineage.pyi b/sqllineage-python/sqllineage.pyi index 04f025b..1f8bde5 100644 --- a/sqllineage-python/sqllineage.pyi +++ b/sqllineage-python/sqllineage.pyi @@ -17,7 +17,8 @@ class ColumnOrigin: Check ``kind`` to determine the variant: - ``"concrete"``: ``table`` and ``column`` are set. - - ``"ambiguous"``: ``column`` and ``candidates`` are set. + - ``"ambiguous"``: ``column`` and ``candidates`` are set. ``candidates`` + may be an empty list when the column is unresolved. - ``"wildcard"``: ``table`` is set. - ``"recursive"``: ``base_sources`` is set. """ diff --git a/sqllineage-python/src/lib.rs b/sqllineage-python/src/lib.rs index 337dfc1..60db414 100644 --- a/sqllineage-python/src/lib.rs +++ b/sqllineage-python/src/lib.rs @@ -111,6 +111,15 @@ impl PyColumnOrigin { "ColumnOrigin.wildcard({}.*)", self.table.as_ref().map_or("?", |t| &t.table), ), + "named_wildcard" => format!( + "ColumnOrigin.named_wildcard({}.{})", + self.table.as_ref().map_or("?", |t| &t.table), + self.column.as_deref().unwrap_or("?"), + ), + "source_free" => format!( + "ColumnOrigin.source_free({})", + self.column.as_deref().unwrap_or("?"), + ), "ambiguous" => format!( "ColumnOrigin.ambiguous({})", self.column.as_deref().unwrap_or("?"), @@ -144,6 +153,20 @@ fn convert_origin(o: &sqllineage_core::ColumnOrigin) -> PyColumnOrigin { candidates: None, base_sources: None, }, + sqllineage_core::ColumnOrigin::NamedWildcard { table, column } => PyColumnOrigin { + kind: "named_wildcard".into(), + table: Some(PyTableRef::from(table)), + column: Some(column.clone()), + candidates: None, + base_sources: None, + }, + sqllineage_core::ColumnOrigin::SourceFree { column } => PyColumnOrigin { + kind: "source_free".into(), + table: None, + column: Some(column.clone()), + candidates: None, + base_sources: None, + }, sqllineage_core::ColumnOrigin::Recursive { base_sources } => PyColumnOrigin { kind: "recursive".into(), table: None, @@ -151,6 +174,13 @@ fn convert_origin(o: &sqllineage_core::ColumnOrigin) -> PyColumnOrigin { candidates: None, base_sources: Some(base_sources.iter().map(convert_origin).collect()), }, + _ => PyColumnOrigin { + kind: "unknown".into(), + table: None, + column: None, + candidates: None, + base_sources: None, + }, } } @@ -210,6 +240,8 @@ struct PyLineageResult { tables: PyTableLineage, #[pyo3(get)] columns: Vec, + #[pyo3(get)] + has_unresolved_stars: bool, } #[pymethods] @@ -282,22 +314,15 @@ fn analyze( catalog: Option>, normalize_case: bool, ) -> PyResult> { - let d = match dialect.to_lowercase().as_str() { - "generic" => sqllineage_core::Dialect::Generic, - "ansi" => sqllineage_core::Dialect::Ansi, - "postgresql" | "postgres" => sqllineage_core::Dialect::PostgreSql, - "mysql" => sqllineage_core::Dialect::MySql, - "hive" => sqllineage_core::Dialect::Hive, - "databricks" => sqllineage_core::Dialect::Databricks, - "snowflake" => sqllineage_core::Dialect::Snowflake, - "bigquery" => sqllineage_core::Dialect::BigQuery, - other => { + let d = match parse_dialect_name(dialect) { + Some(dialect) => dialect, + None => { + let normalized = dialect.to_lowercase(); return Err(pyo3::exceptions::PyValueError::new_err(format!( - "unknown dialect: '{other}'" + "unknown dialect: '{normalized}'" ))); } }; - let catalog_box: Option> = catalog.map(|obj| Box::new(PyCatalog { obj }) as Box); @@ -329,10 +354,58 @@ fn analyze( transform: convert_transform(&m.transform).into(), }) .collect(), + has_unresolved_stars: result.columns.has_unresolved_stars, }) .collect()) } +fn parse_dialect_name(name: &str) -> Option { + match name.to_lowercase().as_str() { + "generic" => Some(sqllineage_core::Dialect::Generic), + "ansi" => Some(sqllineage_core::Dialect::Ansi), + "postgresql" | "postgres" => Some(sqllineage_core::Dialect::PostgreSql), + "mysql" => Some(sqllineage_core::Dialect::MySql), + "hive" => Some(sqllineage_core::Dialect::Hive), + "databricks" => Some(sqllineage_core::Dialect::Databricks), + "snowflake" => Some(sqllineage_core::Dialect::Snowflake), + "bigquery" => Some(sqllineage_core::Dialect::BigQuery), + "duckdb" | "duck_db" => Some(sqllineage_core::Dialect::DuckDb), + "redshift" => Some(sqllineage_core::Dialect::Redshift), + "trino" => Some(sqllineage_core::Dialect::Trino), + "spark" | "spark2" | "sparksql" => Some(sqllineage_core::Dialect::Spark), + "clickhouse" | "click_house" => Some(sqllineage_core::Dialect::ClickHouse), + "sqlite" => Some(sqllineage_core::Dialect::SQLite), + "mssql" | "ms_sql" | "tsql" | "t-sql" | "sqlserver" => { + Some(sqllineage_core::Dialect::MsSql) + } + _ => None, + } +} + +#[cfg(test)] +mod tests { + use super::parse_dialect_name; + use sqllineage_core::Dialect; + + #[test] + fn parses_added_dialect_names_and_tsql_aliases() { + for (name, expected) in [ + ("duckdb", Dialect::DuckDb), + ("redshift", Dialect::Redshift), + ("trino", Dialect::Trino), + ("spark", Dialect::Spark), + ("clickhouse", Dialect::ClickHouse), + ("sqlite", Dialect::SQLite), + ("tsql", Dialect::MsSql), + ] { + assert!( + matches!(parse_dialect_name(name), Some(actual) if std::mem::discriminant(&actual) == std::mem::discriminant(&expected)) + ); + } + assert!(parse_dialect_name("not-a-dialect").is_none()); + } +} + #[pymodule] mod sqllineage { #[pymodule_export] diff --git a/sqllineage/src/bin/sqllineage.rs b/sqllineage/src/bin/sqllineage.rs index 49445c6..22602ef 100644 --- a/sqllineage/src/bin/sqllineage.rs +++ b/sqllineage/src/bin/sqllineage.rs @@ -34,7 +34,7 @@ fn main() { Some(d) => d, None => { eprintln!( - "error: unknown dialect '{}'. valid: generic, ansi, postgresql, mysql, hive, databricks, snowflake, bigquery", + "error: unknown dialect '{}'. valid: generic, ansi, postgresql, mysql, hive, databricks, snowflake, bigquery, duckdb, redshift, trino, spark, clickhouse, sqlite, mssql/tsql", cli.dialect ); process::exit(1); @@ -82,6 +82,13 @@ fn parse_dialect(s: &str) -> Option { "databricks" => Some(Dialect::Databricks), "snowflake" => Some(Dialect::Snowflake), "bigquery" => Some(Dialect::BigQuery), + "duckdb" | "duck_db" => Some(Dialect::DuckDb), + "redshift" => Some(Dialect::Redshift), + "trino" => Some(Dialect::Trino), + "spark" | "spark2" | "sparksql" => Some(Dialect::Spark), + "clickhouse" | "click_house" => Some(Dialect::ClickHouse), + "sqlite" => Some(Dialect::SQLite), + "mssql" | "ms_sql" | "tsql" | "t-sql" | "sqlserver" => Some(Dialect::MsSql), _ => None, } } @@ -158,10 +165,13 @@ fn format_origin(origin: &ColumnOrigin) -> String { ColumnOrigin::Concrete { table, column } => format!("{table}.{column}"), ColumnOrigin::Ambiguous { column, .. } => format!("?{column}?"), ColumnOrigin::Wildcard { table } => format!("{table}.*"), + ColumnOrigin::NamedWildcard { table, column } => format!("{table}.*({column})"), + ColumnOrigin::SourceFree { column } => format!(""), ColumnOrigin::Recursive { base_sources } => { let inner: Vec = base_sources.iter().map(format_origin).collect(); format!("recursive({})", inner.join(", ")) } + _ => "".to_string(), } } @@ -205,3 +215,24 @@ fn format_dot(result: &AnalyzeResult, columns: bool) -> String { out.push('}'); out } + +#[cfg(test)] +mod tests { + use super::parse_dialect; + use sqllineage::Dialect; + + #[test] + fn parses_added_dialects_and_tsql_aliases() { + assert!(matches!(parse_dialect("duckdb"), Some(Dialect::DuckDb))); + assert!(matches!(parse_dialect("redshift"), Some(Dialect::Redshift))); + assert!(matches!(parse_dialect("trino"), Some(Dialect::Trino))); + assert!(matches!(parse_dialect("spark"), Some(Dialect::Spark))); + assert!(matches!( + parse_dialect("clickhouse"), + Some(Dialect::ClickHouse) + )); + assert!(matches!(parse_dialect("sqlite"), Some(Dialect::SQLite))); + assert!(matches!(parse_dialect("tsql"), Some(Dialect::MsSql))); + assert!(parse_dialect("not-a-dialect").is_none()); + } +} diff --git a/sqllineage/src/build/expr.rs b/sqllineage/src/build/expr.rs index 2b2cb3e..97e37e1 100644 --- a/sqllineage/src/build/expr.rs +++ b/sqllineage/src/build/expr.rs @@ -1,30 +1,50 @@ -use sqlparser::ast::{self, Expr, FunctionArguments, WindowType}; +use sqlparser::ast::{self, AccessExpr, Expr, FunctionArguments, Subscript, WindowType}; use crate::build::LineageBuilder; -use crate::build::select::split_compound; +use crate::build::function_semantics::{self, ArgumentSemantic}; use crate::graph::edge::EdgeKind; use crate::graph::node::NodeId; -use crate::graph::scope::ScopeKind; +use crate::graph::scope::{Binding, ScopeKind}; +use crate::types::TableRef; impl LineageBuilder { pub(crate) fn collect_ancestors(&mut self, expr: &Expr) -> Vec { match expr { Expr::Identifier(ident) => { - let node = self + let binding = self .graph - .add_unqualified(ident.value.clone(), self.current_scope); + .scopes + .lookup(self.current_scope, &ident.value) + .cloned(); + if binding.as_ref().is_some_and(|binding| { + self.dialect.supports_relation_alias_row_value() + && matches!( + binding, + Binding::Table(_) | Binding::Cte(_) | Binding::DerivedTable(_) + ) + }) { + return vec![self.graph.add_row_value_candidate( + ident.value.clone(), + self.current_scope, + binding, + )]; + } + let binding = + binding.filter(|binding| matches!(binding, Binding::VirtualSource(_))); + let node = self.graph.add_unqualified_with_binding( + ident.value.clone(), + self.current_scope, + binding, + ); vec![node] } - Expr::CompoundIdentifier(parts) => { - let (qualifier, column) = split_compound(parts); - let node = self - .graph - .add_ref(column, Some(qualifier), self.current_scope); - vec![node] - } + Expr::CompoundIdentifier(parts) => self.collect_compound_identifier_ancestors(parts), - Expr::Value(_) | Expr::TypedString { .. } | Expr::Wildcard(..) | Expr::QualifiedWildcard(..) => vec![], + Expr::Value(_) + | Expr::TypedString { .. } + | Expr::Wildcard(..) + | Expr::QualifiedWildcard(..) => vec![], Expr::Cast { expr, .. } | Expr::Nested(expr) @@ -49,7 +69,12 @@ impl LineageBuilder { Expr::Extract { expr, .. } => self.collect_ancestors(expr), - Expr::Trim { expr, trim_what, trim_characters, .. } => { + Expr::Trim { + expr, + trim_what, + trim_characters, + .. + } => { let mut v = self.collect_ancestors(expr); if let Some(what) = trim_what { v.extend(self.collect_ancestors(what)); @@ -62,7 +87,12 @@ impl LineageBuilder { v } - Expr::Substring { expr, substring_from, substring_for, .. } => { + Expr::Substring { + expr, + substring_from, + substring_for, + .. + } => { let mut v = self.collect_ancestors(expr); if let Some(from) = substring_from { v.extend(self.collect_ancestors(from)); @@ -73,7 +103,13 @@ impl LineageBuilder { v } - Expr::Overlay { expr, overlay_what, overlay_from, overlay_for, .. } => { + Expr::Overlay { + expr, + overlay_what, + overlay_from, + overlay_for, + .. + } => { let mut v = self.collect_ancestors(expr); v.extend(self.collect_ancestors(overlay_what)); v.extend(self.collect_ancestors(overlay_from)); @@ -89,17 +125,36 @@ impl LineageBuilder { v } - Expr::AtTimeZone { timestamp, time_zone } => { + Expr::AtTimeZone { + timestamp, + time_zone, + } => { let mut v = self.collect_ancestors(timestamp); v.extend(self.collect_ancestors(time_zone)); v } Expr::BinaryOp { left, right, .. } - | Expr::Like { expr: left, pattern: right, .. } - | Expr::ILike { expr: left, pattern: right, .. } - | Expr::SimilarTo { expr: left, pattern: right, .. } - | Expr::RLike { expr: left, pattern: right, .. } + | Expr::Like { + expr: left, + pattern: right, + .. + } + | Expr::ILike { + expr: left, + pattern: right, + .. + } + | Expr::SimilarTo { + expr: left, + pattern: right, + .. + } + | Expr::RLike { + expr: left, + pattern: right, + .. + } | Expr::IsDistinctFrom(left, right) | Expr::IsNotDistinctFrom(left, right) => { let mut v = self.collect_ancestors(left); @@ -113,7 +168,9 @@ impl LineageBuilder { v } - Expr::InUnnest { expr, array_expr, .. } => { + Expr::InUnnest { + expr, array_expr, .. + } => { let mut v = self.collect_ancestors(expr); v.extend(self.collect_ancestors(array_expr)); v @@ -158,21 +215,38 @@ impl LineageBuilder { v } - Expr::CompoundFieldAccess { root, .. } => self.collect_ancestors(root), + Expr::CompoundFieldAccess { root, access_chain } => { + self.collect_compound_field_ancestors(root, access_chain) + } Expr::JsonAccess { value, .. } => self.collect_ancestors(value), Expr::Function(func) => { let mut ancestors = Vec::new(); if let FunctionArguments::List(list) = &func.args { - for arg in &list.args { - match arg { + let semantics = function_semantics::classify_function(self.dialect, func) + .map(|signature| signature.arguments); + for (index, arg) in list.args.iter().enumerate() { + let arg_expr = match arg { ast::FunctionArg::Unnamed(arg_expr) | ast::FunctionArg::Named { arg: arg_expr, .. } - | ast::FunctionArg::ExprNamed { arg: arg_expr, .. } => { - if let ast::FunctionArgExpr::Expr(e) = arg_expr { - ancestors.extend(self.collect_ancestors(e)); - } - } + | ast::FunctionArg::ExprNamed { arg: arg_expr, .. } => arg_expr, + }; + let ast::FunctionArgExpr::Expr(expr) = arg_expr else { + continue; + }; + let is_static_syntax = semantics + .and_then(|semantics| semantics.get(index)) + .is_some_and(|semantic| { + matches!( + semantic, + ArgumentSemantic::DatePart(grammar) + if function_semantics::expression_is_static_date_part( + expr, *grammar + ) + ) + }); + if !is_static_syntax { + ancestors.extend(self.collect_ancestors(expr)); } } } @@ -192,7 +266,12 @@ impl LineageBuilder { ancestors } - Expr::Case { operand, conditions, else_result, .. } => { + Expr::Case { + operand, + conditions, + else_result, + .. + } => { let mut v = Vec::new(); if let Some(op) = operand { v.extend(self.collect_ancestors(op)); @@ -229,7 +308,9 @@ impl LineageBuilder { vec![] } - Expr::Between { expr, low, high, .. } => { + Expr::Between { + expr, low, high, .. + } => { let mut v = self.collect_ancestors(expr); v.extend(self.collect_ancestors(low)); v.extend(self.collect_ancestors(high)); @@ -251,11 +332,221 @@ impl LineageBuilder { | Expr::Interval(_) | Expr::Lambda(_) | Expr::MatchAgainst { .. } => vec![], + } + } + + /// Collect the physical column at the root of a structured access chain. + /// + /// `base.items[0]` is ambiguous at the syntax level: `base` can be a + /// visible relation binding, in which case `items` is its physical column, + /// or it can be an unqualified top-level column (`payload.items[0]`). The + /// scope binding, rather than rendered SQL text or dialect-specific names, + /// is the structural distinction between those cases. + fn collect_compound_field_ancestors( + &mut self, + root: &Expr, + access_chain: &[AccessExpr], + ) -> Vec { + let mut ancestors = match (root, access_chain.first()) { + (Expr::Identifier(binding_name), Some(AccessExpr::Dot(Expr::Identifier(field)))) + if self + .graph + .scopes + .lookup(self.current_scope, &binding_name.value) + .is_some() => + { + let binding = self + .graph + .scopes + .lookup(self.current_scope, &binding_name.value) + .cloned(); + vec![self.add_bound_field_ancestor( + binding_name.value.clone(), + field.value.clone(), + binding, + )] + } + (Expr::Identifier(column), Some(AccessExpr::Dot(Expr::Identifier(_)))) => { + let binding = self + .graph + .scopes + .lookup(self.current_scope, &column.value) + .cloned(); + let binding = + binding.filter(|binding| matches!(binding, Binding::VirtualSource(_))); + vec![self.graph.add_unqualified_with_binding( + column.value.clone(), + self.current_scope, + binding, + )] + } + _ => self.collect_ancestors(root), + }; + + for access in access_chain { + if let AccessExpr::Subscript(subscript) = access { + ancestors.extend(self.collect_subscript_ancestors(subscript)); + } + } + ancestors + } + + /// Resolve a plain dotted identifier by separating its relation binding + /// from the top-level physical column. The parser represents both + /// `alias.column` and `alias.struct.field` as a flat compound identifier, + /// so rendering all but the final component as one qualifier loses the + /// distinction between a relation name and a nested field path. + fn collect_compound_identifier_ancestors( + &mut self, + parts: &[sqlparser::ast::Ident], + ) -> Vec { + if let Some((prefix_len, binding)) = self.find_compound_binding(parts) { + let qualifier = parts[..prefix_len] + .iter() + .map(|part| part.value.as_str()) + .collect::>() + .join("."); + let column = parts[prefix_len].value.clone(); + return vec![self.add_bound_field_ancestor(qualifier, column, Some(binding))]; + } + + // With no visible relation bindings, preserve the traditional + // qualified-reference fallback. This is used by callers that feed + // already-qualified expressions without a FROM clause (for example + // `orders.id`): the qualifier is still a physical relation name, + // rather than an unqualified struct root. + if self + .graph + .scopes + .visible_bindings(self.current_scope) + .is_empty() + && parts.len() >= 2 + { + let relation_parts = &parts[..parts.len() - 1]; + let qualifier = parts[..parts.len() - 1] + .iter() + .map(|part| part.value.as_str()) + .collect::>() + .join("."); + let binding = match relation_parts { + [table] => Some(Binding::Table(TableRef::new(table.value.clone()))), + [schema, table] => Some(Binding::Table(TableRef::with_schema( + schema.value.clone(), + table.value.clone(), + ))), + [catalog, schema, table] => Some(Binding::Table(TableRef { + catalog: Some(catalog.value.clone()), + schema: Some(schema.value.clone()), + table: table.value.clone(), + })), + // Keep the legacy display-only fallback for an unsupported + // number of relation components. Standard SQL relation + // names are at most catalog.schema.table, and this branch + // avoids inventing a lossy structured interpretation beyond + // that shape. + _ => None, + }; + return vec![self.graph.add_ref_with_binding( + parts[parts.len() - 1].value.clone(), + Some(qualifier), + self.current_scope, + binding, + )]; + } + + // No relation prefix was found: `struct.field` is an unqualified + // top-level column followed by a nested field path. Only the + // top-level column can be represented by the public ColumnOrigin API. + let column = parts[0].value.clone(); + let binding = self + .graph + .scopes + .lookup(self.current_scope, &column) + .cloned() + .filter(|binding| matches!(binding, Binding::VirtualSource(_))); + vec![ + self.graph + .add_unqualified_with_binding(column, self.current_scope, binding), + ] + } + + /// Find the longest visible relation prefix in a compound identifier. + /// + /// A one-component prefix is a SQL alias. Longer prefixes are matched + /// against the physical parts of a table binding, allowing references such + /// as `catalog.schema.table.column` without turning the relation into a + /// single quoted string containing dots. + fn find_compound_binding(&self, parts: &[sqlparser::ast::Ident]) -> Option<(usize, Binding)> { + let visible = self.graph.scopes.visible_bindings(self.current_scope); + (1..parts.len()).rev().find_map(|prefix_len| { + let prefix = parts[..prefix_len] + .iter() + .map(|part| part.value.as_str()) + .collect::>(); + + // Aliases are single identifiers and therefore only match the + // first component of a compound identifier. + if prefix_len == 1 + && let Some((_, binding)) = visible.iter().find(|(name, _)| name == prefix[0]) + { + return Some((prefix_len, binding.clone())); + } + + visible.iter().find_map(|(_, binding)| { + let Binding::Table(table) = binding else { + return None; + }; + (table_parts(table) == prefix).then(|| (prefix_len, binding.clone())) + }) + }) + } + + fn add_bound_field_ancestor( + &mut self, + qualifier: String, + column: String, + binding: Option, + ) -> NodeId { + self.graph + .add_ref_with_binding(column, Some(qualifier), self.current_scope, binding) + } + fn collect_subscript_ancestors(&mut self, subscript: &Subscript) -> Vec { + match subscript { + Subscript::Index { index } => self.collect_ancestors(index), + Subscript::Slice { + lower_bound, + upper_bound, + stride, + } => { + let mut ancestors = Vec::new(); + if let Some(lower) = lower_bound { + ancestors.extend(self.collect_ancestors(lower)); + } + if let Some(upper) = upper_bound { + ancestors.extend(self.collect_ancestors(upper)); + } + if let Some(step) = stride { + ancestors.extend(self.collect_ancestors(step)); + } + ancestors + } } } } +fn table_parts(table: &TableRef) -> Vec<&str> { + let mut parts = Vec::with_capacity(3); + if let Some(catalog) = &table.catalog { + parts.push(catalog.as_str()); + } + if let Some(schema) = &table.schema { + parts.push(schema.as_str()); + } + parts.push(table.table.as_str()); + parts +} + pub(crate) fn determine_edge_kind(expr: &Expr) -> EdgeKind { match expr { Expr::Identifier(_) | Expr::CompoundIdentifier(_) | Expr::Value(_) => EdgeKind::Direct, diff --git a/sqllineage/src/build/function_semantics.rs b/sqllineage/src/build/function_semantics.rs new file mode 100644 index 0000000..b003ee2 --- /dev/null +++ b/sqllineage/src/build/function_semantics.rs @@ -0,0 +1,1252 @@ +//! Compatibility semantics for temporal function arguments. +//! +//! sqlparser 0.62 retains date-part arguments of generic `Function` nodes as +//! ordinary `Expr::Identifier` values, which creates a semantic gap for +//! lineage. This small, profile-scoped layer bridges that gap until sqlparser +//! provides typed temporal function arguments; at that point these profiles +//! and the registry should be removed or migrated to the typed AST. The +//! approach follows the direction discussed in [PR 1191](https://github.com/apache/datafusion-sqlparser-rs/pull/1191) +//! (including `BigQuery` `WEEK(MONDAY)`) and [issue 1983](https://github.com/apache/datafusion-sqlparser-rs/issues/1983) +//! / [PR 2030](https://github.com/apache/datafusion-sqlparser-rs/pull/2030) +//! (dialect-specific `EXTRACT` date-part parsing). +//! +//! The tables below are an exception inventory for official static syntax that +//! sqlparser represents as generic identifiers, not a registry of every +//! temporal function. Unknown functions and dynamic expressions retain the +//! generic all-values fallback. + +use sqlparser::ast::{ + Expr, Function, FunctionArg, FunctionArgExpr, FunctionArguments, ObjectNamePart, +}; + +use crate::types::Dialect; + +/// The lineage role of an argument in a known function grammar. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ArgumentSemantic { + ValueExpression, + DatePart(DatePartGrammar), +} + +/// A dialect-specific function signature. This deliberately lives outside +/// sqlparser's AST: the AST describes syntax, while this layer describes +/// which syntax contributes data lineage. +#[derive(Debug, Clone, Copy)] +pub(crate) struct FunctionSignature { + pub names: &'static [&'static str], + pub arity: usize, + pub arguments: &'static [ArgumentSemantic], +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct DatePartGrammar { + profile: &'static DatePartProfile, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct DatePartToken { + canonical: &'static str, + aliases: &'static [&'static str], +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct DatePartProfile { + token_groups: &'static [&'static [DatePartToken]], + allows_weekday_modifier: bool, +} + +const BQ_VALUE_PART: &[ArgumentSemantic] = &[ + ArgumentSemantic::ValueExpression, + ArgumentSemantic::DatePart(BIGQUERY_DATE_PART), +]; +const BQ_VALUE_PART_VALUE: &[ArgumentSemantic] = &[ + ArgumentSemantic::ValueExpression, + ArgumentSemantic::DatePart(BIGQUERY_DATE_PART), + ArgumentSemantic::ValueExpression, +]; +const BQ_VALUE_VALUE_PART: &[ArgumentSemantic] = &[ + ArgumentSemantic::ValueExpression, + ArgumentSemantic::ValueExpression, + ArgumentSemantic::DatePart(BIGQUERY_DATE_PART), +]; +const BQ_VALUE: &[ArgumentSemantic] = &[ArgumentSemantic::ValueExpression]; +const SNOW_PART_VALUE_VALUE: &[ArgumentSemantic] = &[ + ArgumentSemantic::DatePart(SNOWFLAKE_DATE_PART), + ArgumentSemantic::ValueExpression, + ArgumentSemantic::ValueExpression, +]; +const SNOW_PART_VALUE: &[ArgumentSemantic] = &[ + ArgumentSemantic::DatePart(SNOWFLAKE_DATE_PART), + ArgumentSemantic::ValueExpression, +]; +const SNOW_VALUE_PART: &[ArgumentSemantic] = &[ + ArgumentSemantic::ValueExpression, + ArgumentSemantic::DatePart(SNOWFLAKE_DATE_PART), +]; +const DB_PART_VALUE_VALUE: &[ArgumentSemantic] = &[ + ArgumentSemantic::DatePart(DATABRICKS_DATE_PART), + ArgumentSemantic::ValueExpression, + ArgumentSemantic::ValueExpression, +]; +const DB_DIFF_PART_VALUE_VALUE: &[ArgumentSemantic] = &[ + ArgumentSemantic::DatePart(DATABRICKS_DIFF_DATE_PART), + ArgumentSemantic::ValueExpression, + ArgumentSemantic::ValueExpression, +]; +const MYSQL_PART_VALUE_VALUE: &[ArgumentSemantic] = &[ + ArgumentSemantic::DatePart(MYSQL_DATE_PART), + ArgumentSemantic::ValueExpression, + ArgumentSemantic::ValueExpression, +]; +const REDSHIFT_PART_VALUE_VALUE: &[ArgumentSemantic] = &[ + ArgumentSemantic::DatePart(REDSHIFT_ADD_DIFF_DATE_PART), + ArgumentSemantic::ValueExpression, + ArgumentSemantic::ValueExpression, +]; +const REDSHIFT_PART_VALUE: &[ArgumentSemantic] = &[ + ArgumentSemantic::DatePart(REDSHIFT_DATE_PART), + ArgumentSemantic::ValueExpression, +]; +const MSSQL_PART_VALUE_VALUE: &[ArgumentSemantic] = &[ + ArgumentSemantic::DatePart(MSSQL_DATEADD_DATE_PART), + ArgumentSemantic::ValueExpression, + ArgumentSemantic::ValueExpression, +]; +const MSSQL_DATEDIFF_PART_VALUE_VALUE: &[ArgumentSemantic] = &[ + ArgumentSemantic::DatePart(MSSQL_DATEDIFF_DATE_PART), + ArgumentSemantic::ValueExpression, + ArgumentSemantic::ValueExpression, +]; +const MSSQL_PART_VALUE: &[ArgumentSemantic] = &[ + ArgumentSemantic::DatePart(MSSQL_DATEPART_DATE_PART), + ArgumentSemantic::ValueExpression, +]; +const MSSQL_DATETRUNC_PART_VALUE: &[ArgumentSemantic] = &[ + ArgumentSemantic::DatePart(MSSQL_DATETRUNC_DATE_PART), + ArgumentSemantic::ValueExpression, +]; +const MSSQL_DATE_BUCKET_PART_VALUE_VALUE: &[ArgumentSemantic] = &[ + ArgumentSemantic::DatePart(MSSQL_DATE_BUCKET_DATE_PART), + ArgumentSemantic::ValueExpression, + ArgumentSemantic::ValueExpression, +]; +const MSSQL_DATE_BUCKET_PART_VALUE_VALUE_VALUE: &[ArgumentSemantic] = &[ + ArgumentSemantic::DatePart(MSSQL_DATE_BUCKET_DATE_PART), + ArgumentSemantic::ValueExpression, + ArgumentSemantic::ValueExpression, + ArgumentSemantic::ValueExpression, +]; + +const EMPTY_ALIASES: &[&str] = &[]; +const BQ_PARTS: &[DatePartToken] = &[ + DatePartToken { + canonical: "MICROSECOND", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "MILLISECOND", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "SECOND", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "MINUTE", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "HOUR", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "DAY", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "WEEK", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "ISOWEEK", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "MONTH", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "QUARTER", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "YEAR", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "ISOYEAR", + aliases: EMPTY_ALIASES, + }, +]; +const SNOWFLAKE_PARTS: &[DatePartToken] = &[ + DatePartToken { + canonical: "YEAR", + aliases: &["y", "yy", "yyy", "yyyy", "yr", "years", "yrs"], + }, + DatePartToken { + canonical: "QUARTER", + aliases: &["q", "qtr", "qtrs", "quarters"], + }, + DatePartToken { + canonical: "MONTH", + aliases: &["mm", "mon", "mons", "months"], + }, + DatePartToken { + canonical: "DAY", + aliases: &["d", "dd", "days", "dayofmonth"], + }, + DatePartToken { + canonical: "DAYOFWEEK", + aliases: &["weekday", "dow", "dw"], + }, + DatePartToken { + canonical: "DAYOFWEEKISO", + aliases: &["weekday_iso", "dow_iso", "dw_iso", "dayofweek_iso"], + }, + DatePartToken { + canonical: "DAYOFYEAR", + aliases: &["doy", "dy", "yearday"], + }, + DatePartToken { + canonical: "WEEK", + aliases: &["w", "wk", "ww", "weekofyear", "woy", "wy"], + }, + DatePartToken { + canonical: "WEEKISO", + aliases: &["isoweek", "week_iso", "weekofyeariso", "weekofyear_iso"], + }, + DatePartToken { + canonical: "HOUR", + aliases: &["h", "hh", "hr", "hours", "hrs"], + }, + DatePartToken { + canonical: "MINUTE", + aliases: &["m", "mi", "min", "minutes", "mins"], + }, + DatePartToken { + canonical: "SECOND", + aliases: &["s", "sec", "seconds", "secs"], + }, + DatePartToken { + canonical: "MILLISECOND", + aliases: &[ + "ms", + "msec", + "msecs", + "msecond", + "mseconds", + "millisec", + "millisecs", + "millisecon", + "milliseconds", + ], + }, + DatePartToken { + canonical: "MICROSECOND", + aliases: &[ + "us", + "usec", + "usecs", + "microsec", + "microsecs", + "usecond", + "useconds", + "microseconds", + ], + }, + DatePartToken { + canonical: "NANOSECOND", + aliases: &["ns", "nsec", "nanosec", "nsecond", "nseconds", "nanosecs"], + }, + DatePartToken { + canonical: "EPOCH", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "EPOCH_SECOND", + aliases: &["epoch_second", "epoch_seconds"], + }, + DatePartToken { + canonical: "EPOCH_MILLISECOND", + aliases: &["epoch_milliseconds"], + }, + DatePartToken { + canonical: "EPOCH_MICROSECOND", + aliases: &["epoch_microseconds"], + }, + DatePartToken { + canonical: "EPOCH_NANOSECOND", + aliases: &["epoch_nanoseconds"], + }, + DatePartToken { + canonical: "TIMEZONE_HOUR", + aliases: &["tzh"], + }, + DatePartToken { + canonical: "TIMEZONE_MINUTE", + aliases: &["tzm"], + }, + DatePartToken { + canonical: "DECADE", + aliases: &["dec", "decs", "decades"], + }, + DatePartToken { + canonical: "MILLENNIUM", + aliases: &["mil", "mils", "millenia"], + }, + DatePartToken { + canonical: "CENTURY", + aliases: &["c", "cent", "cents", "centuries"], + }, + DatePartToken { + canonical: "YEAROFWEEK", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "YEAROFWEEKISO", + aliases: EMPTY_ALIASES, + }, +]; +const MYSQL_PARTS: &[DatePartToken] = &[ + DatePartToken { + canonical: "MICROSECOND", + aliases: &["SQL_TSI_MICROSECOND"], + }, + DatePartToken { + canonical: "SECOND", + aliases: &["SQL_TSI_SECOND"], + }, + DatePartToken { + canonical: "MINUTE", + aliases: &["SQL_TSI_MINUTE"], + }, + DatePartToken { + canonical: "HOUR", + aliases: &["SQL_TSI_HOUR"], + }, + DatePartToken { + canonical: "DAY", + aliases: &["SQL_TSI_DAY"], + }, + DatePartToken { + canonical: "WEEK", + aliases: &["SQL_TSI_WEEK"], + }, + DatePartToken { + canonical: "MONTH", + aliases: &["SQL_TSI_MONTH"], + }, + DatePartToken { + canonical: "QUARTER", + aliases: &["SQL_TSI_QUARTER"], + }, + DatePartToken { + canonical: "YEAR", + aliases: &["SQL_TSI_YEAR"], + }, +]; +const DATABRICKS_ADD_PARTS: &[DatePartToken] = &[ + DatePartToken { + canonical: "MICROSECOND", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "MILLISECOND", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "SECOND", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "MINUTE", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "HOUR", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "DAY", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "DAYOFYEAR", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "WEEK", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "MONTH", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "QUARTER", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "YEAR", + aliases: EMPTY_ALIASES, + }, +]; +const DATABRICKS_DIFF_PARTS: &[DatePartToken] = &[ + DatePartToken { + canonical: "MICROSECOND", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "MILLISECOND", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "SECOND", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "MINUTE", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "HOUR", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "DAY", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "WEEK", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "MONTH", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "QUARTER", + aliases: EMPTY_ALIASES, + }, + DatePartToken { + canonical: "YEAR", + aliases: EMPTY_ALIASES, + }, +]; +const REDSHIFT_COMMON_PARTS: &[DatePartToken] = &[ + DatePartToken { + canonical: "MILLENNIUM", + aliases: &["millennia", "mil", "mils"], + }, + DatePartToken { + canonical: "CENTURY", + aliases: &["centuries", "c", "cent", "cents"], + }, + DatePartToken { + canonical: "DECADE", + aliases: &["decades", "dec", "decs"], + }, + DatePartToken { + canonical: "YEAR", + aliases: &["years", "y", "yr", "yrs"], + }, + DatePartToken { + canonical: "QUARTER", + aliases: &["quarters", "qtr", "qtrs"], + }, + DatePartToken { + canonical: "MONTH", + aliases: &["months", "mon", "mons"], + }, + DatePartToken { + canonical: "WEEK", + aliases: &["weeks", "w"], + }, + DatePartToken { + canonical: "DAY", + aliases: &["days", "d"], + }, + DatePartToken { + canonical: "HOUR", + aliases: &["hours", "h", "hr", "hrs"], + }, + DatePartToken { + canonical: "MINUTE", + aliases: &["minutes", "m", "min", "mins"], + }, + DatePartToken { + canonical: "SECOND", + aliases: &["seconds", "s", "sec", "secs"], + }, + DatePartToken { + canonical: "MILLISECOND", + aliases: &[ + "ms", + "msec", + "msecs", + "msecond", + "mseconds", + "millisec", + "millisecs", + "millisecon", + "milliseconds", + ], + }, + DatePartToken { + canonical: "MICROSECOND", + aliases: &[ + "microsec", + "microsecs", + "usecond", + "useconds", + "us", + "usec", + "usecs", + "microseconds", + ], + }, +]; +const REDSHIFT_DATE_PART_EXTRA: &[DatePartToken] = &[DatePartToken { + canonical: "DAYOFWEEK", + aliases: &["dow", "dw", "weekday"], +}]; +const MSSQL_COMMON_PARTS: &[DatePartToken] = &[ + DatePartToken { + canonical: "YEAR", + aliases: &["yy", "yyyy"], + }, + DatePartToken { + canonical: "QUARTER", + aliases: &["qq", "q"], + }, + DatePartToken { + canonical: "MONTH", + aliases: &["mm", "m"], + }, + DatePartToken { + canonical: "DAYOFYEAR", + aliases: &["dy", "y"], + }, + DatePartToken { + canonical: "DAY", + aliases: &["dd", "d"], + }, + DatePartToken { + canonical: "WEEK", + aliases: &["wk", "ww"], + }, + DatePartToken { + canonical: "HOUR", + aliases: &["hh"], + }, + DatePartToken { + canonical: "MINUTE", + aliases: &["mi", "n"], + }, + DatePartToken { + canonical: "SECOND", + aliases: &["ss", "s"], + }, + DatePartToken { + canonical: "MILLISECOND", + aliases: &["ms"], + }, + DatePartToken { + canonical: "MICROSECOND", + aliases: &["mcs"], + }, +]; +const MSSQL_ADD_DIFF_EXTRA: &[DatePartToken] = &[ + DatePartToken { + canonical: "WEEKDAY", + aliases: &["dw", "w"], + }, + DatePartToken { + canonical: "NANOSECOND", + aliases: &["ns"], + }, +]; +const MSSQL_DATEPART_EXTRA: &[DatePartToken] = &[ + DatePartToken { + canonical: "WEEKDAY", + aliases: &["dw"], + }, + DatePartToken { + canonical: "NANOSECOND", + aliases: &["ns"], + }, + DatePartToken { + canonical: "TZOFFSET", + aliases: &["tz"], + }, + DatePartToken { + canonical: "ISO_WEEK", + aliases: &["isowk", "isoww"], + }, +]; +const MSSQL_DATETRUNC_EXTRA: &[DatePartToken] = &[DatePartToken { + canonical: "ISO_WEEK", + aliases: &["isowk", "isoww"], +}]; +const MSSQL_DATE_BUCKET_PARTS: &[DatePartToken] = &[ + DatePartToken { + canonical: "YEAR", + aliases: &["yy", "yyyy"], + }, + DatePartToken { + canonical: "QUARTER", + aliases: &["qq", "q"], + }, + DatePartToken { + canonical: "MONTH", + aliases: &["mm", "m"], + }, + DatePartToken { + canonical: "DAY", + aliases: &["dd", "d"], + }, + DatePartToken { + canonical: "WEEK", + aliases: &["wk", "ww"], + }, + DatePartToken { + canonical: "HOUR", + aliases: &["hh"], + }, + DatePartToken { + canonical: "MINUTE", + aliases: &["mi", "n"], + }, + DatePartToken { + canonical: "SECOND", + aliases: &["ss", "s"], + }, + DatePartToken { + canonical: "MILLISECOND", + aliases: &["ms"], + }, +]; + +static BIGQUERY_PROFILE: DatePartProfile = DatePartProfile { + token_groups: &[BQ_PARTS], + allows_weekday_modifier: true, +}; +static SNOWFLAKE_PROFILE: DatePartProfile = DatePartProfile { + token_groups: &[SNOWFLAKE_PARTS], + allows_weekday_modifier: false, +}; +static MYSQL_PROFILE: DatePartProfile = DatePartProfile { + token_groups: &[MYSQL_PARTS], + allows_weekday_modifier: false, +}; +static DATABRICKS_ADD_PROFILE: DatePartProfile = DatePartProfile { + token_groups: &[DATABRICKS_ADD_PARTS], + allows_weekday_modifier: false, +}; +static DATABRICKS_DIFF_PROFILE: DatePartProfile = DatePartProfile { + token_groups: &[DATABRICKS_DIFF_PARTS], + allows_weekday_modifier: false, +}; +static REDSHIFT_ADD_DIFF_PROFILE: DatePartProfile = DatePartProfile { + token_groups: &[REDSHIFT_COMMON_PARTS], + allows_weekday_modifier: false, +}; +static REDSHIFT_DATE_PART_PROFILE: DatePartProfile = DatePartProfile { + token_groups: &[REDSHIFT_COMMON_PARTS, REDSHIFT_DATE_PART_EXTRA], + allows_weekday_modifier: false, +}; +static MSSQL_DATEADD_PROFILE: DatePartProfile = DatePartProfile { + token_groups: &[MSSQL_COMMON_PARTS, MSSQL_ADD_DIFF_EXTRA], + allows_weekday_modifier: false, +}; +static MSSQL_DATEDIFF_PROFILE: DatePartProfile = DatePartProfile { + token_groups: &[MSSQL_COMMON_PARTS, MSSQL_ADD_DIFF_EXTRA], + allows_weekday_modifier: false, +}; +static MSSQL_DATEPART_PROFILE: DatePartProfile = DatePartProfile { + token_groups: &[MSSQL_COMMON_PARTS, MSSQL_DATEPART_EXTRA], + allows_weekday_modifier: false, +}; +static MSSQL_DATETRUNC_PROFILE: DatePartProfile = DatePartProfile { + token_groups: &[MSSQL_COMMON_PARTS, MSSQL_DATETRUNC_EXTRA], + allows_weekday_modifier: false, +}; +static MSSQL_DATE_BUCKET_PROFILE: DatePartProfile = DatePartProfile { + token_groups: &[MSSQL_DATE_BUCKET_PARTS], + allows_weekday_modifier: false, +}; + +const BIGQUERY_DATE_PART: DatePartGrammar = DatePartGrammar { + profile: &BIGQUERY_PROFILE, +}; +const SNOWFLAKE_DATE_PART: DatePartGrammar = DatePartGrammar { + profile: &SNOWFLAKE_PROFILE, +}; +const MYSQL_DATE_PART: DatePartGrammar = DatePartGrammar { + profile: &MYSQL_PROFILE, +}; +const DATABRICKS_DATE_PART: DatePartGrammar = DatePartGrammar { + profile: &DATABRICKS_ADD_PROFILE, +}; +const DATABRICKS_DIFF_DATE_PART: DatePartGrammar = DatePartGrammar { + profile: &DATABRICKS_DIFF_PROFILE, +}; +const REDSHIFT_ADD_DIFF_DATE_PART: DatePartGrammar = DatePartGrammar { + profile: &REDSHIFT_ADD_DIFF_PROFILE, +}; +const REDSHIFT_DATE_PART: DatePartGrammar = DatePartGrammar { + profile: &REDSHIFT_DATE_PART_PROFILE, +}; +const MSSQL_DATEADD_DATE_PART: DatePartGrammar = DatePartGrammar { + profile: &MSSQL_DATEADD_PROFILE, +}; +const MSSQL_DATEDIFF_DATE_PART: DatePartGrammar = DatePartGrammar { + profile: &MSSQL_DATEDIFF_PROFILE, +}; +const MSSQL_DATEPART_DATE_PART: DatePartGrammar = DatePartGrammar { + profile: &MSSQL_DATEPART_PROFILE, +}; +const MSSQL_DATETRUNC_DATE_PART: DatePartGrammar = DatePartGrammar { + profile: &MSSQL_DATETRUNC_PROFILE, +}; +const MSSQL_DATE_BUCKET_DATE_PART: DatePartGrammar = DatePartGrammar { + profile: &MSSQL_DATE_BUCKET_PROFILE, +}; + +const BQ_TRUNC_NAMES: &[&str] = &[ + "DATE_TRUNC", + "DATETIME_TRUNC", + "TIME_TRUNC", + "TIMESTAMP_TRUNC", +]; +const BQ_TRUNC_WITH_TIMEZONE_NAMES: &[&str] = &["TIMESTAMP_TRUNC"]; +const BQ_DIFF_NAMES: &[&str] = &["DATE_DIFF", "DATETIME_DIFF", "TIME_DIFF", "TIMESTAMP_DIFF"]; +const LAST_DAY_NAMES: &[&str] = &["LAST_DAY"]; +const SNOW_ADD_DIFF_NAMES: &[&str] = &[ + "DATEADD", + "TIMEADD", + "TIMESTAMPADD", + "DATEDIFF", + "TIMEDIFF", + "TIMESTAMPDIFF", +]; +const SNOW_PART_NAMES: &[&str] = &["DATE_PART", "DATE_TRUNC"]; +const SNOW_VALUE_PART_NAMES: &[&str] = &["LAST_DAY", "TRUNC"]; +const MYSQL_ADD_DIFF_NAMES: &[&str] = &["TIMESTAMPADD", "TIMESTAMPDIFF"]; +const DATABRICKS_ADD_NAMES: &[&str] = &["DATEADD", "DATE_ADD", "TIMESTAMPADD"]; +const DATABRICKS_DIFF_NAMES: &[&str] = &["DATEDIFF", "DATE_DIFF", "TIMESTAMPDIFF"]; +const REDSHIFT_ADD_DIFF_NAMES: &[&str] = &["DATEADD", "DATEDIFF"]; +const REDSHIFT_DATE_PART_NAMES: &[&str] = &["DATE_PART", "PGDATE_PART"]; +const MSSQL_DATEADD_NAMES: &[&str] = &["DATEADD"]; +const MSSQL_DATEDIFF_NAMES: &[&str] = &["DATEDIFF", "DATEDIFF_BIG"]; +const MSSQL_DATEPART_NAMES: &[&str] = &["DATEPART", "DATENAME"]; +const MSSQL_DATETRUNC_NAMES: &[&str] = &["DATETRUNC"]; +const MSSQL_DATE_BUCKET_NAMES: &[&str] = &["DATE_BUCKET"]; + +const BIGQUERY_SIGNATURES: &[FunctionSignature] = &[ + FunctionSignature { + names: BQ_TRUNC_NAMES, + arity: 2, + arguments: BQ_VALUE_PART, + }, + FunctionSignature { + names: BQ_TRUNC_WITH_TIMEZONE_NAMES, + arity: 3, + arguments: BQ_VALUE_PART_VALUE, + }, + FunctionSignature { + names: BQ_DIFF_NAMES, + arity: 3, + arguments: BQ_VALUE_VALUE_PART, + }, + FunctionSignature { + names: LAST_DAY_NAMES, + arity: 1, + arguments: BQ_VALUE, + }, + FunctionSignature { + names: LAST_DAY_NAMES, + arity: 2, + arguments: BQ_VALUE_PART, + }, +]; + +const SNOWFLAKE_SIGNATURES: &[FunctionSignature] = &[ + FunctionSignature { + names: SNOW_ADD_DIFF_NAMES, + arity: 3, + arguments: SNOW_PART_VALUE_VALUE, + }, + FunctionSignature { + names: SNOW_PART_NAMES, + arity: 2, + arguments: SNOW_PART_VALUE, + }, + FunctionSignature { + names: SNOW_VALUE_PART_NAMES, + arity: 2, + arguments: SNOW_VALUE_PART, + }, +]; + +const MYSQL_SIGNATURES: &[FunctionSignature] = &[FunctionSignature { + names: MYSQL_ADD_DIFF_NAMES, + arity: 3, + arguments: MYSQL_PART_VALUE_VALUE, +}]; + +const DATABRICKS_SIGNATURES: &[FunctionSignature] = &[ + FunctionSignature { + names: DATABRICKS_ADD_NAMES, + arity: 3, + arguments: DB_PART_VALUE_VALUE, + }, + FunctionSignature { + names: DATABRICKS_DIFF_NAMES, + arity: 3, + arguments: DB_DIFF_PART_VALUE_VALUE, + }, +]; +const REDSHIFT_SIGNATURES: &[FunctionSignature] = &[ + FunctionSignature { + names: REDSHIFT_ADD_DIFF_NAMES, + arity: 3, + arguments: REDSHIFT_PART_VALUE_VALUE, + }, + FunctionSignature { + names: REDSHIFT_DATE_PART_NAMES, + arity: 2, + arguments: REDSHIFT_PART_VALUE, + }, +]; +const MSSQL_SIGNATURES: &[FunctionSignature] = &[ + FunctionSignature { + names: MSSQL_DATEADD_NAMES, + arity: 3, + arguments: MSSQL_PART_VALUE_VALUE, + }, + FunctionSignature { + names: MSSQL_DATEDIFF_NAMES, + arity: 3, + arguments: MSSQL_DATEDIFF_PART_VALUE_VALUE, + }, + FunctionSignature { + names: MSSQL_DATEPART_NAMES, + arity: 2, + arguments: MSSQL_PART_VALUE, + }, + FunctionSignature { + names: MSSQL_DATETRUNC_NAMES, + arity: 2, + arguments: MSSQL_DATETRUNC_PART_VALUE, + }, + FunctionSignature { + names: MSSQL_DATE_BUCKET_NAMES, + arity: 3, + arguments: MSSQL_DATE_BUCKET_PART_VALUE_VALUE, + }, + FunctionSignature { + names: MSSQL_DATE_BUCKET_NAMES, + arity: 4, + arguments: MSSQL_DATE_BUCKET_PART_VALUE_VALUE_VALUE, + }, +]; + +/// Classify a function only when its name, arity, and argument forms exactly +/// match a supported grammar. Returning `None` is intentional: unknown, +/// qualified, quoted, named, or malformed calls retain the generic behavior +/// of walking every expression argument. +pub(crate) fn classify_function( + dialect: Dialect, + function: &Function, +) -> Option<&'static FunctionSignature> { + let args = match &function.args { + FunctionArguments::List(list) => &list.args, + _ => return None, + }; + if args + .iter() + .any(|arg| !matches!(arg, FunctionArg::Unnamed(FunctionArgExpr::Expr(_)))) + { + return None; + } + + let name = simple_function_name(function)?; + signature_for(dialect, name, args.len()) +} + +pub(crate) fn expression_is_static_date_part(expr: &Expr, grammar: DatePartGrammar) -> bool { + match expr { + Expr::Identifier(ident) => { + ident.quote_style.is_none() && grammar.is_part_name(&ident.value) + } + Expr::Function(function) if grammar.allows_weekday_modifier() => { + let Some(name) = simple_function_name(function) else { + return false; + }; + if !name.eq_ignore_ascii_case("WEEK") { + return false; + } + let FunctionArguments::List(args) = &function.args else { + return false; + }; + let [FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Identifier(weekday)))] = + args.args.as_slice() + else { + return false; + }; + weekday.quote_style.is_none() && DatePartGrammar::is_weekday(&weekday.value) + } + _ => false, + } +} + +fn simple_function_name(function: &Function) -> Option<&str> { + let [ObjectNamePart::Identifier(ident)] = function.name.0.as_slice() else { + return None; + }; + ident.quote_style.is_none().then_some(ident.value.as_str()) +} + +fn signature_for(dialect: Dialect, name: &str, arity: usize) -> Option<&'static FunctionSignature> { + profile(dialect).iter().find(|signature| { + signature.arity == arity + && signature + .names + .iter() + .any(|candidate| candidate.eq_ignore_ascii_case(name)) + }) +} + +fn profile(dialect: Dialect) -> &'static [FunctionSignature] { + match dialect { + Dialect::BigQuery => BIGQUERY_SIGNATURES, + Dialect::Snowflake => SNOWFLAKE_SIGNATURES, + Dialect::MySql => MYSQL_SIGNATURES, + Dialect::Databricks => DATABRICKS_SIGNATURES, + Dialect::Redshift => REDSHIFT_SIGNATURES, + Dialect::MsSql => MSSQL_SIGNATURES, + Dialect::Generic + | Dialect::Ansi + | Dialect::PostgreSql + | Dialect::Hive + | Dialect::DuckDb + | Dialect::Trino + | Dialect::Spark + | Dialect::ClickHouse + | Dialect::SQLite => &[], + } +} + +impl DatePartGrammar { + fn allows_weekday_modifier(self) -> bool { + self.profile.allows_weekday_modifier + } + + fn is_part_name(self, value: &str) -> bool { + self.profile.token_groups.iter().any(|group| { + group.iter().any(|token| { + token.canonical.eq_ignore_ascii_case(value) + || token + .aliases + .iter() + .any(|alias| alias.eq_ignore_ascii_case(value)) + }) + }) + } + + fn is_weekday(value: &str) -> bool { + matches!( + value.to_ascii_uppercase().as_str(), + "SUNDAY" | "MONDAY" | "TUESDAY" | "WEDNESDAY" | "THURSDAY" | "FRIDAY" | "SATURDAY" + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::Dialect; + use sqlparser::ast::{FunctionArg, FunctionArgOperator, ObjectNamePart, Statement}; + use sqlparser::dialect::BigQueryDialect; + use sqlparser::parser::Parser; + + fn function(sql: &str) -> Function { + let statement = Parser::parse_sql(&BigQueryDialect, sql) + .expect("SQL should parse") + .remove(0); + let Statement::Query(query) = statement else { + panic!("expected query") + }; + let sqlparser::ast::SetExpr::Select(select) = *query.body else { + panic!("expected SELECT") + }; + let sqlparser::ast::SelectItem::UnnamedExpr(Expr::Function(function)) = + select.projection.into_iter().next().unwrap() + else { + panic!("expected function") + }; + function + } + + #[test] + fn bigquery_signatures_classify_exact_roles() { + let function = function("SELECT DATE_DIFF(a, b, ISOWEEK) FROM t"); + let signature = classify_function(Dialect::BigQuery, &function).unwrap(); + let roles = signature.arguments; + assert_eq!(signature.arity, 3); + assert_eq!(roles[0], ArgumentSemantic::ValueExpression); + assert!(matches!(roles[2], ArgumentSemantic::DatePart(_))); + } + + #[test] + fn profile_table_covers_representative_signatures() { + let cases = [ + (Dialect::BigQuery, "SELECT DATE_DIFF(a, b, DAY) FROM t"), + (Dialect::Snowflake, "SELECT DATEADD(DAY, a, b) FROM t"), + (Dialect::MySql, "SELECT TIMESTAMPDIFF(DAY, a, b) FROM t"), + (Dialect::Databricks, "SELECT DATEDIFF(DAY, a, b) FROM t"), + (Dialect::Redshift, "SELECT DATEADD(DAY, a, b) FROM t"), + (Dialect::MsSql, "SELECT DATEADD(DAY, a, b) FROM t"), + ]; + for (dialect, sql) in cases { + let function = function(sql); + let signature = classify_function(dialect, &function).unwrap(); + let roles = signature.arguments; + assert_eq!(signature.arity, roles.len()); + assert!( + roles + .iter() + .any(|role| { matches!(role, ArgumentSemantic::DatePart(_)) }) + ); + } + } + + #[test] + fn unsupported_dialect_falls_back() { + let function = function("SELECT DATE_DIFF(a, b, ISOWEEK) FROM t"); + for dialect in [ + Dialect::Generic, + Dialect::Ansi, + Dialect::PostgreSql, + Dialect::Hive, + Dialect::DuckDb, + Dialect::Trino, + Dialect::Spark, + Dialect::ClickHouse, + Dialect::SQLite, + ] { + assert!( + classify_function(dialect, &function).is_none(), + "unexpected temporal profile for {dialect:?}" + ); + } + } + + #[test] + fn qualified_quoted_named_and_wrong_arity_calls_fall_back() { + let mut qualified = function("SELECT DATE_DIFF(a, b, ISOWEEK) FROM t"); + qualified.name.0.insert( + 0, + ObjectNamePart::Identifier(sqlparser::ast::Ident::new("project")), + ); + assert!(classify_function(Dialect::BigQuery, &qualified).is_none()); + + let mut quoted = function("SELECT DATE_DIFF(a, b, ISOWEEK) FROM t"); + let ObjectNamePart::Identifier(ident) = &mut quoted.name.0[0] else { + panic!("expected identifier function name") + }; + ident.quote_style = Some('"'); + assert!(classify_function(Dialect::BigQuery, "ed).is_none()); + + let mut named = function("SELECT DATE_DIFF(a, b, ISOWEEK) FROM t"); + let FunctionArguments::List(args) = &mut named.args else { + panic!("expected argument list") + }; + let FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) = args.args[0].clone() else { + panic!("expected expression argument") + }; + args.args[0] = FunctionArg::Named { + name: sqlparser::ast::Ident::new("date"), + arg: FunctionArgExpr::Expr(expr), + operator: FunctionArgOperator::Equals, + }; + assert!(classify_function(Dialect::BigQuery, &named).is_none()); + + let mut wrong_arity = function("SELECT DATE_DIFF(a, b, ISOWEEK) FROM t"); + let FunctionArguments::List(args) = &mut wrong_arity.args else { + panic!("expected argument list") + }; + args.args.pop(); + assert!(classify_function(Dialect::BigQuery, &wrong_arity).is_none()); + } + + #[test] + fn static_date_part_requires_known_grammar() { + assert!(expression_is_static_date_part( + &Expr::Identifier(sqlparser::ast::Ident::new("ISOYEAR")), + BIGQUERY_DATE_PART + )); + assert!(!expression_is_static_date_part( + &Expr::Identifier(sqlparser::ast::Ident::new("made_up_part")), + BIGQUERY_DATE_PART + )); + } + + #[test] + fn date_part_alias_tables_are_profile_scoped() { + let aliases = [ + (SNOWFLAKE_DATE_PART, "yyyy"), + (SNOWFLAKE_DATE_PART, "yyy"), + (SNOWFLAKE_DATE_PART, "years"), + (SNOWFLAKE_DATE_PART, "mon"), + (SNOWFLAKE_DATE_PART, "months"), + (SNOWFLAKE_DATE_PART, "dayofmonth"), + (SNOWFLAKE_DATE_PART, "dow"), + (SNOWFLAKE_DATE_PART, "dayofweek_iso"), + (SNOWFLAKE_DATE_PART, "wk"), + (SNOWFLAKE_DATE_PART, "weekofyeariso"), + (SNOWFLAKE_DATE_PART, "qtr"), + (SNOWFLAKE_DATE_PART, "quarters"), + (SNOWFLAKE_DATE_PART, "hh"), + (SNOWFLAKE_DATE_PART, "hours"), + (SNOWFLAKE_DATE_PART, "min"), + (SNOWFLAKE_DATE_PART, "minutes"), + (SNOWFLAKE_DATE_PART, "us"), + (SNOWFLAKE_DATE_PART, "microseconds"), + (SNOWFLAKE_DATE_PART, "epoch_milliseconds"), + (SNOWFLAKE_DATE_PART, "epoch_second"), + (SNOWFLAKE_DATE_PART, "yearofweek"), + (SNOWFLAKE_DATE_PART, "yearofweekiso"), + (MYSQL_DATE_PART, "SQL_TSI_DAY"), + (DATABRICKS_DATE_PART, "DAYOFYEAR"), + (REDSHIFT_ADD_DIFF_DATE_PART, "m"), + (REDSHIFT_ADD_DIFF_DATE_PART, "w"), + (REDSHIFT_ADD_DIFF_DATE_PART, "mon"), + (MSSQL_DATEPART_DATE_PART, "tz"), + (MSSQL_DATEPART_DATE_PART, "isoww"), + ]; + for (grammar, alias) in aliases { + assert!( + grammar.is_part_name(alias), + "alias {alias} was not recognized" + ); + } + assert!(!SNOWFLAKE_DATE_PART.is_part_name("fortnight")); + assert!(!MYSQL_DATE_PART.is_part_name("yyyy")); + assert!(!MYSQL_DATE_PART.is_part_name("DAY_SECOND")); + assert!(!DATABRICKS_DATE_PART.is_part_name("DAYOFWEEK")); + assert!(!DATABRICKS_DATE_PART.is_part_name("WEEKOFYEAR")); + assert!(!REDSHIFT_ADD_DIFF_DATE_PART.is_part_name("mm")); + assert!(!REDSHIFT_ADD_DIFF_DATE_PART.is_part_name("wk")); + assert!(!REDSHIFT_ADD_DIFF_DATE_PART.is_part_name("DOW")); + assert!(REDSHIFT_DATE_PART.is_part_name("DOW")); + assert!(!REDSHIFT_DATE_PART.is_part_name("DAYOFYEAR")); + assert!(!MSSQL_DATE_BUCKET_DATE_PART.is_part_name("mcs")); + assert!(!MSSQL_DATETRUNC_DATE_PART.is_part_name("NANOSECOND")); + } + + #[test] + fn overloads_and_profile_specific_date_parts_are_conservative() { + let bq_timezone = function("SELECT DATE_TRUNC(value, DAY, timezone) FROM t"); + assert!(classify_function(Dialect::BigQuery, &bq_timezone).is_none()); + + let mysql_composite = function("SELECT TIMESTAMPDIFF(DAY_SECOND, a, b) FROM t"); + assert!(classify_function(Dialect::MySql, &mysql_composite).is_some()); + let mysql_signature = classify_function(Dialect::MySql, &mysql_composite).unwrap(); + assert!(!expression_is_static_date_part( + &Expr::Identifier(sqlparser::ast::Ident::new("DAY_SECOND")), + match mysql_signature.arguments[0] { + ArgumentSemantic::DatePart(grammar) => grammar, + ArgumentSemantic::ValueExpression => panic!("expected date-part role"), + } + )); + + let databricks_add = function("SELECT DATEADD(DAYOFYEAR, amount, ts) FROM t"); + let add_signature = classify_function(Dialect::Databricks, &databricks_add).unwrap(); + assert!(matches!( + add_signature.arguments[0], + ArgumentSemantic::DatePart(_) + )); + let databricks_diff = function("SELECT DATEDIFF(DAYOFYEAR, start_ts, end_ts) FROM t"); + let diff_signature = classify_function(Dialect::Databricks, &databricks_diff).unwrap(); + assert!(matches!( + diff_signature.arguments[0], + ArgumentSemantic::DatePart(_) + )); + let ArgumentSemantic::DatePart(diff_grammar) = diff_signature.arguments[0] else { + unreachable!() + }; + assert!(!expression_is_static_date_part( + &Expr::Identifier(sqlparser::ast::Ident::new("DAYOFYEAR")), + diff_grammar + )); + } + + #[test] + fn redshift_and_mssql_profiles_keep_family_specific_static_parts() { + let redshift_part = function("SELECT DATE_PART(DOW, event_ts) FROM t"); + let redshift_signature = + classify_function(Dialect::Redshift, &redshift_part).expect("Redshift DATE_PART"); + let ArgumentSemantic::DatePart(redshift_grammar) = redshift_signature.arguments[0] else { + panic!("expected Redshift date-part role") + }; + assert!(expression_is_static_date_part( + &Expr::Identifier(sqlparser::ast::Ident::new("DOW")), + redshift_grammar + )); + assert!( + classify_function( + Dialect::Redshift, + &function("SELECT DATE_TRUNC(event_ts, WEEK) FROM t") + ) + .is_none() + ); + + let mssql_trunc = function("SELECT DATETRUNC(WEEK, event_ts) FROM t"); + let mssql_signature = + classify_function(Dialect::MsSql, &mssql_trunc).expect("MsSql DATETRUNC"); + let ArgumentSemantic::DatePart(mssql_grammar) = mssql_signature.arguments[0] else { + panic!("expected MsSql date-part role") + }; + assert!(expression_is_static_date_part( + &Expr::Identifier(sqlparser::ast::Ident::new("WEEK")), + mssql_grammar + )); + for unsupported in ["WEEKDAY", "TZOFFSET", "NANOSECOND"] { + assert!(!expression_is_static_date_part( + &Expr::Identifier(sqlparser::ast::Ident::new(unsupported)), + mssql_grammar + )); + } + let mssql_dateadd = function("SELECT DATEADD(TZOFFSET, amount, event_ts) FROM t"); + let ArgumentSemantic::DatePart(mssql_dateadd_grammar) = + classify_function(Dialect::MsSql, &mssql_dateadd) + .expect("MsSql DATEADD") + .arguments[0] + else { + panic!("expected MsSql DATEADD date-part role") + }; + assert!(!expression_is_static_date_part( + &Expr::Identifier(sqlparser::ast::Ident::new("TZOFFSET")), + mssql_dateadd_grammar + )); + assert!( + classify_function( + Dialect::Spark, + &function("SELECT DATEDIFF(DAY, a, b) FROM t") + ) + .is_none() + ); + } +} diff --git a/sqllineage/src/build/mod.rs b/sqllineage/src/build/mod.rs index 605b29c..f4f243d 100644 --- a/sqllineage/src/build/mod.rs +++ b/sqllineage/src/build/mod.rs @@ -1,11 +1,12 @@ pub(crate) mod expr; +pub(crate) mod function_semantics; pub(crate) mod query; pub(crate) mod select; pub(crate) mod statement; use crate::graph::RawGraph; use crate::graph::scope::{Binding, ScopeId, ScopeKind, ScopeTree}; -use crate::types::{StatementType, Warning}; +use crate::types::{Dialect, StatementType, Warning}; use sqlparser::ast::Statement; pub(crate) struct LineageBuilder { @@ -15,10 +16,11 @@ pub(crate) struct LineageBuilder { pub(crate) warnings: Vec, pub(crate) normalize_case: bool, pub(crate) inner_statement_type: Option, + pub(crate) dialect: Dialect, } impl LineageBuilder { - pub fn new(normalize_case: bool) -> Self { + pub fn new(normalize_case: bool, dialect: Dialect) -> Self { let graph = RawGraph::new(); let root = ScopeTree::root(); Self { @@ -28,6 +30,7 @@ impl LineageBuilder { warnings: Vec::new(), normalize_case, inner_statement_type: None, + dialect, } } diff --git a/sqllineage/src/build/query.rs b/sqllineage/src/build/query.rs index b558d5c..59b7b96 100644 --- a/sqllineage/src/build/query.rs +++ b/sqllineage/src/build/query.rs @@ -1,9 +1,7 @@ use sqlparser::ast::{Query, SetExpr}; use crate::build::LineageBuilder; -use crate::graph::edge::RawEdge; -use crate::graph::node::NodeId; -use crate::graph::scope::{Binding, ScopeKind}; +use crate::graph::scope::{Binding, OutputPlan, ScopeKind}; impl LineageBuilder { pub(crate) fn visit_query(&mut self, query: &Query) { @@ -39,7 +37,13 @@ impl LineageBuilder { .scopes .output_columns(self.current_scope) .to_vec(); + let body_scope = self.current_scope; self.pop_scope(); + // The parent owns the query's public output. Keep the child plan + // intact and delegate through it after returning to the parent. + self.graph + .scopes + .set_output_plan(self.current_scope, OutputPlan::Delegate(body_scope)); for col in body_outputs { self.graph.scopes.add_output_column(self.current_scope, col); } @@ -54,7 +58,7 @@ impl LineageBuilder { SetExpr::SetOperation { left, right, .. } => { let left_scope = self.push_scope(ScopeKind::SetOperation); self.visit_set_expr(left); - let left_outputs: Vec<(String, NodeId)> = self + let left_outputs: Vec<(String, crate::graph::node::NodeId)> = self .graph .scopes .output_columns(left_scope) @@ -65,39 +69,17 @@ impl LineageBuilder { let right_scope = self.push_scope(ScopeKind::SetOperation); self.visit_set_expr(right); - let right_outputs: Vec<(String, NodeId)> = self - .graph - .scopes - .output_columns(right_scope) - .iter() - .map(|c| (c.name.clone(), c.node_id)) - .collect(); self.pop_scope(); let is_recursive = self.recursive_cte_name.is_some(); - - let pair_count = left_outputs.len().min(right_outputs.len()); - for i in 0..pair_count { - let left_out = left_outputs[i].1; - let right_out = right_outputs[i].1; - - let redirected: Vec<(NodeId, _)> = self - .graph - .edges - .iter() - .filter(|e| e.to == right_out) - .map(|e| (e.from, e.kind.clone())) - .collect(); - - for (from, kind) in redirected { - self.graph.edges.push(RawEdge { - from, - to: left_out, - kind, - is_recursive_back_edge: is_recursive, - }); - } - } + self.graph.scopes.set_output_plan( + self.current_scope, + OutputPlan::SetOperation { + left: left_scope, + right: right_scope, + recursive: is_recursive, + }, + ); for (name, node_id) in &left_outputs { self.graph.scopes.add_output_column( diff --git a/sqllineage/src/build/select.rs b/sqllineage/src/build/select.rs index 803400f..9700307 100644 --- a/sqllineage/src/build/select.rs +++ b/sqllineage/src/build/select.rs @@ -1,10 +1,12 @@ use sqlparser::ast::{ - Expr, Ident, Select, SelectItem, SelectItemQualifiedWildcardKind, TableFactor, TableWithJoins, + Expr, ObjectName, Select, SelectItem, SelectItemQualifiedWildcardKind, TableFactor, + TableWithJoins, WildcardAdditionalOptions, }; use crate::build::LineageBuilder; use crate::build::expr::determine_edge_kind; -use crate::graph::scope::{Binding, ScopeColumn, ScopeKind}; +use crate::graph::node::{StarBase, StarColumnName, StarOptions, StarReplacement}; +use crate::graph::scope::{Binding, ScopeColumn, ScopeKind, VirtualColumn, VirtualColumnState}; impl LineageBuilder { /// Process a SELECT — FROM first, then projection. @@ -24,7 +26,7 @@ impl LineageBuilder { let ancestors = self.collect_ancestors(expr); let kind = determine_edge_kind(expr); let name = infer_column_name(expr); - let output = self.graph.add_output(name.clone()); + let output = self.graph.add_output(name.clone(), kind.clone()); for &anc in &ancestors { self.graph.add_edge(anc, output, kind.clone()); } @@ -40,7 +42,7 @@ impl LineageBuilder { let ancestors = self.collect_ancestors(expr); let kind = determine_edge_kind(expr); let name = alias.value.clone(); - let output = self.graph.add_output(name.clone()); + let output = self.graph.add_output(name.clone(), kind.clone()); for &anc in &ancestors { self.graph.add_edge(anc, output, kind.clone()); } @@ -57,7 +59,7 @@ impl LineageBuilder { let kind = determine_edge_kind(expr); for alias in aliases { let name = alias.value.clone(); - let output = self.graph.add_output(name.clone()); + let output = self.graph.add_output(name.clone(), kind.clone()); for &anc in &ancestors { self.graph.add_edge(anc, output, kind.clone()); } @@ -70,8 +72,13 @@ impl LineageBuilder { ); } } - SelectItem::Wildcard(_) => { - let star = self.graph.add_star(None, self.current_scope); + SelectItem::Wildcard(options) => { + let star_options = self.star_options(options); + let star = self.graph.add_star( + StarBase::Unqualified, + star_options, + self.current_scope, + ); self.graph.scopes.add_output_column( self.current_scope, ScopeColumn { @@ -80,23 +87,107 @@ impl LineageBuilder { }, ); } - SelectItem::QualifiedWildcard(kind, _) => { - if let SelectItemQualifiedWildcardKind::ObjectName(obj_name) = kind { - let table = self.table_ref_from_object_name(obj_name); - let star = self.graph.add_star(Some(table), self.current_scope); - self.graph.scopes.add_output_column( - self.current_scope, - ScopeColumn { - name: "*".to_string(), - node_id: star, - }, - ); - } + SelectItem::QualifiedWildcard(kind, options) => { + let base = match kind { + SelectItemQualifiedWildcardKind::ObjectName(obj_name) => { + StarBase::Qualified(self.object_name_parts(obj_name)) + } + SelectItemQualifiedWildcardKind::Expr(expr) => { + StarBase::Expr(self.collect_ancestors(expr)) + } + }; + let star_options = self.star_options(options); + let star = self.graph.add_star(base, star_options, self.current_scope); + self.graph.scopes.add_output_column( + self.current_scope, + ScopeColumn { + name: "*".to_string(), + node_id: star, + }, + ); } } } } + fn object_name_parts(&self, name: &ObjectName) -> Vec { + name.0 + .iter() + .map(|part| { + part.as_ident() + .map_or_else(|| part.to_string(), |ident| self.normalize_ident(ident)) + }) + .collect() + } + + fn star_options(&mut self, options: &WildcardAdditionalOptions) -> StarOptions { + let mut result = StarOptions { + ilike: options + .opt_ilike + .as_ref() + .map(|ilike| ilike.pattern.clone()), + ..StarOptions::default() + }; + + if let Some(exclude) = &options.opt_exclude { + result.exclude.extend(match exclude { + sqlparser::ast::ExcludeSelectItem::Single(name) => { + vec![StarColumnName { + parts: self.object_name_parts(name), + }] + } + sqlparser::ast::ExcludeSelectItem::Multiple(names) => names + .iter() + .map(|name| StarColumnName { + parts: self.object_name_parts(name), + }) + .collect(), + }); + } + if let Some(except) = &options.opt_except { + result.exclude.push(StarColumnName { + parts: vec![self.normalize_ident(&except.first_element)], + }); + result.exclude.extend( + except + .additional_elements + .iter() + .map(|ident| StarColumnName { + parts: vec![self.normalize_ident(ident)], + }), + ); + } + if let Some(rename) = &options.opt_rename { + let entries = match rename { + sqlparser::ast::RenameSelectItem::Single(entry) => vec![entry], + sqlparser::ast::RenameSelectItem::Multiple(entries) => entries.iter().collect(), + }; + result.rename.extend(entries.into_iter().map(|entry| { + ( + self.normalize_ident(&entry.ident), + self.normalize_ident(&entry.alias), + ) + })); + } + if let Some(replace) = &options.opt_replace { + for element in &replace.items { + let node_id = self.graph.add_output( + "?wildcard-replace".to_string(), + determine_edge_kind(&element.expr), + ); + for ancestor in self.collect_ancestors(&element.expr) { + self.graph + .add_edge(ancestor, node_id, determine_edge_kind(&element.expr)); + } + result.replace.push(StarReplacement { + column: self.normalize_ident(&element.column_name), + node_id, + }); + } + } + result + } + /// Process FROM clause items (including JOINs). pub(crate) fn visit_from(&mut self, from: &[TableWithJoins]) { for table_with_joins in from { @@ -167,9 +258,69 @@ impl LineageBuilder { let _ = alias; } + TableFactor::UNNEST { + alias, + array_exprs, + with_offset, + with_offset_alias, + with_ordinality, + } => { + // The array expressions are evaluated in the scope visible + // before this FROM item is introduced. Capture their nodes + // first, then install the range-variable binding so it is + // visible to subsequent lateral FROM items and projection. + let dependencies = array_exprs + .iter() + .map(|expr| self.collect_ancestors(expr)) + .collect::>(); + + let Some(alias) = alias else { + return; + }; + + let mut columns = Vec::with_capacity( + array_exprs.len() + usize::from(*with_offset) + usize::from(*with_ordinality), + ); + for (index, deps) in dependencies.into_iter().enumerate() { + let name = alias.columns.get(index).map_or_else( + || alias.name.value.clone(), + |column| column.name.value.clone(), + ); + columns.push(VirtualColumn { + name, + state: if deps.is_empty() { + VirtualColumnState::KnownEmpty + } else { + VirtualColumnState::Unknown + }, + dependencies: deps, + }); + } + if *with_offset { + columns.push(VirtualColumn { + name: with_offset_alias + .as_ref() + .map_or_else(|| "offset".to_string(), |ident| ident.value.clone()), + dependencies: Vec::new(), + state: VirtualColumnState::KnownEmpty, + }); + } else if *with_ordinality { + columns.push(VirtualColumn { + name: "ordinality".to_string(), + dependencies: Vec::new(), + state: VirtualColumnState::KnownEmpty, + }); + } + + let virtual_id = self + .graph + .scopes + .add_virtual_source(self.current_scope, columns); + self.add_binding(alias.name.value.clone(), Binding::VirtualSource(virtual_id)); + } + TableFactor::TableFunction { .. } | TableFactor::Function { .. } - | TableFactor::UNNEST { .. } | TableFactor::JsonTable { .. } | TableFactor::OpenJsonTable { .. } | TableFactor::Pivot { .. } @@ -193,15 +344,3 @@ fn infer_column_name(expr: &Expr) -> String { _ => "?column?".to_string(), } } - -/// Split a compound identifier into (qualifier, `column_name`). -pub(crate) fn split_compound(parts: &[Ident]) -> (String, String) { - let len = parts.len(); - let column = parts[len - 1].value.clone(); - let qualifier = parts[..len - 1] - .iter() - .map(|p| p.value.as_str()) - .collect::>() - .join("."); - (qualifier, column) -} diff --git a/sqllineage/src/build/statement.rs b/sqllineage/src/build/statement.rs index 684f052..e42613a 100644 --- a/sqllineage/src/build/statement.rs +++ b/sqllineage/src/build/statement.rs @@ -50,7 +50,7 @@ impl LineageBuilder { let col_name = assignment_target_name(&assignment.target); let ancestors = self.collect_ancestors(&assignment.value); let kind = determine_edge_kind(&assignment.value); - let output = self.graph.add_output(col_name.clone()); + let output = self.graph.add_output(col_name.clone(), kind.clone()); for &anc in &ancestors { self.graph.add_edge(anc, output, kind.clone()); } @@ -101,7 +101,7 @@ impl LineageBuilder { let col_name = assignment_target_name(&assignment.target); let ancestors = self.collect_ancestors(&assignment.value); let kind = determine_edge_kind(&assignment.value); - let output = self.graph.add_output(col_name.clone()); + let output = self.graph.add_output(col_name.clone(), kind.clone()); for &anc in &ancestors { self.graph.add_edge(anc, output, kind.clone()); } @@ -134,7 +134,8 @@ impl LineageBuilder { .unwrap_or_else(|| format!("col{i}")); let ancestors = self.collect_ancestors(expr); let kind = determine_edge_kind(expr); - let output = self.graph.add_output(col_name.clone()); + let output = + self.graph.add_output(col_name.clone(), kind.clone()); for &anc in &ancestors { self.graph.add_edge(anc, output, kind.clone()); } @@ -280,7 +281,7 @@ impl LineageBuilder { | Statement::UNCache { .. } | Statement::UNLISTEN { .. } | Statement::Unload { .. } - | Statement::UnlockTables { .. } + | Statement::UnlockTables | Statement::Use(_) | Statement::Vacuum { .. } | Statement::WaitFor { .. } @@ -371,7 +372,7 @@ impl LineageBuilder { } } - fn normalize_ident(&self, ident: &Ident) -> String { + pub(crate) fn normalize_ident(&self, ident: &Ident) -> String { if self.normalize_case && ident.quote_style.is_none() { ident.value.to_lowercase() } else { diff --git a/sqllineage/src/dialect.rs b/sqllineage/src/dialect.rs index e6f1d1f..3a33fd6 100644 --- a/sqllineage/src/dialect.rs +++ b/sqllineage/src/dialect.rs @@ -1,7 +1,8 @@ use crate::types::Dialect; use sqlparser::dialect::{ - self, AnsiDialect, BigQueryDialect, DatabricksDialect, GenericDialect, HiveDialect, - MySqlDialect, PostgreSqlDialect, SnowflakeDialect, + self, AnsiDialect, BigQueryDialect, ClickHouseDialect, DatabricksDialect, DuckDbDialect, + GenericDialect, HiveDialect, MsSqlDialect, MySqlDialect, PostgreSqlDialect, RedshiftSqlDialect, + SQLiteDialect, SnowflakeDialect, SparkSqlDialect, }; impl Dialect { @@ -15,6 +16,16 @@ impl Dialect { Dialect::Databricks => Box::new(DatabricksDialect), Dialect::Snowflake => Box::new(SnowflakeDialect), Dialect::BigQuery => Box::new(BigQueryDialect), + Dialect::DuckDb => Box::new(DuckDbDialect), + Dialect::Redshift => Box::new(RedshiftSqlDialect {}), + // sqlparser 0.62 has no TrinoDialect; GenericDialect accepts the + // common Trino grammar without pretending to provide dialect-only + // validation. + Dialect::Trino => Box::new(GenericDialect), + Dialect::Spark => Box::new(SparkSqlDialect), + Dialect::ClickHouse => Box::new(ClickHouseDialect {}), + Dialect::SQLite => Box::new(SQLiteDialect {}), + Dialect::MsSql => Box::new(MsSqlDialect {}), } } } diff --git a/sqllineage/src/graph/mod.rs b/sqllineage/src/graph/mod.rs index 6cc33cd..d5a0375 100644 --- a/sqllineage/src/graph/mod.rs +++ b/sqllineage/src/graph/mod.rs @@ -3,9 +3,8 @@ pub(crate) mod node; pub(crate) mod scope; use crate::types::TableLineage; -use crate::types::TableRef; use edge::{EdgeKind, RawEdge}; -use node::{NodeId, RawNode}; +use node::{NodeId, RawNode, StarBase, StarOptions}; use scope::{ScopeId, ScopeTree}; pub(crate) struct RawGraph { @@ -31,24 +30,60 @@ impl RawGraph { id } - pub fn add_output(&mut self, name: String) -> NodeId { - self.add_node(RawNode::Output { name }) + pub fn add_output(&mut self, name: String, intrinsic_kind: EdgeKind) -> NodeId { + self.add_node(RawNode::Output { + name, + intrinsic_kind, + }) } - pub fn add_ref(&mut self, name: String, qualifier: Option, scope: ScopeId) -> NodeId { + pub fn add_ref_with_binding( + &mut self, + name: String, + qualifier: Option, + scope: ScopeId, + binding: Option, + ) -> NodeId { self.add_node(RawNode::Ref { name, qualifier, scope, + binding, + }) + } + + pub fn add_unqualified_with_binding( + &mut self, + name: String, + scope: ScopeId, + binding: Option, + ) -> NodeId { + self.add_node(RawNode::Unqualified { + name, + scope, + binding, }) } - pub fn add_unqualified(&mut self, name: String, scope: ScopeId) -> NodeId { - self.add_node(RawNode::Unqualified { name, scope }) + pub fn add_row_value_candidate( + &mut self, + name: String, + scope: ScopeId, + binding: Option, + ) -> NodeId { + self.add_node(RawNode::RowValueCandidate { + name, + scope, + binding, + }) } - pub fn add_star(&mut self, table: Option, scope: ScopeId) -> NodeId { - self.add_node(RawNode::Star { table, scope }) + pub fn add_star(&mut self, base: StarBase, options: StarOptions, scope: ScopeId) -> NodeId { + self.add_node(RawNode::Star { + base, + options, + scope, + }) } pub fn add_edge(&mut self, from: NodeId, to: NodeId, kind: EdgeKind) { diff --git a/sqllineage/src/graph/node.rs b/sqllineage/src/graph/node.rs index 3227495..9d7fe3f 100644 --- a/sqllineage/src/graph/node.rs +++ b/sqllineage/src/graph/node.rs @@ -1,23 +1,78 @@ -use crate::graph::scope::ScopeId; -use crate::types::TableRef; - +use crate::graph::edge::EdgeKind; +use crate::graph::scope::{Binding, ScopeId}; pub(crate) type NodeId = usize; +/// The expression which a wildcard expands from. Keeping the original parts +/// lets resolution distinguish a relation prefix from a nested field path. +#[derive(Debug, Clone)] +pub(crate) enum StarBase { + Unqualified, + Qualified(Vec), + Expr(Vec), +} + +#[derive(Debug, Clone)] +pub(crate) struct StarColumnName { + pub parts: Vec, +} + +#[derive(Debug, Clone, Default)] +pub(crate) struct StarOptions { + pub exclude: Vec, + pub ilike: Option, + pub replace: Vec, + pub rename: Vec<(String, String)>, +} + +#[derive(Debug, Clone)] +pub(crate) struct StarReplacement { + pub column: String, + pub node_id: NodeId, +} + #[derive(Debug, Clone)] pub(crate) enum RawNode { /// Output column — produced by a projection or assignment. - Output { name: String }, + Output { + name: String, + /// The edge kind the defining expression would carry to its own + /// ancestors, kept even when it has none (e.g. `COUNT(*)` has no + /// column ancestor but is still an aggregate). Used as a fallback + /// classification when no ancestor edge exists to classify from. + intrinsic_kind: EdgeKind, + }, /// Named reference — alias, CTE reference, derived table column. Ref { name: String, qualifier: Option, scope: ScopeId, + /// Binding captured while building a FROM expression. This prevents + /// a later table alias from changing the meaning of a lateral + /// dependency (for example `base, UNNEST(base.items) AS base`). + binding: Option, }, /// SELECT * or table.* — expandable with catalog. Star { - table: Option, + base: StarBase, + options: StarOptions, scope: ScopeId, }, /// Unqualified column in multi-table scope. - Unqualified { name: String, scope: ScopeId }, + Unqualified { + name: String, + scope: ScopeId, + /// Binding visible while the expression was built. This keeps a + /// lateral FROM dependency attached to the preceding relation even + /// when a later range variable shadows its name. + binding: Option, + }, + /// A relation alias used where the dialect permits a whole-row value. + /// Resolution must distinguish this from a source-free expression: a + /// catalog or derived scope can still prove that the alias is an ordinary + /// physical/output column with the same name. + RowValueCandidate { + name: String, + scope: ScopeId, + binding: Option, + }, } diff --git a/sqllineage/src/graph/scope.rs b/sqllineage/src/graph/scope.rs index de9f4ab..2073e60 100644 --- a/sqllineage/src/graph/scope.rs +++ b/sqllineage/src/graph/scope.rs @@ -5,6 +5,30 @@ use crate::types::TableRef; pub(crate) type ScopeId = usize; +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(crate) struct VirtualSourceId { + scope: ScopeId, + index: usize, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum VirtualColumnState { + KnownEmpty, + Unknown, +} + +#[derive(Debug, Clone)] +pub(crate) struct VirtualColumn { + pub name: String, + pub dependencies: Vec, + pub state: VirtualColumnState, +} + +#[derive(Debug, Clone)] +pub(crate) struct VirtualSource { + pub columns: Vec, +} + pub(crate) struct ScopeTree { scopes: Vec, } @@ -14,6 +38,8 @@ struct Scope { bindings: HashMap, anonymous_derived: Vec, output_columns: Vec, + output_plan: OutputPlan, + virtual_sources: Vec, } #[derive(Debug, Clone)] @@ -25,11 +51,25 @@ pub(crate) enum ScopeKind { SetOperation, } +/// Describes how a scope's output columns are assembled. This is retained in +/// the raw graph until catalog expansion and positional set-operation merge. +#[derive(Debug, Clone)] +pub(crate) enum OutputPlan { + Projection, + SetOperation { + left: ScopeId, + right: ScopeId, + recursive: bool, + }, + Delegate(ScopeId), +} + #[derive(Debug, Clone)] pub(crate) enum Binding { Table(TableRef), Cte(ScopeId), DerivedTable(ScopeId), + VirtualSource(VirtualSourceId), } #[derive(Debug, Clone)] @@ -46,6 +86,8 @@ impl ScopeTree { bindings: HashMap::new(), anonymous_derived: Vec::new(), output_columns: Vec::new(), + output_plan: OutputPlan::Projection, + virtual_sources: Vec::new(), }], } } @@ -61,6 +103,8 @@ impl ScopeTree { bindings: HashMap::new(), anonymous_derived: Vec::new(), output_columns: Vec::new(), + output_plan: OutputPlan::Projection, + virtual_sources: Vec::new(), }); id } @@ -73,6 +117,22 @@ impl ScopeTree { self.scopes[scope].bindings.insert(name, binding); } + pub fn add_virtual_source( + &mut self, + scope: ScopeId, + columns: Vec, + ) -> VirtualSourceId { + let index = self.scopes[scope].virtual_sources.len(); + self.scopes[scope] + .virtual_sources + .push(VirtualSource { columns }); + VirtualSourceId { scope, index } + } + + pub fn virtual_source(&self, id: VirtualSourceId) -> &VirtualSource { + &self.scopes[id.scope].virtual_sources[id.index] + } + pub fn add_output_column(&mut self, scope: ScopeId, col: ScopeColumn) { self.scopes[scope].output_columns.push(col); } @@ -91,6 +151,14 @@ impl ScopeTree { &self.scopes[scope].output_columns } + pub fn set_output_plan(&mut self, scope: ScopeId, plan: OutputPlan) { + self.scopes[scope].output_plan = plan; + } + + pub fn output_plan(&self, scope: ScopeId) -> &OutputPlan { + &self.scopes[scope].output_plan + } + pub fn add_anonymous_derived(&mut self, parent: ScopeId, child: ScopeId) { self.scopes[parent].anonymous_derived.push(child); } diff --git a/sqllineage/src/lib.rs b/sqllineage/src/lib.rs index 80b2b58..e98860b 100644 --- a/sqllineage/src/lib.rs +++ b/sqllineage/src/lib.rs @@ -59,7 +59,9 @@ use sqlparser::parser::Parser; /// /// # Errors /// -/// Returns [`ParseError`] if the SQL string cannot be parsed. +/// Returns [`ParseError`] if the SQL string cannot be parsed or fails semantic +/// validation during lineage analysis (for example, exact set-operation +/// branches with different column counts). #[allow(clippy::needless_pass_by_value)] pub fn analyze(sql: &str, opts: AnalyzeOptions) -> Result, ParseError> { let dialect = opts.dialect.to_sqlparser_dialect(); @@ -68,12 +70,12 @@ pub fn analyze(sql: &str, opts: AnalyzeOptions) -> Result, Pa })?; let catalog = opts.catalog; - Ok(statements + statements .iter() .map(|stmt| { - let builder = build::LineageBuilder::new(opts.normalize_case); + let builder = build::LineageBuilder::new(opts.normalize_case, opts.dialect); let (raw_graph, warnings, statement_type) = builder.build(stmt); resolve::resolve(raw_graph, catalog.as_deref(), warnings, statement_type) }) - .collect()) + .collect() } diff --git a/sqllineage/src/resolve/catalog.rs b/sqllineage/src/resolve/catalog.rs index 97dcc71..52363ea 100644 --- a/sqllineage/src/resolve/catalog.rs +++ b/sqllineage/src/resolve/catalog.rs @@ -5,6 +5,7 @@ pub(crate) fn apply_catalog(mappings: &mut Vec, catalog: &dyn Cat for mapping in mappings.iter_mut() { for source in &mut mapping.sources { if let ColumnOrigin::Ambiguous { column, candidates } = source + && !candidates.is_empty() && let Some(owner) = catalog.resolve_column(column, candidates) { *source = ColumnOrigin::Concrete { diff --git a/sqllineage/src/resolve/mod.rs b/sqllineage/src/resolve/mod.rs index c544d0e..8300ee4 100644 --- a/sqllineage/src/resolve/mod.rs +++ b/sqllineage/src/resolve/mod.rs @@ -1,34 +1,41 @@ mod catalog; mod topo; +#[cfg(test)] +use std::cell::Cell; use std::collections::{HashMap, HashSet}; +use std::sync::Arc; use crate::graph::RawGraph; use crate::graph::edge::EdgeKind; -use crate::graph::node::{NodeId, RawNode}; -use crate::graph::scope::{Binding, ScopeTree}; +use crate::graph::node::{NodeId, RawNode, StarBase, StarColumnName, StarOptions}; +use crate::graph::scope::{Binding, OutputPlan, ScopeTree, VirtualColumnState, VirtualSourceId}; use crate::types::{ AnalyzeResult, CatalogProvider, ColumnLineage, ColumnMapping, ColumnOrigin, ColumnRef, - StatementType, TableRef, TransformKind, Warning, WarningKind, + ParseError, StatementType, TableRef, TransformKind, Warning, WarningKind, }; /// Resolve `RawGraph` into `AnalyzeResult`. +#[allow(clippy::if_not_else, clippy::useless_let_if_seq)] pub(crate) fn resolve( mut graph: RawGraph, catalog: Option<&dyn CatalogProvider>, mut warnings: Vec, statement_type: StatementType, -) -> AnalyzeResult { +) -> Result { graph.tables.inputs.sort(); graph.tables.inputs.dedup(); if graph.nodes.is_empty() { - return AnalyzeResult { + return Ok(AnalyzeResult { statement_type, tables: graph.tables, - columns: ColumnLineage::default(), + columns: ColumnLineage { + mappings: Vec::new(), + has_unresolved_stars: false, + }, warnings, - }; + }); } if topo::topological_sort(&graph.nodes, &graph.edges).is_err() { @@ -36,14 +43,20 @@ pub(crate) fn resolve( kind: WarningKind::UnexpectedCycle, location: None, }); - return AnalyzeResult { + return Ok(AnalyzeResult { statement_type, tables: graph.tables, - columns: ColumnLineage::default(), + columns: ColumnLineage { + mappings: Vec::new(), + has_unresolved_stars: false, + }, warnings, - }; + }); } + validate_set_arities(&graph, catalog)?; + let has_unresolved_stars = graph_has_unresolved_stars(&graph, catalog); + let mut incoming: Vec> = vec![vec![]; graph.nodes.len()]; for (idx, edge) in graph.edges.iter().enumerate() { incoming[edge.to].push(idx); @@ -52,70 +65,370 @@ pub(crate) fn resolve( let mut resolved: Vec> = vec![None; graph.nodes.len()]; let root = ScopeTree::root(); - let ordered_cols = graph.scopes.output_columns(root).to_vec(); - let final_ids: HashSet = ordered_cols.iter().map(|c| c.node_id).collect(); - let output_table = graph.tables.output.clone(); - let mut mappings = Vec::new(); - - for &node_id in &final_ids { - match &graph.nodes[node_id] { - RawNode::Output { name, .. } => { - let mut visited = HashSet::new(); - let (sources, edge_kinds, has_back) = - collect_output_sources(node_id, &graph, &mut resolved, &incoming, &mut visited); - let transform = derive_transform(&edge_kinds); - - if has_back { - mappings.push(ColumnMapping { - target: ColumnRef { table: output_table.clone(), column: name.clone() }, - sources: vec![ColumnOrigin::Recursive { base_sources: sources }], - transform, - }); + let mut mapping_cache = ScopeMappingCache::default(); + // Scope mappings are cached in canonical form without an output table. + // Internal CTE/derived references need that form, while the root output + // table is only presentation metadata and is attached once here. Keeping + // it out of the cache key avoids materializing the same scope once per + // output column and cannot alter source resolution. + let mut mappings = resolve_scope_mappings( + root, + &graph, + &mut resolved, + &incoming, + catalog, + &mut mapping_cache, + ) + .iter() + .cloned() + .collect::>(); + for mapping in &mut mappings { + mapping.target.table.clone_from(&output_table); + } + + Ok(AnalyzeResult { + statement_type, + tables: graph.tables, + columns: ColumnLineage { + mappings, + has_unresolved_stars, + }, + warnings, + }) +} + +/// Inspect the complete graph rather than only the root projection. A star +/// under a JOIN/CTE/derived-table boundary is still an unresolved schema +/// dependency and must be visible to consumers of the public result. +fn graph_has_unresolved_stars(graph: &RawGraph, catalog: Option<&dyn CatalogProvider>) -> bool { + graph.nodes.iter().any(|node| { + let RawNode::Star { + base, + options, + scope, + } = node + else { + return false; + }; + star_arity(base, options, *scope, graph, catalog, &mut HashSet::new()).is_none() + }) +} + +const SET_ARITY_ERROR_PREFIX: &str = "set operation arity mismatch"; + +fn validate_set_arities( + graph: &RawGraph, + catalog: Option<&dyn CatalogProvider>, +) -> Result<(), ParseError> { + let mut active = HashSet::new(); + let _ = scope_arity(ScopeTree::root(), graph, catalog, &mut active)?; + Ok(()) +} + +/// Return an exact output width when every star in a scope can be expanded; +/// otherwise return `None` and leave the eventual merge conservative. +fn scope_arity( + scope: usize, + graph: &RawGraph, + catalog: Option<&dyn CatalogProvider>, + active: &mut HashSet, +) -> Result, ParseError> { + if !active.insert(scope) { + return Ok(None); + } + let result = match graph.scopes.output_plan(scope).clone() { + OutputPlan::Projection => { + let mut width = 0; + let mut exact = true; + for col in graph.scopes.output_columns(scope) { + if let RawNode::Star { + base, + options, + scope: star_scope, + } = &graph.nodes[col.node_id] + { + match star_arity(base, options, *star_scope, graph, catalog, active) { + Some(star_width) => width += star_width, + None => exact = false, + } } else { - mappings.push(ColumnMapping { - target: ColumnRef { table: output_table.clone(), column: name.clone() }, - sources, - transform, + width += 1; + } + } + exact.then_some(width) + } + OutputPlan::Delegate(child) => scope_arity(child, graph, catalog, active)?, + OutputPlan::SetOperation { left, right, .. } => { + let left_width = scope_arity(left, graph, catalog, active)?; + let right_width = scope_arity(right, graph, catalog, active)?; + match (left_width, right_width) { + (Some(left), Some(right)) if left != right => { + return Err(ParseError { + message: format!( + "{SET_ARITY_ERROR_PREFIX}: left has {left} columns, right has {right} columns" + ), }); } + (Some(width), Some(_)) => Some(width), + _ => None, } - RawNode::Star { table, scope } => { - expand_star( - table.as_ref(), - *scope, - &graph, - &mut resolved, - &incoming, - output_table.as_ref(), - &mut mappings, - &mut HashSet::new(), - ); + } + }; + active.remove(&scope); + Ok(result) +} + +fn star_arity( + base: &StarBase, + options: &StarOptions, + scope: usize, + graph: &RawGraph, + catalog: Option<&dyn CatalogProvider>, + active: &mut HashSet, +) -> Option { + let target = resolve_star_target(base, scope, graph); + let names = match target { + StarTarget::All => { + let mut names = Vec::new(); + for (binding_name, binding) in effective_bindings(scope, graph) { + let binding_names = binding_column_names(&binding, graph, catalog, active)?; + let qualifier = [binding_name]; + names.extend(apply_name_options( + binding_names, + options, + Some(base), + Some(&qualifier), + )); + } + for &child in graph.scopes.anonymous_derived(scope) { + let child_names = scope_output_names(child, graph, catalog, active)?; + names.extend(apply_name_options(child_names, options, Some(base), None)); } - _ => {} + Some(names) } + StarTarget::Binding(binding) => binding_column_names(&binding, graph, catalog, active) + .map(|names| apply_name_options(names, options, Some(base), None)), + StarTarget::Unknown(table) => catalog + .and_then(|catalog| catalog.list_columns(&table)) + .map(|names| apply_name_options(names, options, Some(base), None)), + StarTarget::FieldPath { .. } | StarTarget::Expr => None, + }; + names.map(|names| names.len()) +} + +enum StarTarget { + All, + Binding(Binding), + Unknown(TableRef), + FieldPath { binding: Binding, path: Vec }, + Expr, +} + +fn binding_column_names( + binding: &Binding, + graph: &RawGraph, + catalog: Option<&dyn CatalogProvider>, + active: &mut HashSet, +) -> Option> { + match binding { + Binding::Table(table) => catalog.and_then(|catalog| catalog.list_columns(table)), + Binding::Cte(child) | Binding::DerivedTable(child) => { + scope_output_names(*child, graph, catalog, active) + } + Binding::VirtualSource(source) => Some( + graph + .scopes + .virtual_source(*source) + .columns + .iter() + .map(|column| column.name.clone()) + .collect(), + ), } +} - let name_order: HashMap = ordered_cols - .iter() - .enumerate() - .filter_map(|(i, c)| match &graph.nodes[c.node_id] { - RawNode::Output { name, .. } => Some((name.clone(), i)), - RawNode::Star { .. } => Some(("*".to_string(), i)), - _ => None, - }) - .collect(); - mappings.sort_by_key(|m| name_order.get(&m.target.column).copied().unwrap_or(usize::MAX)); +fn scope_output_names( + scope: usize, + graph: &RawGraph, + catalog: Option<&dyn CatalogProvider>, + active: &mut HashSet, +) -> Option> { + if !active.insert(scope) { + return None; + } + let result = match graph.scopes.output_plan(scope) { + OutputPlan::Projection => { + let mut names = Vec::new(); + for column in graph.scopes.output_columns(scope) { + match &graph.nodes[column.node_id] { + RawNode::Star { + base, + options, + scope: star_scope, + } => { + let target = resolve_star_target(base, *star_scope, graph); + let child_names = match target { + StarTarget::All => { + let mut values = Vec::new(); + for (binding_name, binding) in + effective_bindings(*star_scope, graph) + { + let binding_names = + binding_column_names(&binding, graph, catalog, active)?; + let qualifier = [binding_name]; + values.extend(apply_name_options( + binding_names, + options, + Some(base), + Some(&qualifier), + )); + } + Some(values) + } + StarTarget::Binding(binding) => binding_column_names( + &binding, graph, catalog, active, + ) + .map(|names| apply_name_options(names, options, Some(base), None)), + StarTarget::Unknown(table) => catalog + .and_then(|catalog| catalog.list_columns(&table)) + .map(|names| apply_name_options(names, options, Some(base), None)), + StarTarget::FieldPath { .. } | StarTarget::Expr => None, + }?; + names.extend(child_names); + } + _ => names.push(column.name.clone()), + } + } + Some(names) + } + OutputPlan::Delegate(child) => scope_output_names(*child, graph, catalog, active), + OutputPlan::SetOperation { left, .. } => scope_output_names(*left, graph, catalog, active), + }; + active.remove(&scope); + result +} - if let Some(cat) = catalog { - catalog::apply_catalog(&mut mappings, cat); +fn apply_name_options( + mut names: Vec, + options: &StarOptions, + base: Option<&StarBase>, + relation_qualifier: Option<&[String]>, +) -> Vec { + names.retain(|name| { + !options.exclude.iter().any(|excluded| { + excluded_matches_name_with_context(excluded, name, base, relation_qualifier) + }) && options + .ilike + .as_deref() + .is_none_or(|pattern| ilike_matches(pattern, name)) + }); + for (old, new) in &options.rename { + if let Some(name) = names.iter_mut().find(|name| same_column_name(name, old)) { + name.clone_from(new); + } } + names +} - AnalyzeResult { - statement_type, - tables: graph.tables, - columns: ColumnLineage { mappings }, - warnings, +fn excluded_matches_name_with_context( + excluded: &StarColumnName, + name: &str, + base: Option<&StarBase>, + relation_qualifier: Option<&[String]>, +) -> bool { + let Some((excluded_column, qualifier)) = excluded.parts.split_last() else { + return false; + }; + same_column_name(name, excluded_column) + && (qualifier.is_empty() + || base.is_some_and( + |base| matches!(base, StarBase::Qualified(parts) if parts == qualifier), + ) + || relation_qualifier.is_some_and(|parts| parts == qualifier)) +} + +fn same_column_name(left: &str, right: &str) -> bool { + left == right +} + +fn ilike_matches(pattern: &str, value: &str) -> bool { + fn matches(pattern: &[char], value: &[char]) -> bool { + match pattern.split_first() { + None => value.is_empty(), + Some(('%', rest)) => { + matches(rest, value) || value.first().is_some_and(|_| matches(pattern, &value[1..])) + } + Some(('_', rest)) => value + .split_first() + .is_some_and(|(_, tail)| matches(rest, tail)), + Some((part, rest)) => value.split_first().is_some_and(|(value, tail)| { + part.eq_ignore_ascii_case(value) && matches(rest, tail) + }), + } + } + matches( + &pattern.chars().collect::>(), + &value.chars().collect::>(), + ) +} + +fn resolve_star_target(base: &StarBase, scope: usize, graph: &RawGraph) -> StarTarget { + match base { + StarBase::Unqualified => StarTarget::All, + StarBase::Expr(_) => StarTarget::Expr, + StarBase::Qualified(parts) => { + let mut best: Option<(usize, Binding)> = None; + for (name, binding) in graph.scopes.visible_bindings(scope) { + let candidates = match &binding { + Binding::Table(table) => { + let mut values = Vec::new(); + if let Some(catalog) = &table.catalog { + values.push(catalog.clone()); + } + if let Some(schema) = &table.schema { + values.push(schema.clone()); + } + values.push(table.table.clone()); + vec![values, vec![name.clone()]] + } + _ => vec![vec![name.clone()]], + }; + for candidate in candidates { + if parts.len() >= candidate.len() && parts[..candidate.len()] == candidate[..] { + let matched = candidate.len(); + if best.as_ref().is_none_or(|(length, _)| matched > *length) { + best = Some((matched, binding.clone())); + } + } + } + } + if let Some((matched, binding)) = best { + if matched == parts.len() { + StarTarget::Binding(binding) + } else { + StarTarget::FieldPath { + binding, + path: parts[matched..].to_vec(), + } + } + } else { + StarTarget::Unknown(table_ref_from_parts(parts)) + } + } + } +} + +fn table_ref_from_parts(parts: &[String]) -> TableRef { + match parts { + [table] => TableRef::new(table.clone()), + [schema, table] => TableRef::with_schema(schema.clone(), table.clone()), + [catalog, schema, table] => TableRef { + catalog: Some(catalog.clone()), + schema: Some(schema.clone()), + table: table.clone(), + }, + _ => TableRef::new(parts.join(".")), } } @@ -128,143 +441,1648 @@ fn effective_bindings(scope: usize, graph: &RawGraph) -> Vec<(String, Binding)> } } -fn wildcard_mapping(output_table: Option<&TableRef>, source_table: TableRef) -> ColumnMapping { - ColumnMapping { - target: ColumnRef { - table: output_table.cloned(), - column: "*".to_string(), - }, - sources: vec![ColumnOrigin::Wildcard { - table: source_table, - }], - transform: TransformKind::Direct, +fn wildcard_mapping(output_table: Option<&TableRef>, source_table: TableRef) -> ColumnMapping { + ColumnMapping { + target: ColumnRef { + table: output_table.cloned(), + column: "*".to_string(), + }, + sources: vec![ColumnOrigin::Wildcard { + table: source_table, + }], + transform: TransformKind::Direct, + } +} + +#[allow(clippy::too_many_arguments)] +fn expand_virtual_source( + source: VirtualSourceId, + graph: &RawGraph, + resolved: &mut Vec>, + incoming: &[Vec], + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, + mappings: &mut Vec, +) { + for column in &graph.scopes.virtual_source(source).columns { + let mut origins = Vec::new(); + let mut visited = HashSet::new(); + for &dependency in &column.dependencies { + let (dependency_origins, _, _) = collect_leaf_origins( + dependency, + graph, + resolved, + incoming, + &mut visited, + catalog, + mapping_cache, + ); + origins.extend(dependency_origins); + } + if origins.is_empty() + && matches!(column.state, VirtualColumnState::Unknown) + && !column + .dependencies + .iter() + .all(|&dependency| known_empty_dependency(dependency, graph)) + { + origins.push(ColumnOrigin::Ambiguous { + column: column.name.clone(), + candidates: Vec::new(), + }); + } + mappings.push(ColumnMapping { + target: ColumnRef { + table: None, + column: column.name.clone(), + }, + sources: origins, + transform: TransformKind::Direct, + }); + } +} + +fn known_empty_dependency(node_id: NodeId, graph: &RawGraph) -> bool { + let (name, binding) = match &graph.nodes[node_id] { + RawNode::Ref { name, binding, .. } | RawNode::Unqualified { name, binding, .. } => { + (name, binding.as_ref()) + } + _ => return false, + }; + let Some(Binding::VirtualSource(source)) = binding else { + return false; + }; + let Some(column) = graph + .scopes + .virtual_source(*source) + .columns + .iter() + .find(|column| column.name == *name) + else { + return false; + }; + matches!(column.state, VirtualColumnState::KnownEmpty) + && column + .dependencies + .iter() + .all(|&dependency| known_empty_dependency(dependency, graph)) +} + +#[derive(Default)] +struct ScopeMappingCache { + entries: HashMap, +} + +enum ScopeMappingEntry { + Computing, + Resolved(Arc<[ColumnMapping]>), +} + +#[cfg(test)] +thread_local! { + static SCOPE_MAPPING_COMPUTATIONS: Cell = const { Cell::new(0) }; +} + +#[cfg(test)] +fn reset_scope_mapping_stats() { + SCOPE_MAPPING_COMPUTATIONS.with(|computations| computations.set(0)); +} + +#[cfg(test)] +fn scope_mapping_computations() -> usize { + SCOPE_MAPPING_COMPUTATIONS.with(Cell::get) +} + +/// Resolve a scope's output plan into ordered mappings. Set-operation +/// branches are resolved independently so catalog expansion happens before +/// their positional merge. +#[allow(clippy::too_many_arguments)] +fn resolve_scope_mappings( + scope: usize, + graph: &RawGraph, + resolved: &mut Vec>, + incoming: &[Vec], + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, +) -> Arc<[ColumnMapping]> { + if let Some(entry) = mapping_cache.entries.get(&scope) { + return match entry { + ScopeMappingEntry::Resolved(mappings) => mappings.clone(), + // A recursive scope cannot safely publish a partially materialized + // result. Returning no mappings preserves the existing fallback to + // raw output resolution and, importantly, does not cache an + // incomplete entry as resolved. + ScopeMappingEntry::Computing => Arc::from([]), + }; + } + mapping_cache + .entries + .insert(scope, ScopeMappingEntry::Computing); + #[cfg(test)] + SCOPE_MAPPING_COMPUTATIONS.with(|computations| computations.set(computations.get() + 1)); + + let mappings = match graph.scopes.output_plan(scope).clone() { + OutputPlan::Projection => { + let mut mappings = Vec::new(); + for col in graph.scopes.output_columns(scope) { + match &graph.nodes[col.node_id] { + RawNode::Output { name, .. } => { + let mut visited = HashSet::new(); + let (sources, edge_kinds, has_back, inherited_transform) = + collect_output_sources( + col.node_id, + graph, + resolved, + incoming, + &mut visited, + catalog, + mapping_cache, + ); + let transform = merge_transform( + &derive_transform(&graph.nodes[col.node_id], &edge_kinds), + &inherited_transform, + ); + mappings.push(ColumnMapping { + target: ColumnRef { + table: None, + column: name.clone(), + }, + sources: if has_back { + vec![ColumnOrigin::Recursive { + base_sources: sources, + }] + } else { + sources + }, + transform, + }); + } + RawNode::Star { + base, + options, + scope, + } => expand_star( + base, + options, + *scope, + graph, + resolved, + incoming, + catalog, + mapping_cache, + &mut mappings, + &mut HashSet::new(), + ), + _ => {} + } + } + if let Some(cat) = catalog { + catalog::apply_catalog(&mut mappings, cat); + } + mappings + } + OutputPlan::Delegate(child) => { + resolve_scope_mappings(child, graph, resolved, incoming, catalog, mapping_cache) + .iter() + .cloned() + .collect() + } + OutputPlan::SetOperation { + left, + right, + recursive, + } => { + let left_mappings = + resolve_scope_mappings(left, graph, resolved, incoming, catalog, mapping_cache) + .iter() + .cloned() + .collect::>(); + let right_mappings = if recursive { + Vec::new() + } else { + resolve_scope_mappings(right, graph, resolved, incoming, catalog, mapping_cache) + .iter() + .cloned() + .collect() + }; + // A wildcard without catalog metadata is a variable-width slot. + // Positional alignment at or after it would fabricate an ordinal + // and lose the remaining branch columns. Merge only the exact + // prefix before the first wildcard, then retain both tails and + // expose the Wildcard origin. + if !recursive + && (mappings_have_unknown_shape(&left_mappings) + || mappings_have_unknown_shape(&right_mappings)) + { + merge_unknown_shape_mappings(left_mappings, right_mappings) + } else { + // The branch resolvers expand stars independently. Preserve the + // left branch's names and order, as SQL set operations do. + let mut merged = Vec::with_capacity(left_mappings.len()); + for (idx, left_mapping) in left_mappings.into_iter().enumerate() { + if let Some(right_mapping) = right_mappings.get(idx) { + let (sources, transform) = + merge_branch_sources(&left_mapping, right_mapping); + if recursive { + merged.push(ColumnMapping { + target: left_mapping.target, + sources: vec![ColumnOrigin::Recursive { + base_sources: sources, + }], + transform, + }); + } else { + merged.push(ColumnMapping { + target: left_mapping.target, + sources, + transform, + }); + } + } else { + let sources = left_mapping.sources; + let transform = left_mapping.transform; + if recursive { + merged.push(ColumnMapping { + target: left_mapping.target, + sources: vec![ColumnOrigin::Recursive { + base_sources: sources, + }], + transform, + }); + } else { + merged.push(ColumnMapping { + target: left_mapping.target, + sources, + transform, + }); + } + } + } + merged + } + } + }; + let mappings: Arc<[ColumnMapping]> = mappings.into(); + mapping_cache + .entries + .insert(scope, ScopeMappingEntry::Resolved(mappings.clone())); + mappings +} + +fn mappings_have_unknown_shape(mappings: &[ColumnMapping]) -> bool { + mappings.iter().any(|mapping| { + // A named output that happens to flow through an unresolved star has + // a known ordinal/name and must not act as a variable-width barrier. + mapping.target.column == "*" + && mapping.sources.iter().any(|source| match source { + ColumnOrigin::Wildcard { .. } => true, + ColumnOrigin::Ambiguous { column, .. } if column == "*" => true, + ColumnOrigin::Recursive { base_sources } => base_sources.iter().any(|source| { + matches!(source, ColumnOrigin::Wildcard { .. }) + || matches!(source, ColumnOrigin::Ambiguous { column, .. } if column == "*") + }), + _ => false, + }) + }) +} + +fn merge_unknown_shape_mappings( + left: Vec, + right: Vec, +) -> Vec { + let left_barrier = first_unknown_mapping(&left).unwrap_or(left.len()); + let right_barrier = first_unknown_mapping(&right).unwrap_or(right.len()); + + // A leading unknown star determines the output names for the set + // operation. When the left side is already a merged set (and therefore + // contains named slots contributed by an earlier operand), a right-only + // tail would publish names that the leading operand never declared. Keep + // the left candidates and their existing wildcard provenance; a direct + // two-branch `SELECT * UNION SELECT a, b` still retains the right branch + // names because there are no prior named slots to preserve. + if left_barrier == 0 && left.len() > 1 && right_barrier == right.len() { + return left; + } + + let prefix_len = left_barrier.min(right_barrier); + let left_unknown = wildcard_sources(&left); + let right_unknown = wildcard_sources(&right); + let mut merged = Vec::with_capacity(left.len() + right.len() - prefix_len); + + for (left_mapping, right_mapping) in left.iter().zip(right.iter()).take(prefix_len) { + let (sources, transform) = merge_branch_sources(left_mapping, right_mapping); + merged.push(ColumnMapping { + target: left_mapping.target.clone(), + sources, + transform, + }); + } + merged.extend( + left.into_iter() + .skip(prefix_len) + .map(|mapping| append_unknown_sources(mapping, &right_unknown)), + ); + merged.extend( + right + .into_iter() + .skip(prefix_len) + .map(|mapping| append_unknown_sources(mapping, &left_unknown)), + ); + merged +} + +/// Merge two positional branch mappings while retaining the fact that one +/// branch was source-free. An empty source list is meaningful for literals; +/// dropping it would overclaim complete lineage from the other branch. +fn merge_branch_sources( + left: &ColumnMapping, + right: &ColumnMapping, +) -> (Vec, TransformKind) { + let mut sources = left.sources.clone(); + sources.extend(right.sources.clone()); + if left.sources.is_empty() != right.sources.is_empty() { + sources.push(ColumnOrigin::SourceFree { + column: left.target.column.clone(), + }); + } + (sources, merge_transform(&left.transform, &right.transform)) +} + +fn first_unknown_mapping(mappings: &[ColumnMapping]) -> Option { + mappings + .iter() + .position(|mapping| mappings_have_unknown_shape(std::slice::from_ref(mapping))) +} + +fn wildcard_sources(mappings: &[ColumnMapping]) -> Vec { + let mut sources = Vec::new(); + for mapping in mappings { + for source in &mapping.sources { + match source { + ColumnOrigin::Wildcard { .. } => sources.push(source.clone()), + ColumnOrigin::Ambiguous { column, .. } if column == "*" => { + sources.push(source.clone()); + } + ColumnOrigin::Recursive { base_sources } => sources.extend( + base_sources + .iter() + .filter(|source| { + matches!(source, ColumnOrigin::Wildcard { .. }) + || matches!(source, ColumnOrigin::Ambiguous { column, .. } if column == "*") + }) + .cloned(), + ), + _ => {} + } + } + } + sources +} + +fn append_unknown_sources( + mut mapping: ColumnMapping, + unknown_sources: &[ColumnOrigin], +) -> ColumnMapping { + if mapping.sources.is_empty() && !unknown_sources.is_empty() { + mapping.sources.push(ColumnOrigin::SourceFree { + column: mapping.target.column.clone(), + }); + } + mapping.sources.extend(unknown_sources.iter().cloned()); + mapping +} + +fn merge_transform(left: &TransformKind, right: &TransformKind) -> TransformKind { + if matches!(left, TransformKind::Aggregation) || matches!(right, TransformKind::Aggregation) { + TransformKind::Aggregation + } else if matches!(left, TransformKind::Conditional) + || matches!(right, TransformKind::Conditional) + { + TransformKind::Conditional + } else if matches!(left, TransformKind::Expression) + || matches!(right, TransformKind::Expression) + { + TransformKind::Expression + } else if matches!(left, TransformKind::Window) || matches!(right, TransformKind::Window) { + TransformKind::Window + } else if matches!(left, TransformKind::Unknown) || matches!(right, TransformKind::Unknown) { + TransformKind::Unknown + } else { + TransformKind::Direct + } +} + +/// Expand a Star node (qualified or unqualified) into `ColumnMapping`s. +#[allow(clippy::too_many_arguments)] +fn expand_star( + base: &StarBase, + options: &StarOptions, + scope: usize, + graph: &RawGraph, + resolved: &mut Vec>, + incoming: &[Vec], + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, + mappings: &mut Vec, + visited_scopes: &mut HashSet, +) { + let start = mappings.len(); + match resolve_star_target(base, scope, graph) { + StarTarget::All => { + for (_, binding) in effective_bindings(scope, graph) { + expand_star_binding( + &binding, + graph, + resolved, + incoming, + catalog, + mapping_cache, + mappings, + visited_scopes, + ); + } + for &child in graph.scopes.anonymous_derived(scope) { + expand_scope_columns( + child, + graph, + resolved, + incoming, + catalog, + mapping_cache, + mappings, + visited_scopes, + ); + } + } + StarTarget::Binding(binding) => expand_star_binding( + &binding, + graph, + resolved, + incoming, + catalog, + mapping_cache, + mappings, + visited_scopes, + ), + StarTarget::Unknown(table) => expand_unknown_relation(table, catalog, mappings), + StarTarget::FieldPath { binding, path } => { + let mut sources = Vec::new(); + if let Some(field) = path.first() + && let Some(origin) = resolve_captured_binding( + field, + binding, + graph, + resolved, + incoming, + &mut HashSet::new(), + catalog, + mapping_cache, + ) + { + sources.push(origin); + } + sources.push(ColumnOrigin::Ambiguous { + column: "*".to_string(), + candidates: Vec::new(), + }); + mappings.push(ColumnMapping { + target: ColumnRef { + table: None, + column: "*".to_string(), + }, + sources, + transform: TransformKind::Direct, + }); + } + StarTarget::Expr => { + let mut sources = Vec::new(); + if let StarBase::Expr(dependencies) = base { + for &dependency in dependencies { + let (origins, _, _) = collect_leaf_origins( + dependency, + graph, + resolved, + incoming, + &mut HashSet::new(), + catalog, + mapping_cache, + ); + sources.extend(origins); + } + } + if !sources.iter().any( + |source| matches!(source, ColumnOrigin::Ambiguous { column, .. } if column == "*"), + ) { + sources.push(ColumnOrigin::Ambiguous { + column: "*".to_string(), + candidates: Vec::new(), + }); + } + mappings.push(ColumnMapping { + target: ColumnRef { + table: None, + column: "*".to_string(), + }, + sources, + transform: TransformKind::Direct, + }); + } + } + apply_star_options( + mappings, + start, + base, + scope, + options, + graph, + resolved, + incoming, + catalog, + mapping_cache, + ); +} + +#[allow(clippy::too_many_arguments)] +fn expand_star_binding( + binding: &Binding, + graph: &RawGraph, + resolved: &mut Vec>, + incoming: &[Vec], + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, + mappings: &mut Vec, + visited_scopes: &mut HashSet, +) { + match binding { + Binding::Table(table) => { + if let Some(columns) = catalog.and_then(|catalog| catalog.list_columns(table)) { + mappings.extend(columns.into_iter().map(|column| ColumnMapping { + target: ColumnRef { + table: None, + column: column.clone(), + }, + sources: vec![ColumnOrigin::Concrete { + table: table.clone(), + column, + }], + transform: TransformKind::Direct, + })); + } else { + mappings.push(wildcard_mapping(None, table.clone())); + } + } + Binding::Cte(scope) | Binding::DerivedTable(scope) => expand_scope_columns( + *scope, + graph, + resolved, + incoming, + catalog, + mapping_cache, + mappings, + visited_scopes, + ), + Binding::VirtualSource(source) => expand_virtual_source( + *source, + graph, + resolved, + incoming, + catalog, + mapping_cache, + mappings, + ), + } +} + +fn expand_unknown_relation( + table: TableRef, + catalog: Option<&dyn CatalogProvider>, + mappings: &mut Vec, +) { + if let Some(columns) = catalog.and_then(|catalog| catalog.list_columns(&table)) { + mappings.extend(columns.into_iter().map(|column| ColumnMapping { + target: ColumnRef { + table: None, + column: column.clone(), + }, + sources: vec![ColumnOrigin::Concrete { + table: table.clone(), + column, + }], + transform: TransformKind::Direct, + })); + } else { + mappings.push(wildcard_mapping(None, table)); + } +} + +#[allow(clippy::too_many_arguments)] +fn apply_star_options( + mappings: &mut Vec, + start: usize, + base: &StarBase, + scope: usize, + options: &StarOptions, + graph: &RawGraph, + resolved: &mut Vec>, + incoming: &[Vec], + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, +) { + if mappings_have_unknown_shape(&mappings[start..]) { + append_unknown_star_options( + mappings, + start, + base, + scope, + options, + graph, + resolved, + incoming, + catalog, + mapping_cache, + ); + return; + } + let original = mappings.split_off(start); + let mut retained = Vec::with_capacity(original.len()); + for mapping in original { + let name = &mapping.target.column; + let excluded = options + .exclude + .iter() + .any(|excluded| excluded_matches_mapping(excluded, name, &mapping, base)); + let ilike_mismatch = options + .ilike + .as_deref() + .is_some_and(|pattern| !ilike_matches(pattern, name)); + if !excluded && !ilike_mismatch { + retained.push(mapping); + } + } + for replacement in &options.replace { + let Some(index) = retained + .iter() + .position(|mapping| same_column_name(&mapping.target.column, &replacement.column)) + else { + continue; + }; + let mut visited = HashSet::new(); + let (sources, edge_kinds, _, inherited_transform) = collect_output_sources( + replacement.node_id, + graph, + resolved, + incoming, + &mut visited, + catalog, + mapping_cache, + ); + retained[index].sources = sources; + retained[index].transform = merge_transform( + &derive_transform(&graph.nodes[replacement.node_id], &edge_kinds), + &inherited_transform, + ); + } + for (old, new) in &options.rename { + if let Some(mapping) = retained + .iter_mut() + .find(|mapping| same_column_name(&mapping.target.column, old)) + { + mapping.target.column.clone_from(new); + } + } + mappings.extend(retained); +} + +#[allow(clippy::too_many_arguments)] +fn append_unknown_star_options( + mappings: &mut Vec, + start: usize, + base: &StarBase, + scope: usize, + options: &StarOptions, + graph: &RawGraph, + resolved: &mut Vec>, + incoming: &[Vec], + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, +) { + let unknown_sources = mappings[start..].to_vec(); + for replacement in &options.replace { + if !name_passes_star_filters(&replacement.column, options, base) { + continue; + } + let mut visited = HashSet::new(); + let (sources, edge_kinds, _, inherited_transform) = collect_output_sources( + replacement.node_id, + graph, + resolved, + incoming, + &mut visited, + catalog, + mapping_cache, + ); + mappings.push(ColumnMapping { + target: ColumnRef { + table: None, + column: replacement.column.clone(), + }, + sources, + transform: merge_transform( + &derive_transform(&graph.nodes[replacement.node_id], &edge_kinds), + &inherited_transform, + ), + }); + } + for (old, new) in &options.rename { + if !name_passes_star_filters(old, options, base) { + continue; + } + let sources = + named_wildcard_sources_for_star(&unknown_sources, old, options, base, scope, graph); + if !sources + .iter() + .any(|source| matches!(source, ColumnOrigin::NamedWildcard { .. })) + { + continue; + } + mappings.push(ColumnMapping { + target: ColumnRef { + table: None, + column: new.clone(), + }, + sources, + transform: TransformKind::Direct, + }); + } +} + +fn name_passes_star_filters(name: &str, options: &StarOptions, base: &StarBase) -> bool { + !options + .exclude + .iter() + .any(|excluded| excluded_matches_name(excluded, name, base)) + && options + .ilike + .as_deref() + .is_none_or(|pattern| ilike_matches(pattern, name)) +} + +fn excluded_matches_name(excluded: &StarColumnName, name: &str, base: &StarBase) -> bool { + let Some((excluded_column, qualifier)) = excluded.parts.split_last() else { + return false; + }; + same_column_name(name, excluded_column) + && (qualifier.is_empty() + || matches!(base, StarBase::Qualified(parts) if parts == qualifier)) +} + +fn excluded_matches_mapping( + excluded: &StarColumnName, + column: &str, + mapping: &ColumnMapping, + base: &StarBase, +) -> bool { + let Some((excluded_column, qualifier)) = excluded.parts.split_last() else { + return false; + }; + if !same_column_name(column, excluded_column) { + return false; + } + if qualifier.is_empty() { + return true; + } + if let StarBase::Qualified(parts) = base + && parts == qualifier + { + return true; + } + mapping.sources.iter().any(|source| { + let table = match source { + ColumnOrigin::Concrete { table, .. } + | ColumnOrigin::Wildcard { table } + | ColumnOrigin::NamedWildcard { table, .. } => table, + ColumnOrigin::Recursive { base_sources } => { + return base_sources + .iter() + .any(|source| relation_matches_qualifier(source, qualifier)); + } + _ => return false, + }; + table_matches_qualifier(table, qualifier) + }) +} + +fn relation_matches_qualifier(source: &ColumnOrigin, qualifier: &[String]) -> bool { + match source { + ColumnOrigin::Concrete { table, .. } + | ColumnOrigin::Wildcard { table } + | ColumnOrigin::NamedWildcard { table, .. } => table_matches_qualifier(table, qualifier), + _ => false, + } +} + +fn table_matches_qualifier(table: &TableRef, qualifier: &[String]) -> bool { + match qualifier { + [name] => table.table == *name, + [schema, name] => table.schema.as_deref() == Some(schema) && table.table == *name, + [catalog, schema, name] => { + table.catalog.as_deref() == Some(catalog) + && table.schema.as_deref() == Some(schema) + && table.table == *name + } + _ => false, + } +} + +/// Recursively expand a scope's output columns into `ColumnMapping`s. +#[allow(clippy::too_many_arguments)] +fn expand_scope_columns( + scope_id: usize, + graph: &RawGraph, + resolved: &mut Vec>, + incoming: &[Vec], + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, + mappings: &mut Vec, + visited_scopes: &mut HashSet, +) { + if !visited_scopes.insert(scope_id) { + return; + } + if !matches!(graph.scopes.output_plan(scope_id), OutputPlan::Projection) { + let nested = + resolve_scope_mappings(scope_id, graph, resolved, incoming, catalog, mapping_cache); + mappings.extend(nested.iter().cloned()); + return; + } + for col in graph.scopes.output_columns(scope_id) { + if let RawNode::Star { + base, + options, + scope, + } = &graph.nodes[col.node_id] + { + expand_star( + base, + options, + *scope, + graph, + resolved, + incoming, + catalog, + mapping_cache, + mappings, + visited_scopes, + ); + } else { + let mut visited = HashSet::new(); + let (sources, edge_kinds, _, inherited_transform) = collect_output_sources( + col.node_id, + graph, + resolved, + incoming, + &mut visited, + catalog, + mapping_cache, + ); + let transform = merge_transform( + &derive_transform(&graph.nodes[col.node_id], &edge_kinds), + &inherited_transform, + ); + mappings.push(ColumnMapping { + target: ColumnRef { + table: None, + column: col.name.clone(), + }, + sources, + transform, + }); + } + } +} + +/// Collect source origins for one logical output slot, retaining both sides +/// of a set operation. Unlike the public mapping path this returns origins so +/// a later CTE/derived-table reference can continue through the slot. +fn scope_column_sources( + scope: usize, + index: usize, + graph: &RawGraph, + resolved: &mut Vec>, + incoming: &[Vec], + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, +) -> (Vec, Vec, bool, TransformKind) { + let mappings = resolve_scope_mappings(scope, graph, resolved, incoming, catalog, mapping_cache); + let Some(mapping) = mappings.get(index) else { + return (vec![], vec![], false, TransformKind::Direct); + }; + let mut sources = Vec::new(); + let mut has_back = false; + for source in &mapping.sources { + match source { + ColumnOrigin::Recursive { base_sources } => { + sources.extend(base_sources.clone()); + has_back = true; + } + source => sources.push(source.clone()), + } + } + (sources, vec![], has_back, mapping.transform.clone()) +} + +fn collect_output_sources( + node_id: NodeId, + graph: &RawGraph, + resolved: &mut Vec>, + incoming: &[Vec], + visited: &mut HashSet, + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, +) -> (Vec, Vec, bool, TransformKind) { + if !visited.insert(node_id) { + return (vec![], vec![], false, TransformKind::Direct); + } + + let mut sources = Vec::new(); + let mut kinds = Vec::new(); + let mut has_back = false; + let mut inherited_transform = TransformKind::Direct; + + for &edge_idx in &incoming[node_id] { + let edge = &graph.edges[edge_idx]; + if edge.is_recursive_back_edge { + has_back = true; + continue; + } + let (sub_sources, sub_back, sub_transform) = collect_leaf_origins( + edge.from, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ); + for _ in &sub_sources { + kinds.push(edge.kind.clone()); + } + sources.extend(sub_sources); + has_back |= sub_back; + inherited_transform = merge_transform(&inherited_transform, &sub_transform); + } + + (sources, kinds, has_back, inherited_transform) +} + +fn collect_leaf_origins( + node_id: NodeId, + graph: &RawGraph, + resolved: &mut Vec>, + incoming: &[Vec], + visited: &mut HashSet, + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, +) -> (Vec, bool, TransformKind) { + if let Some(result) = + resolve_virtual_reference(node_id, graph, resolved, incoming, catalog, mapping_cache) + { + return result; + } + + if let Some((sources, has_back, transform)) = + resolve_named_scope_reference(node_id, graph, resolved, incoming, catalog, mapping_cache) + { + return (sources, has_back, transform); + } + + if let Some((target_output, scope)) = find_cte_redirect(node_id, graph) { + if let Some(index) = graph + .scopes + .output_columns(scope) + .iter() + .position(|c| c.node_id == target_output) + && !matches!(graph.scopes.output_plan(scope), OutputPlan::Projection) + { + let (sources, _, has_back, transform) = scope_column_sources( + scope, + index, + graph, + resolved, + incoming, + catalog, + mapping_cache, + ); + return (sources, has_back, transform); + } + let (sources, _, has_back, transform) = collect_output_sources( + target_output, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ); + return (sources, has_back, transform); + } + + if let RawNode::Output { .. } = &graph.nodes[node_id] { + let (sources, _, has_back, transform) = collect_output_sources( + node_id, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ); + (sources, has_back, transform) + } else { + let origin = resolve_node( + node_id, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ); + match origin { + Some(o) => (vec![o], false, TransformKind::Direct), + None => (vec![], false, TransformKind::Direct), + } + } +} + +#[allow(clippy::too_many_arguments)] +fn resolve_virtual_reference( + node_id: NodeId, + graph: &RawGraph, + resolved: &mut Vec>, + incoming: &[Vec], + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, +) -> Option<(Vec, bool, TransformKind)> { + let (name, binding, scope) = match &graph.nodes[node_id] { + RawNode::Ref { + name, + qualifier, + scope, + binding, + } => ( + name, + binding + .clone() + .or_else(|| { + qualifier + .as_deref() + .and_then(|qualifier| graph.scopes.lookup(*scope, qualifier).cloned()) + }) + .or_else(|| { + graph + .scopes + .lookup(*scope, name) + .filter(|binding| matches!(binding, Binding::VirtualSource(_))) + .cloned() + }), + *scope, + ), + RawNode::Unqualified { + name, + scope, + binding, + } => ( + name, + binding.clone().or_else(|| { + graph + .scopes + .lookup(*scope, name) + .filter(|binding| matches!(binding, Binding::VirtualSource(_))) + .cloned() + }), + *scope, + ), + _ => return None, + }; + let source = match binding { + Some(Binding::VirtualSource(source)) => source, + Some(_) => return None, + None => match find_virtual_sources_for_column(scope, name, graph).as_slice() { + [source] => *source, + [] => return None, + _ => { + return Some(( + vec![ColumnOrigin::Ambiguous { + column: name.clone(), + candidates: Vec::new(), + }], + false, + TransformKind::Direct, + )); + } + }, + }; + resolve_virtual_column_sources( + name, + source, + graph, + resolved, + incoming, + catalog, + mapping_cache, + ) +} + +#[allow(clippy::too_many_arguments)] +fn resolve_virtual_column_sources( + name: &str, + source: VirtualSourceId, + graph: &RawGraph, + resolved: &mut Vec>, + incoming: &[Vec], + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, +) -> Option<(Vec, bool, TransformKind)> { + let column = graph + .scopes + .virtual_source(source) + .columns + .iter() + .find(|column| column.name == name)?; + let mut origins = Vec::new(); + for &dependency in &column.dependencies { + let (dependency_origins, _, _) = collect_leaf_origins( + dependency, + graph, + resolved, + incoming, + &mut HashSet::new(), + catalog, + mapping_cache, + ); + origins.extend(dependency_origins); + } + if origins.is_empty() + && matches!(column.state, VirtualColumnState::Unknown) + && !column + .dependencies + .iter() + .all(|&dependency| known_empty_dependency(dependency, graph)) + { + origins.push(ColumnOrigin::Ambiguous { + column: column.name.clone(), + candidates: Vec::new(), + }); + } + Some((origins, false, TransformKind::Direct)) +} + +fn virtual_column_origin( + name: &str, + source: VirtualSourceId, + graph: &RawGraph, +) -> Option { + let column = graph + .scopes + .virtual_source(source) + .columns + .iter() + .find(|column| column.name == name)?; + match column.state { + VirtualColumnState::KnownEmpty => None, + VirtualColumnState::Unknown => Some(ColumnOrigin::Ambiguous { + column: column.name.clone(), + candidates: Vec::new(), + }), + } +} + +/// Resolve a named reference through a set-operation/Delegate scope using the +/// expanded mappings. Raw scope columns intentionally do not contain names +/// for individual catalog-expanded star outputs, so lookup must happen here. +fn resolve_named_scope_reference( + node_id: NodeId, + graph: &RawGraph, + resolved: &mut Vec>, + incoming: &[Vec], + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, +) -> Option<(Vec, bool, TransformKind)> { + let (name, qualifier, scope) = match &graph.nodes[node_id] { + RawNode::Ref { + name, + qualifier, + scope, + .. + } => (name, qualifier.as_ref(), *scope), + RawNode::Unqualified { name, scope, .. } => (name, None, *scope), + _ => return None, + }; + let binding = match &graph.nodes[node_id] { + RawNode::Ref { + binding: Some(binding), + .. + } + | RawNode::Unqualified { + binding: Some(binding), + .. + } => Some(binding.clone()), + _ => None, + } + .or_else(|| qualifier.and_then(|qualifier| graph.scopes.lookup(scope, qualifier).cloned())) + .or_else(|| find_single_binding(scope, graph)); + let Some(Binding::Cte(target_scope) | Binding::DerivedTable(target_scope)) = binding else { + return None; + }; + let mappings = resolve_scope_mappings( + target_scope, + graph, + resolved, + incoming, + catalog, + mapping_cache, + ); + if let Some(mapping) = mappings + .iter() + .find(|mapping| mapping.target.column == *name) + { + let (sources, has_back) = flatten_mapping_sources(&mapping.sources); + return Some((sources, has_back, mapping.transform.clone())); } + let requested_name = + match scope_star_name_decision(target_scope, name, graph, &mut HashSet::new()) { + StarNameDecision::Denied | StarNameDecision::Ambiguous => { + return Some(( + vec![ColumnOrigin::Ambiguous { + column: name.clone(), + candidates: Vec::new(), + }], + false, + TransformKind::Direct, + )); + } + StarNameDecision::Replaced(node_id) => { + let (sources, _, has_back, transform) = collect_output_sources( + node_id, + graph, + resolved, + incoming, + &mut HashSet::new(), + catalog, + mapping_cache, + ); + return Some((sources, has_back, transform)); + } + StarNameDecision::Renamed(old) => old, + StarNameDecision::Allowed => name.clone(), + }; + let wildcard_sources = named_wildcard_sources_for_scope( + target_scope, + &requested_name, + graph, + resolved, + incoming, + catalog, + mapping_cache, + ); + (!wildcard_sources.is_empty()).then_some((wildcard_sources, false, TransformKind::Direct)) } -/// Expand a Star node (qualified or unqualified) into `ColumnMapping`s. #[allow(clippy::too_many_arguments)] -fn expand_star( - table: Option<&TableRef>, +fn named_wildcard_sources_for_scope( scope: usize, + column: &str, graph: &RawGraph, resolved: &mut Vec>, incoming: &[Vec], - output_table: Option<&TableRef>, - mappings: &mut Vec, - visited_scopes: &mut HashSet, -) { - if let Some(t) = table { - let binding = graph.scopes.lookup(scope, &t.table).cloned(); - if let Some(Binding::Cte(s) | Binding::DerivedTable(s)) = binding { - expand_scope_columns(s, graph, resolved, incoming, output_table, mappings, visited_scopes); - } else { - mappings.push(wildcard_mapping(output_table, t.clone())); - } - } else { - for (_, binding) in effective_bindings(scope, graph) { - match binding { - Binding::Table(tref) => mappings.push(wildcard_mapping(output_table, tref)), - Binding::Cte(s) | Binding::DerivedTable(s) => { - expand_scope_columns(s, graph, resolved, incoming, output_table, mappings, visited_scopes); - } + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, +) -> Vec { + match graph.scopes.output_plan(scope) { + OutputPlan::Projection => { + let mut sources = Vec::new(); + for output in graph.scopes.output_columns(scope) { + let RawNode::Star { + base, + options, + scope: star_scope, + } = &graph.nodes[output.node_id] + else { + continue; + }; + let source_column = match star_name_decision(options, column, Some(base)) { + StarNameDecision::Allowed => column.to_string(), + StarNameDecision::Renamed(old) => old, + StarNameDecision::Denied + | StarNameDecision::Replaced(_) + | StarNameDecision::Ambiguous => continue, + }; + let mut expanded = Vec::new(); + expand_star( + base, + &StarOptions::default(), + *star_scope, + graph, + resolved, + incoming, + catalog, + mapping_cache, + &mut expanded, + &mut HashSet::new(), + ); + sources.extend(named_wildcard_sources_for_star( + &expanded, + &source_column, + options, + base, + *star_scope, + graph, + )); } + sources } - for &child in graph.scopes.anonymous_derived(scope) { - expand_scope_columns(child, graph, resolved, incoming, output_table, mappings, visited_scopes); + OutputPlan::Delegate(child) => named_wildcard_sources_for_scope( + *child, + column, + graph, + resolved, + incoming, + catalog, + mapping_cache, + ), + // Set-operation output names and modifiers come from the left branch, + // but lineage provenance is collected from every branch. A denied + // right branch remains an explicit uncertainty marker because its + // unknown-width slot cannot be aligned safely. + OutputPlan::SetOperation { left, right, .. } => { + let mut sources = named_wildcard_sources_for_scope( + *left, + column, + graph, + resolved, + incoming, + catalog, + mapping_cache, + ); + let right_sources = named_wildcard_sources_for_scope( + *right, + column, + graph, + resolved, + incoming, + catalog, + mapping_cache, + ); + if right_sources.is_empty() + && matches!( + scope_star_name_decision(*right, column, graph, &mut HashSet::new()), + StarNameDecision::Denied | StarNameDecision::Ambiguous + ) + { + sources.push(ColumnOrigin::Ambiguous { + column: column.to_string(), + candidates: Vec::new(), + }); + } else { + sources.extend(right_sources); + } + sources } } } -/// Recursively expand a scope's output columns into `ColumnMapping`s. -fn expand_scope_columns( - scope_id: usize, +fn named_wildcard_sources_for_star( + mappings: &[ColumnMapping], + column: &str, + options: &StarOptions, + base: &StarBase, + scope: usize, graph: &RawGraph, - resolved: &mut Vec>, - incoming: &[Vec], - output_table: Option<&TableRef>, - mappings: &mut Vec, - visited_scopes: &mut HashSet, -) { - if !visited_scopes.insert(scope_id) { - return; - } - for col in graph.scopes.output_columns(scope_id) { - if let RawNode::Star { table, scope } = &graph.nodes[col.node_id] { - expand_star(table.as_ref(), *scope, graph, resolved, incoming, output_table, mappings, visited_scopes); - } else { - let mut visited = HashSet::new(); - let (sources, edge_kinds, _) = - collect_output_sources(col.node_id, graph, resolved, incoming, &mut visited); - let transform = derive_transform(&edge_kinds); - mappings.push(ColumnMapping { - target: ColumnRef { - table: output_table.cloned(), - column: col.name.clone(), - }, - sources, - transform, - }); +) -> Vec { + let mut sources = Vec::new(); + let mut wildcard_seen = false; + let mut denied = false; + for mapping in mappings { + for source in &mapping.sources { + match source { + ColumnOrigin::Wildcard { table } => { + wildcard_seen = true; + if options.exclude.iter().any(|excluded| { + excluded_matches_source(excluded, column, source, base, scope, graph) + }) { + denied = true; + } else { + sources.push(ColumnOrigin::NamedWildcard { + table: table.clone(), + column: column.to_string(), + }); + } + } + ColumnOrigin::Recursive { base_sources } => { + for nested in base_sources { + if let ColumnOrigin::Wildcard { table } = nested { + wildcard_seen = true; + if options.exclude.iter().any(|excluded| { + excluded_matches_source( + excluded, column, nested, base, scope, graph, + ) + }) { + denied = true; + } else { + sources.push(ColumnOrigin::NamedWildcard { + table: table.clone(), + column: column.to_string(), + }); + } + } + } + } + _ => {} + } } } + if sources.is_empty() && wildcard_seen && denied { + sources.push(ColumnOrigin::Ambiguous { + column: column.to_string(), + candidates: Vec::new(), + }); + } + sources } -fn collect_output_sources( - node_id: NodeId, +fn excluded_matches_source( + excluded: &StarColumnName, + column: &str, + source: &ColumnOrigin, + base: &StarBase, + scope: usize, graph: &RawGraph, - resolved: &mut Vec>, - incoming: &[Vec], - visited: &mut HashSet, -) -> (Vec, Vec, bool) { - if !visited.insert(node_id) { - return (vec![], vec![], false); +) -> bool { + let Some((excluded_column, qualifier)) = excluded.parts.split_last() else { + return false; + }; + if !same_column_name(column, excluded_column) { + return false; + } + if qualifier.is_empty() { + return true; + } + if let StarBase::Qualified(parts) = base + && parts == qualifier + { + return true; + } + if qualifier.len() == 1 + && relation_from_origin(source).is_some_and(|table| { + graph.scopes.visible_bindings(scope).iter().any(|(name, binding)| { + name == &qualifier[0] + && matches!(binding, Binding::Table(binding_table) if binding_table == table) + }) + }) + { + return true; } + relation_from_origin(source).is_some_and(|table| table_matches_qualifier(table, qualifier)) +} - let mut sources = Vec::new(); - let mut kinds = Vec::new(); - let mut has_back = false; +fn relation_from_origin(source: &ColumnOrigin) -> Option<&TableRef> { + match source { + ColumnOrigin::Wildcard { table } + | ColumnOrigin::Concrete { table, .. } + | ColumnOrigin::NamedWildcard { table, .. } => Some(table), + _ => None, + } +} - for &edge_idx in &incoming[node_id] { - let edge = &graph.edges[edge_idx]; - if edge.is_recursive_back_edge { - has_back = true; - continue; - } - let (sub_sources, sub_back) = - collect_leaf_origins(edge.from, graph, resolved, incoming, visited); - for _ in &sub_sources { - kinds.push(edge.kind.clone()); +fn flatten_mapping_sources(sources: &[ColumnOrigin]) -> (Vec, bool) { + let mut flattened = Vec::new(); + let mut has_back = false; + for source in sources { + match source { + ColumnOrigin::Recursive { base_sources } => { + flattened.extend(base_sources.clone()); + has_back = true; + } + source => flattened.push(source.clone()), } - sources.extend(sub_sources); - has_back |= sub_back; } + (flattened, has_back) +} - (sources, kinds, has_back) +#[derive(Clone)] +enum StarNameDecision { + Allowed, + Denied, + Renamed(String), + Replaced(NodeId), + Ambiguous, } -fn collect_leaf_origins( - node_id: NodeId, +fn star_name_decision( + options: &StarOptions, + name: &str, + base: Option<&StarBase>, +) -> StarNameDecision { + if options.exclude.iter().any(|excluded| { + excluded.parts.len() == 1 && same_column_name(name, &excluded.parts[0]) + || base.is_some_and(|base| excluded_matches_name(excluded, name, base)) + }) || options + .ilike + .as_deref() + .is_some_and(|pattern| !ilike_matches(pattern, name)) + { + return StarNameDecision::Denied; + } + if let Some(replacement) = options + .replace + .iter() + .find(|replacement| same_column_name(name, &replacement.column)) + { + return StarNameDecision::Replaced(replacement.node_id); + } + if let Some((old, _)) = options + .rename + .iter() + .find(|(_, new)| same_column_name(name, new)) + { + return StarNameDecision::Renamed(old.clone()); + } + if options + .rename + .iter() + .any(|(old, _)| same_column_name(name, old)) + { + return StarNameDecision::Denied; + } + StarNameDecision::Allowed +} + +fn scope_star_name_decision( + scope: usize, + name: &str, graph: &RawGraph, - resolved: &mut Vec>, - incoming: &[Vec], - visited: &mut HashSet, -) -> (Vec, bool) { - if let Some((target_output, _)) = find_cte_redirect(node_id, graph) { - let (sources, _, has_back) = - collect_output_sources(target_output, graph, resolved, incoming, visited); - return (sources, has_back); + visited: &mut HashSet, +) -> StarNameDecision { + if !visited.insert(scope) { + return StarNameDecision::Allowed; + } + match graph.scopes.output_plan(scope) { + OutputPlan::Projection => combine_star_name_decisions( + graph + .scopes + .output_columns(scope) + .iter() + .filter_map(|column| { + if let RawNode::Star { base, options, .. } = &graph.nodes[column.node_id] { + Some(star_name_decision(options, name, Some(base))) + } else { + None + } + }) + .collect(), + ), + OutputPlan::Delegate(child) => scope_star_name_decision(*child, name, graph, visited), + // The left branch defines set-operation output names and modifiers. + OutputPlan::SetOperation { left, .. } => { + scope_star_name_decision(*left, name, graph, visited) + } } +} - if let RawNode::Output { .. } = &graph.nodes[node_id] { - let (sources, _, has_back) = - collect_output_sources(node_id, graph, resolved, incoming, visited); - (sources, has_back) - } else { - let origin = resolve_node(node_id, graph, resolved, incoming, visited); - match origin { - Some(o) => (vec![o], false), - None => (vec![], false), +fn combine_star_name_decisions(decisions: Vec) -> StarNameDecision { + let mut has_allowed = false; + let mut has_denied = false; + let mut specific: Option = None; + for decision in decisions { + match decision { + StarNameDecision::Allowed => has_allowed = true, + StarNameDecision::Denied => has_denied = true, + StarNameDecision::Ambiguous => return StarNameDecision::Ambiguous, + StarNameDecision::Renamed(old) => match &specific { + None => specific = Some(StarNameDecision::Renamed(old)), + Some(StarNameDecision::Renamed(existing)) if existing == &old => {} + Some(_) => return StarNameDecision::Ambiguous, + }, + StarNameDecision::Replaced(node_id) => match specific { + None => specific = Some(StarNameDecision::Replaced(node_id)), + Some(StarNameDecision::Replaced(existing)) if existing == node_id => {} + Some(_) => return StarNameDecision::Ambiguous, + }, } } + if has_allowed && specific.is_some() { + return StarNameDecision::Ambiguous; + } + if has_allowed { + StarNameDecision::Allowed + } else if let Some(specific) = specific { + specific + } else if has_denied { + StarNameDecision::Denied + } else { + StarNameDecision::Allowed + } } fn resolve_node( @@ -273,6 +2091,8 @@ fn resolve_node( resolved: &mut Vec>, incoming: &[Vec], visited: &mut HashSet, + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, ) -> Option { if let Some(ref origin) = resolved[node_id] { return Some(origin.clone()); @@ -283,34 +2103,96 @@ fn resolve_node( name, qualifier, scope, + binding, } => { - if let Some(qual) = qualifier { - let binding = graph.scopes.lookup(*scope, qual).cloned(); - match binding { - Some(Binding::Table(table_ref)) => Some(ColumnOrigin::Concrete { - table: table_ref, - column: name.clone(), - }), - Some(Binding::Cte(cte_scope) | Binding::DerivedTable(cte_scope)) => { - resolve_through_scope(name, cte_scope, graph, resolved, incoming, visited) - } - None => Some(ColumnOrigin::Concrete { - table: TableRef::new(qual.as_str()), - column: name.clone(), - }), - } + let binding = binding.clone().or_else(|| { + qualifier + .as_deref() + .and_then(|qual| graph.scopes.lookup(*scope, qual).cloned()) + }); + if let Some(binding) = binding { + resolve_captured_binding( + name, + binding, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ) + } else if let Some(qual) = qualifier { + Some(ColumnOrigin::Concrete { + table: TableRef::new(qual.as_str()), + column: name.clone(), + }) } else { - resolve_unqualified(name, *scope, graph, resolved, incoming, visited) + resolve_unqualified( + name, + *scope, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ) } } - RawNode::Unqualified { name, scope } => { - resolve_unqualified(name, *scope, graph, resolved, incoming, visited) + RawNode::Unqualified { + name, + scope, + binding, + } => { + if let Some(binding) = binding.clone().or_else(|| { + graph + .scopes + .lookup(*scope, name) + .filter(|binding| matches!(binding, Binding::VirtualSource(_))) + .cloned() + }) { + resolve_captured_binding( + name, + binding, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ) + } else { + resolve_unqualified( + name, + *scope, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ) + } } - RawNode::Star { table, .. } => table - .as_ref() - .map(|t| ColumnOrigin::Wildcard { table: t.clone() }), + RawNode::RowValueCandidate { + name, + scope, + binding, + } => resolve_row_value_candidate( + name, + *scope, + binding.clone(), + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ), + + RawNode::Star { .. } => None, RawNode::Output { .. } => None, }; @@ -319,18 +2201,119 @@ fn resolve_node( origin } +#[allow(clippy::too_many_arguments)] +fn resolve_row_value_candidate( + name: &str, + scope: usize, + binding: Option, + graph: &RawGraph, + resolved: &mut Vec>, + incoming: &[Vec], + visited: &mut HashSet, + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, +) -> Option { + let binding = binding.or_else(|| graph.scopes.lookup(scope, name).cloned()); + let Some(binding) = binding else { + return Some(ColumnOrigin::Ambiguous { + column: name.to_string(), + candidates: Vec::new(), + }); + }; + + match binding { + Binding::Table(table) => { + if let Some(owner) = catalog + .and_then(|catalog| catalog.resolve_column(name, std::slice::from_ref(&table))) + { + Some(ColumnOrigin::Concrete { + table: owner, + column: name.to_string(), + }) + } else { + Some(ColumnOrigin::Ambiguous { + column: name.to_string(), + candidates: Vec::new(), + }) + } + } + Binding::Cte(target_scope) | Binding::DerivedTable(target_scope) => { + let is_named_column = graph + .scopes + .output_columns(target_scope) + .iter() + .any(|column| column.name == name); + if is_named_column { + resolve_through_scope( + name, + target_scope, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ) + } else { + Some(ColumnOrigin::Ambiguous { + column: name.to_string(), + candidates: Vec::new(), + }) + } + } + Binding::VirtualSource(_) => Some(ColumnOrigin::Ambiguous { + column: name.to_string(), + candidates: Vec::new(), + }), + } +} + +#[allow(clippy::too_many_arguments)] +fn resolve_captured_binding( + name: &str, + binding: Binding, + graph: &RawGraph, + resolved: &mut Vec>, + incoming: &[Vec], + visited: &mut HashSet, + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, +) -> Option { + match binding { + Binding::Table(table) => Some(ColumnOrigin::Concrete { + table, + column: name.to_string(), + }), + Binding::Cte(scope) | Binding::DerivedTable(scope) => resolve_through_scope( + name, + scope, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ), + Binding::VirtualSource(source) => virtual_column_origin(name, source, graph), + } +} + fn find_cte_redirect(node_id: NodeId, graph: &RawGraph) -> Option<(NodeId, usize)> { match &graph.nodes[node_id] { RawNode::Ref { name, qualifier, scope, + binding, } => { - let binding = if let Some(qual) = qualifier { - graph.scopes.lookup(*scope, qual).cloned() - } else { - find_single_binding(*scope, graph) - }; + let binding = binding + .clone() + .or_else(|| { + qualifier + .as_deref() + .and_then(|qual| graph.scopes.lookup(*scope, qual).cloned()) + }) + .or_else(|| find_single_binding(*scope, graph)); match binding { Some(Binding::Cte(s) | Binding::DerivedTable(s)) => graph .scopes @@ -341,8 +2324,14 @@ fn find_cte_redirect(node_id: NodeId, graph: &RawGraph) -> Option<(NodeId, usize _ => None, } } - RawNode::Unqualified { name, scope } => { - let binding = find_single_binding(*scope, graph); + RawNode::Unqualified { + name, + scope, + binding, + } => { + let binding = binding + .clone() + .or_else(|| find_single_binding(*scope, graph)); match binding { Some(Binding::Cte(s) | Binding::DerivedTable(s)) => graph .scopes @@ -366,6 +2355,7 @@ fn find_single_binding(scope: usize, graph: &RawGraph) -> Option { } } +#[allow(clippy::too_many_arguments)] fn resolve_unqualified( name: &str, scope: usize, @@ -373,10 +2363,22 @@ fn resolve_unqualified( resolved: &mut Vec>, incoming: &[Vec], visited: &mut HashSet, + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, ) -> Option { - resolve_from_bindings(name, &effective_bindings(scope, graph), graph, resolved, incoming, visited) + resolve_from_bindings( + name, + &effective_bindings(scope, graph), + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ) } +#[allow(clippy::too_many_arguments)] fn resolve_from_bindings( name: &str, bindings: &[(String, Binding)], @@ -384,6 +2386,8 @@ fn resolve_from_bindings( resolved: &mut Vec>, incoming: &[Vec], visited: &mut HashSet, + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, ) -> Option { if bindings.len() == 1 { let (_, binding) = &bindings[0]; @@ -392,25 +2396,48 @@ fn resolve_from_bindings( table: table_ref.clone(), column: name.to_string(), }), - Binding::Cte(cte_scope) | Binding::DerivedTable(cte_scope) => { - resolve_through_scope(name, *cte_scope, graph, resolved, incoming, visited) - } + Binding::Cte(cte_scope) | Binding::DerivedTable(cte_scope) => resolve_through_scope( + name, + *cte_scope, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ), + Binding::VirtualSource(_) => None, } } else if bindings.is_empty() { - Some(ColumnOrigin::Concrete { - table: TableRef::new("?unknown?"), + Some(ColumnOrigin::Ambiguous { column: name.to_string(), + candidates: Vec::new(), }) } else { let mut table_candidates = Vec::new(); for (_, binding) in bindings { match binding { Binding::Cte(s) | Binding::DerivedTable(s) => { - if graph.scopes.output_columns(*s).iter().any(|c| c.name == name) { - return resolve_through_scope(name, *s, graph, resolved, incoming, visited); + if graph + .scopes + .output_columns(*s) + .iter() + .any(|c| c.name == name) + { + return resolve_through_scope( + name, + *s, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ); } } Binding::Table(t) => table_candidates.push(t.clone()), + Binding::VirtualSource(_) => {} } } if table_candidates.len() == 1 { @@ -427,6 +2454,32 @@ fn resolve_from_bindings( } } +fn virtual_has_column(name: &str, source: VirtualSourceId, graph: &RawGraph) -> bool { + graph + .scopes + .virtual_source(source) + .columns + .iter() + .any(|column| column.name == name) +} + +fn find_virtual_sources_for_column( + scope: usize, + name: &str, + graph: &RawGraph, +) -> Vec { + effective_bindings(scope, graph) + .into_iter() + .filter_map(|(_, binding)| match binding { + Binding::VirtualSource(source) if virtual_has_column(name, source, graph) => { + Some(source) + } + _ => None, + }) + .collect() +} + +#[allow(clippy::too_many_arguments)] fn resolve_through_scope( column_name: &str, target_scope: usize, @@ -434,15 +2487,97 @@ fn resolve_through_scope( resolved: &mut Vec>, incoming: &[Vec], visited: &mut HashSet, + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, ) -> Option { + // Resolve through the same expanded output mappings used by the public + // projection path. This is important for a qualified CTE/derived + // reference whose name was introduced by a catalog-expanded star. + let mappings = resolve_scope_mappings( + target_scope, + graph, + resolved, + incoming, + catalog, + mapping_cache, + ); + if let Some(mapping) = mappings + .iter() + .find(|mapping| mapping.target.column == column_name) + { + let (origins, has_back) = flatten_mapping_sources(&mapping.sources); + return if has_back { + Some(ColumnOrigin::Recursive { + base_sources: origins, + }) + } else { + origins.into_iter().next() + }; + } + let requested_name = + match scope_star_name_decision(target_scope, column_name, graph, &mut HashSet::new()) { + StarNameDecision::Denied | StarNameDecision::Ambiguous => { + return Some(ColumnOrigin::Ambiguous { + column: column_name.to_string(), + candidates: Vec::new(), + }); + } + StarNameDecision::Replaced(node_id) => { + let (sources, _, has_back, _) = collect_output_sources( + node_id, + graph, + resolved, + incoming, + &mut HashSet::new(), + catalog, + mapping_cache, + ); + return if has_back { + Some(ColumnOrigin::Recursive { + base_sources: sources, + }) + } else if sources.len() == 1 { + sources.into_iter().next() + } else { + Some(ColumnOrigin::Ambiguous { + column: column_name.to_string(), + candidates: Vec::new(), + }) + }; + } + StarNameDecision::Renamed(old) => old, + StarNameDecision::Allowed => column_name.to_string(), + }; + if let Some(source) = named_wildcard_sources_for_scope( + target_scope, + &requested_name, + graph, + resolved, + incoming, + catalog, + mapping_cache, + ) + .into_iter() + .next() + { + return Some(source); + } + if let Some(col) = graph .scopes .output_columns(target_scope) .iter() .find(|c| c.name == column_name) { - let (origins, _, has_back) = - collect_output_sources(col.node_id, graph, resolved, incoming, visited); + let (origins, _, has_back, _) = collect_output_sources( + col.node_id, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ); if has_back { Some(ColumnOrigin::Recursive { base_sources: origins, @@ -456,14 +2591,23 @@ fn resolve_through_scope( origins.into_iter().next() } } else { - Some(ColumnOrigin::Concrete { - table: TableRef::new("?cte?"), + Some(ColumnOrigin::Ambiguous { column: column_name.to_string(), + candidates: Vec::new(), }) } } -fn derive_transform(kinds: &[EdgeKind]) -> TransformKind { +fn derive_transform(node: &RawNode, edge_kinds: &[EdgeKind]) -> TransformKind { + let kinds = if edge_kinds.is_empty() { + match node { + RawNode::Output { intrinsic_kind, .. } => std::slice::from_ref(intrinsic_kind), + _ => edge_kinds, + } + } else { + edge_kinds + }; + if kinds.iter().any(|k| matches!(k, EdgeKind::ViaAggregation)) { TransformKind::Aggregation } else if kinds.iter().any(|k| matches!(k, EdgeKind::ViaConditional)) { @@ -474,3 +2618,56 @@ fn derive_transform(kinds: &[EdgeKind]) -> TransformKind { TransformKind::Direct } } + +#[cfg(test)] +mod tests { + use super::{reset_scope_mapping_stats, scope_mapping_computations}; + use crate::analyze; + use crate::types::{AnalyzeOptions, ColumnOrigin, Dialect}; + + #[test] + fn explicit_projection_reuses_scope_mappings() { + let columns = (0..10).map(|index| format!("c{index}")).collect::>(); + let names = columns.join(", "); + let base_projection = (0..10) + .map(|index| format!("id + 1 AS c{index}")) + .collect::>() + .join(", "); + let sql = format!( + "WITH base AS (SELECT {base_projection} FROM external_table), \ + cte0 AS (SELECT {names} FROM base), \ + cte1 AS (SELECT {names} FROM cte0), \ + cte2 AS (SELECT {names} FROM cte1), \ + cte3 AS (SELECT {names} FROM cte2), \ + cte4 AS (SELECT {names} FROM cte3) \ + SELECT {names} FROM cte4" + ); + + reset_scope_mapping_stats(); + let results = analyze( + &sql, + AnalyzeOptions { + dialect: Dialect::Generic, + ..Default::default() + }, + ) + .expect("explicit projection should resolve"); + + assert_eq!(results.len(), 1); + let mappings = &results[0].columns.mappings; + assert_eq!(mappings.len(), 10); + assert!(mappings.iter().all(|mapping| { + mapping.sources.iter().any(|source| { + matches!( + source, + ColumnOrigin::Concrete { table, column } + if table.table == "external_table" && column == "id" + ) + }) + })); + // One materialization per scope, independent of the ten requested + // columns. The exact scope count is an implementation detail, but it + // must remain bounded by the five wrappers plus the base and root. + assert!(scope_mapping_computations() <= 8); + } +} diff --git a/sqllineage/src/types.rs b/sqllineage/src/types.rs index fbdc051..673f5ce 100644 --- a/sqllineage/src/types.rs +++ b/sqllineage/src/types.rs @@ -104,6 +104,13 @@ pub struct TableLineage { #[derive(Debug, Clone, Serialize, Default)] pub struct ColumnLineage { pub mappings: Vec, + /// Whether any `SELECT *` in the statement (including nested CTEs and + /// derived tables) could not be expanded with the supplied catalog. + /// + /// This is deliberately statement-wide: an unresolved star in a JOIN + /// branch still makes the statement's schema incomplete even when the + /// selected output happens to come from another relation. + pub has_unresolved_stars: bool, } /// One output column and the source columns it derives from. @@ -119,16 +126,28 @@ pub struct ColumnMapping { /// Resolution state of a source column. #[derive(Debug, Clone, Serialize)] +#[non_exhaustive] pub enum ColumnOrigin { /// Fully resolved to a specific table and column. Concrete { table: TableRef, column: String }, - /// Multiple candidate tables; catalog needed to disambiguate. + /// Multiple candidate tables. A non-empty `candidates` list means genuine + /// ambiguity between known tables and can be disambiguated by a catalog. + /// An empty list means the column could not be resolved to any known table; + /// catalog refinement is not attempted. Ambiguous { column: String, candidates: Vec, }, /// `SELECT *` or `table.*`; catalog needed to expand. Wildcard { table: TableRef }, + /// A named output was selected through an unexpanded star in a CTE or + /// derived table. The table is known, but the physical column cannot be + /// asserted to be concrete without schema metadata. + NamedWildcard { table: TableRef, column: String }, + /// One side of a set operation contributes no column source (for example, + /// a literal projection). Kept alongside the other branch's origins so a + /// consumer cannot mistake partial lineage for complete lineage. + SourceFree { column: String }, /// Derived via recursive CTE; base case sources only. Recursive { base_sources: Vec }, } @@ -195,6 +214,7 @@ impl Default for AnalyzeOptions { /// Supported SQL dialects (maps to sqlparser dialects). #[derive(Debug, Clone, Copy, Default)] +#[non_exhaustive] pub enum Dialect { #[default] Generic, @@ -205,9 +225,32 @@ pub enum Dialect { Databricks, Snowflake, BigQuery, + DuckDb, + Redshift, + /// Trino syntax is handled by sqlparser's generic dialect because the + /// pinned sqlparser release does not expose a dedicated Trino dialect. + Trino, + Spark, + ClickHouse, + SQLite, + /// Microsoft SQL Server / T-SQL. + MsSql, +} + +impl Dialect { + /// Whether an unqualified relation alias can denote the complete row. + /// + /// `BigQuery` and PostgreSQL permit expressions such as `ARRAY_AGG(t)` when + /// `t` is a range-variable alias. Such an expression is a row/record + /// value, not a physical column named after the alias. The lineage API + /// has no whole-row origin, so callers must retain honest uncertainty + /// instead of fabricating a concrete `table.alias` source. + pub(crate) const fn supports_relation_alias_row_value(self) -> bool { + matches!(self, Self::BigQuery | Self::PostgreSql) + } } -/// Error returned when SQL parsing fails. +/// Error returned when SQL parsing or semantic validation fails. #[derive(Debug, Clone)] pub struct ParseError { pub message: String, @@ -225,6 +268,7 @@ impl std::error::Error for ParseError {} pub trait CatalogProvider { /// Return the column names of a table. Used to expand `SELECT *`. fn list_columns(&self, table: &TableRef) -> Option>; - /// Given a column name and candidate tables, return the owning table. + /// Given a column name and candidate tables, return the owning table. This + /// is only called with a non-empty candidate slice. fn resolve_column(&self, column: &str, candidates: &[TableRef]) -> Option; } diff --git a/sqllineage/tests/catalog.rs b/sqllineage/tests/catalog.rs index 73bf2a6..d80267b 100644 --- a/sqllineage/tests/catalog.rs +++ b/sqllineage/tests/catalog.rs @@ -1,7 +1,7 @@ mod common; use common::find_mapping; -use sqllineage::{AnalyzeOptions, CatalogProvider, ColumnOrigin, TableRef, analyze}; +use sqllineage::{AnalyzeOptions, CatalogProvider, ColumnOrigin, Dialect, TableRef, analyze}; struct MockCatalog; @@ -23,6 +23,46 @@ impl CatalogProvider for MockCatalog { } } +struct EagerCatalog; + +impl CatalogProvider for EagerCatalog { + fn list_columns(&self, _table: &TableRef) -> Option> { + None + } + + fn resolve_column(&self, _column: &str, _candidates: &[TableRef]) -> Option { + Some(TableRef::new("fabricated")) + } +} + +struct AliasCatalog; + +impl CatalogProvider for AliasCatalog { + fn list_columns(&self, table: &TableRef) -> Option> { + (table.table == "actual_table").then(|| vec!["id".into(), "event".into()]) + } + + fn resolve_column(&self, _column: &str, _candidates: &[TableRef]) -> Option { + None + } +} + +struct WildcardCatalog; + +impl CatalogProvider for WildcardCatalog { + fn list_columns(&self, table: &TableRef) -> Option> { + match table.table.as_str() { + "users" => Some(vec!["id".into(), "name".into(), "secret".into()]), + "other" => Some(vec!["a".into(), "b".into()]), + _ => None, + } + } + + fn resolve_column(&self, _column: &str, _candidates: &[TableRef]) -> Option { + None + } +} + fn opts_with_catalog() -> AnalyzeOptions { AnalyzeOptions { catalog: Some(Box::new(MockCatalog)), @@ -66,6 +106,646 @@ fn select_star_with_catalog_expands() { ); } +#[test] +fn wildcard_except_removes_catalog_column() { + let result = analyze( + "SELECT * EXCEPT (secret) FROM users", + AnalyzeOptions { + dialect: Dialect::BigQuery, + catalog: Some(Box::new(WildcardCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + let names = result + .columns + .mappings + .iter() + .map(|mapping| mapping.target.column.as_str()) + .collect::>(); + assert_eq!(names, ["id", "name"]); +} + +#[test] +fn qualified_wildcard_except_uses_relation_binding() { + let result = analyze( + "SELECT u.* EXCEPT (secret) FROM users AS u", + AnalyzeOptions { + dialect: Dialect::BigQuery, + catalog: Some(Box::new(WildcardCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!( + result + .columns + .mappings + .iter() + .map(|mapping| mapping.target.column.as_str()) + .collect::>(), + ["id", "name"] + ); +} + +#[test] +fn wildcard_ilike_filters_known_columns() { + let result = analyze( + "SELECT * ILIKE '%na%' FROM users", + AnalyzeOptions { + dialect: Dialect::Snowflake, + catalog: Some(Box::new(WildcardCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!( + result + .columns + .mappings + .iter() + .map(|mapping| mapping.target.column.as_str()) + .collect::>(), + ["name"] + ); +} + +#[test] +fn unknown_cte_star_does_not_recompose_excluded_names() { + let result = analyze( + "WITH x AS (SELECT * FROM unknown_source) SELECT * EXCEPT (missing) FROM x", + AnalyzeOptions { + dialect: Dialect::BigQuery, + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert!(result.columns.has_unresolved_stars); + assert!( + !result + .columns + .mappings + .iter() + .any(|mapping| mapping.target.column == "missing") + ); +} + +#[test] +fn unknown_cte_excluded_name_is_not_named_wildcard_fallback() { + let result = analyze( + "WITH x AS (SELECT * EXCEPT (secret) FROM unknown_source) SELECT secret FROM x", + AnalyzeOptions { + dialect: Dialect::BigQuery, + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert!(matches!( + result.columns.mappings[0].sources.as_slice(), + [ColumnOrigin::Ambiguous { candidates, .. }] if candidates.is_empty() + )); +} + +#[test] +fn unknown_cte_allowed_name_keeps_named_wildcard_fallback() { + let result = analyze( + "WITH x AS (SELECT * EXCEPT (secret) FROM unknown_source) SELECT id FROM x", + AnalyzeOptions { + dialect: Dialect::BigQuery, + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert!(matches!( + result.columns.mappings[0].sources.as_slice(), + [ColumnOrigin::NamedWildcard { column, .. }] if column == "id" + )); +} + +#[test] +fn unknown_cte_rename_maps_new_name_without_restoring_old_name() { + let result = analyze( + "WITH x AS (SELECT * RENAME (id AS user_id) FROM unknown_source) \ + SELECT user_id, id FROM x", + AnalyzeOptions { + dialect: Dialect::Snowflake, + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert!(matches!( + result.columns.mappings[0].sources.as_slice(), + [ColumnOrigin::NamedWildcard { column, .. }] if column == "id" + )); + assert!(matches!( + result.columns.mappings[1].sources.as_slice(), + [ColumnOrigin::Ambiguous { candidates, .. }] if candidates.is_empty() + )); +} + +#[test] +fn unknown_cte_replace_keeps_replacement_lineage() { + let result = analyze( + "WITH x AS (SELECT * REPLACE (other AS id) FROM unknown_source) \ + SELECT id FROM x", + AnalyzeOptions { + dialect: Dialect::BigQuery, + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert!(matches!( + result.columns.mappings[0].sources.as_slice(), + [ColumnOrigin::Concrete { column, .. }] if column == "other" + )); +} + +#[test] +fn wildcard_replace_uses_replacement_expression_lineage() { + let result = analyze( + "SELECT * REPLACE (name AS id) FROM users", + AnalyzeOptions { + dialect: Dialect::BigQuery, + catalog: Some(Box::new(WildcardCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + let id = find_mapping(&result.columns.mappings, "id"); + assert_eq!(concrete_sources(id), vec![("users".into(), "name".into())]); +} + +#[test] +fn wildcard_rename_preserves_position() { + let result = analyze( + "SELECT * EXCLUDE (secret) RENAME (id AS user_id) FROM users", + AnalyzeOptions { + dialect: Dialect::Snowflake, + catalog: Some(Box::new(WildcardCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!(result.columns.mappings[0].target.column, "user_id"); + assert_eq!(result.columns.mappings[1].target.column, "name"); +} + +#[test] +fn qualified_exclude_uses_qualified_relation_not_suffix_matching() { + let result = analyze( + "SELECT u.* EXCLUDE (u.secret) FROM users AS u", + AnalyzeOptions { + dialect: Dialect::Snowflake, + catalog: Some(Box::new(WildcardCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!( + result + .columns + .mappings + .iter() + .map(|mapping| mapping.target.column.as_str()) + .collect::>(), + ["id", "name"] + ); +} + +#[test] +fn qualified_exclude_width_is_used_for_set_operation_arity() { + let result = analyze( + "SELECT u.* EXCLUDE (u.secret) FROM users AS u \ + UNION ALL SELECT a, b FROM other", + AnalyzeOptions { + dialect: Dialect::Snowflake, + catalog: Some(Box::new(WildcardCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!(result.columns.mappings.len(), 2); +} + +#[test] +fn qualified_exclude_requires_exact_catalog_qualification() { + let result = analyze( + "SELECT u.* EXCLUDE (cat2.sch.users.secret) FROM cat1.sch.users AS u", + AnalyzeOptions { + dialect: Dialect::Snowflake, + catalog: Some(Box::new(WildcardCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert!( + result + .columns + .mappings + .iter() + .any(|mapping| mapping.target.column == "secret") + ); +} + +#[test] +fn qualified_field_path_star_does_not_forge_table_ref() { + let result = analyze("SELECT base.event.* FROM base", AnalyzeOptions::default()) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert!(result.columns.has_unresolved_stars); + assert!( + !result + .tables + .inputs + .iter() + .any(|table| table.table == "event") + ); + assert!(!result.columns.mappings.iter().any(|mapping| { + mapping.sources.iter().any( + |source| matches!(source, ColumnOrigin::Wildcard { table } if table.table == "event"), + ) + })); + assert!(matches!( + result.columns.mappings[0].sources.as_slice(), + [ + ColumnOrigin::Concrete { table, column }, + ColumnOrigin::Ambiguous { column: marker, .. } + ] if table.table == "base" && column == "event" && marker == "*" + )); +} + +#[test] +fn cte_field_path_star_keeps_upstream_field_ancestry() { + let result = analyze( + "WITH base AS (SELECT event FROM source) SELECT base.event.* FROM base", + AnalyzeOptions::default(), + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert!(result.columns.has_unresolved_stars); + assert!(result.columns.mappings.iter().any(|mapping| { + mapping.sources.iter().any(|source| { + matches!( + source, + ColumnOrigin::Concrete { table, column } + if table.table == "source" && column == "event" + ) + }) + })); + assert!( + !result + .tables + .inputs + .iter() + .any(|table| table.table == "event") + ); +} + +#[test] +fn derived_field_path_star_keeps_upstream_field_ancestry() { + let result = analyze( + "SELECT base.event.* FROM (SELECT event FROM source) AS base", + AnalyzeOptions::default(), + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert!(result.columns.has_unresolved_stars); + assert!(result.columns.mappings.iter().any(|mapping| { + mapping.sources.iter().any(|source| { + matches!( + source, + ColumnOrigin::Concrete { table, column } + if table.table == "source" && column == "event" + ) + }) + })); + assert!( + !result + .tables + .inputs + .iter() + .any(|table| table.table == "event") + ); +} + +#[test] +fn unknown_root_replace_keeps_barrier_and_adds_replacement_mapping() { + let result = analyze( + "SELECT * REPLACE (other AS id) FROM unknown_source", + AnalyzeOptions { + dialect: Dialect::BigQuery, + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert!(result.columns.has_unresolved_stars); + assert!( + result + .columns + .mappings + .iter() + .any(|mapping| mapping.target.column == "*") + ); + assert!(result.columns.mappings.iter().any(|mapping| { + mapping.target.column == "id" + && matches!( + mapping.sources.as_slice(), + [ColumnOrigin::Concrete { column, .. }] if column == "other" + ) + })); +} + +#[test] +fn unknown_root_rename_keeps_barrier_and_adds_named_wildcard() { + let result = analyze( + "SELECT * RENAME (id AS user_id) FROM unknown_source", + AnalyzeOptions { + dialect: Dialect::Snowflake, + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert!(result.columns.has_unresolved_stars); + assert!(result.columns.mappings.iter().any(|mapping| { + mapping.target.column == "user_id" + && matches!( + mapping.sources.as_slice(), + [ColumnOrigin::NamedWildcard { column, .. }] if column == "id" + ) + })); +} + +#[test] +fn multiple_unknown_stars_keep_name_available_from_unfiltered_star() { + let result = analyze( + "WITH x AS (SELECT *, * EXCEPT (secret) FROM unknown_source) SELECT secret FROM x", + AnalyzeOptions { + dialect: Dialect::BigQuery, + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert!(matches!( + result.columns.mappings[0].sources.as_slice(), + [ColumnOrigin::NamedWildcard { column, .. }] if column == "secret" + )); +} + +#[test] +fn unknown_qualified_exclude_does_not_restore_excluded_name() { + let result = analyze( + "WITH x AS (SELECT u.* EXCLUDE (u.secret) FROM unknown_source AS u) \ + SELECT secret FROM x", + AnalyzeOptions { + dialect: Dialect::Snowflake, + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert!(matches!( + result.columns.mappings[0].sources.as_slice(), + [ColumnOrigin::Ambiguous { candidates, .. }] if candidates.is_empty() + )); +} + +#[test] +fn unqualified_qualified_exclude_keeps_other_relation_wildcard() { + let result = analyze( + "WITH x AS (SELECT * EXCLUDE (u.secret) FROM unknown_source AS u, other_source AS v) \ + SELECT secret FROM x", + AnalyzeOptions { + dialect: Dialect::Snowflake, + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert!(matches!( + result.columns.mappings[0].sources.as_slice(), + [ColumnOrigin::NamedWildcard { table, column }] if table.table == "other_source" && column == "secret" + )); +} + +#[test] +fn set_operation_name_decision_uses_left_branch_only() { + let result = analyze( + "WITH x AS (SELECT * FROM unknown_left UNION ALL SELECT * EXCEPT (secret) FROM unknown_right) \ + SELECT secret FROM x", + AnalyzeOptions { + dialect: Dialect::BigQuery, + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + let sources = &result.columns.mappings[0].sources; + assert!(sources.iter().any(|source| { + matches!(source, ColumnOrigin::NamedWildcard { table, column } + if table.table == "unknown_left" && column == "secret") + })); + assert!(sources.iter().any(|source| { + matches!(source, ColumnOrigin::Ambiguous { column, candidates } + if column == "secret" && candidates.is_empty()) + })); +} + +#[test] +fn set_operation_keeps_named_wildcard_sources_from_both_branches() { + let result = analyze( + "WITH x AS (SELECT * FROM unknown_left UNION ALL SELECT * FROM unknown_right) \ + SELECT secret FROM x", + AnalyzeOptions { + dialect: Dialect::BigQuery, + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + let sources = &result.columns.mappings[0].sources; + for table_name in ["unknown_left", "unknown_right"] { + assert!(sources.iter().any(|source| { + matches!(source, ColumnOrigin::NamedWildcard { table, column } + if table.table == table_name && column == "secret") + })); + } + assert!( + !sources + .iter() + .any(|source| matches!(source, ColumnOrigin::Ambiguous { .. })) + ); +} + +#[test] +fn expr_qualified_star_is_an_unknown_shape_set_barrier() { + let result = analyze( + "SELECT STRUCT(1 AS value).* UNION ALL SELECT id FROM known", + AnalyzeOptions { + dialect: Dialect::BigQuery, + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert!(result.columns.has_unresolved_stars); + assert!(result.columns.mappings.iter().any(|mapping| { + mapping.target.column == "*" + && mapping.sources.iter().any( + |source| matches!(source, ColumnOrigin::Ambiguous { column, .. } if column == "*"), + ) + })); +} + +fn assert_qualified_alias_star_expands(sql: &str) { + let result = analyze( + sql, + AnalyzeOptions { + catalog: Some(Box::new(AliasCatalog)), + dialect: Dialect::Generic, + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!(result.columns.mappings.len(), 2); + assert!(result.columns.mappings.iter().all(|mapping| { + mapping + .sources + .iter() + .all(|source| matches!(source, ColumnOrigin::Concrete { .. })) + })); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "id")), + vec![("actual_table".into(), "id".into())] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "event")), + vec![("actual_table".into(), "event".into())] + ); +} + +#[test] +fn qualified_alias_star_uses_catalog_table_binding() { + for sql in [ + "SELECT a.* FROM actual_table AS a", + "WITH x AS (SELECT a.* FROM actual_table AS a) SELECT * FROM x", + ] { + assert_qualified_alias_star_expands(sql); + } +} + +#[test] +fn qualified_alias_star_without_catalog_keeps_actual_table_wildcard() { + let result = analyze( + "SELECT a.* FROM actual_table AS a", + AnalyzeOptions { + dialect: Dialect::Generic, + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!(result.columns.mappings.len(), 1); + match &result.columns.mappings[0].sources[0] { + ColumnOrigin::Wildcard { table } => assert_eq!(table.table, "actual_table"), + other => panic!("expected Wildcard, got {other:?}"), + } +} + #[test] fn select_star_without_catalog_preserved() { let result = analyze("SELECT * FROM users", AnalyzeOptions::default()) @@ -110,6 +790,29 @@ fn ambiguous_column_without_catalog() { } } +#[test] +fn catalog_does_not_fabricate_unresolved_column_owner() { + let result = analyze( + "SELECT missing", + AnalyzeOptions { + catalog: Some(Box::new(EagerCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + let m = find_mapping(&result.columns.mappings, "missing"); + match &m.sources[0] { + ColumnOrigin::Ambiguous { column, candidates } => { + assert_eq!(column, "missing"); + assert!(candidates.is_empty()); + } + other => panic!("expected Ambiguous, got {other:?}"), + } +} + #[test] fn catalog_preserves_qualified_columns() { let sql = @@ -128,3 +831,270 @@ fn catalog_preserves_qualified_columns() { vec![("orders".into(), "amount".into())] ); } + +struct SetOperationCatalog; + +impl CatalogProvider for SetOperationCatalog { + fn list_columns(&self, table: &TableRef) -> Option> { + match table.table.as_str() { + "users" => Some(vec!["id".into(), "name".into(), "email".into()]), + "other" => Some(vec!["a".into(), "b".into(), "c".into(), "d".into()]), + "ext_a" => Some(vec!["col_x".into(), "col_y".into()]), + _ => None, + } + } + + fn resolve_column(&self, _column: &str, _candidates: &[TableRef]) -> Option { + None + } +} + +#[test] +fn set_operation_expands_leading_star_before_positional_merge() { + let sql = "SELECT * FROM users UNION ALL SELECT a, b, c FROM other"; + let result = analyze( + sql, + AnalyzeOptions { + catalog: Some(Box::new(SetOperationCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!(result.columns.mappings.len(), 3); + assert_eq!(result.columns.mappings[0].target.column, "id"); + assert_eq!(result.columns.mappings[1].target.column, "name"); + assert_eq!(result.columns.mappings[2].target.column, "email"); + assert_eq!( + concrete_sources(&result.columns.mappings[0]), + vec![("other".into(), "a".into()), ("users".into(), "id".into())] + ); + assert_eq!( + concrete_sources(&result.columns.mappings[1]), + vec![ + ("other".into(), "b".into()), + ("users".into(), "name".into()) + ] + ); +} + +#[test] +fn set_operation_preserves_non_leading_star_contribution() { + let sql = "SELECT id, * FROM users UNION ALL SELECT a, b, c, d FROM other"; + let result = analyze( + sql, + AnalyzeOptions { + catalog: Some(Box::new(SetOperationCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!(result.columns.mappings.len(), 4); + assert_eq!( + concrete_sources(&result.columns.mappings[1]), + vec![("other".into(), "b".into()), ("users".into(), "id".into())] + ); + assert_eq!( + concrete_sources(&result.columns.mappings[3]), + vec![ + ("other".into(), "d".into()), + ("users".into(), "email".into()) + ] + ); +} + +#[test] +fn set_operation_branches_survive_cte_and_derived_boundaries() { + let sql = "WITH combined AS (SELECT * FROM users UNION ALL SELECT a, b, c FROM other) \ + SELECT * FROM (SELECT * FROM combined) derived"; + let result = analyze( + sql, + AnalyzeOptions { + catalog: Some(Box::new(SetOperationCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!(result.columns.mappings.len(), 3); + assert_eq!( + concrete_sources(&result.columns.mappings[1]), + vec![ + ("other".into(), "b".into()), + ("users".into(), "name".into()) + ] + ); +} + +#[test] +fn non_leading_star_in_right_set_branch_contributes_to_named_output() { + let sql = "WITH lit AS (SELECT 1 AS col_a), \ + u AS (SELECT col_a FROM lit UNION ALL SELECT * FROM ext_a) \ + SELECT col_a FROM u"; + let result = analyze( + sql, + AnalyzeOptions { + catalog: Some(Box::new(SetOperationCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!(result.columns.mappings.len(), 1); + assert!(matches!( + result.columns.mappings[0].sources.as_slice(), + [ColumnOrigin::Concrete { table, column }, ColumnOrigin::SourceFree { column: marker }] + if table.table == "ext_a" && column == "col_x" && marker == "col_a" + )); +} + +#[test] +fn named_lookup_through_set_operation_cte_keeps_all_branches() { + let sql = "WITH combined AS (SELECT * FROM users UNION ALL SELECT a, b, c FROM other) \ + SELECT name FROM combined"; + let result = analyze( + sql, + AnalyzeOptions { + catalog: Some(Box::new(SetOperationCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!( + concrete_sources(&result.columns.mappings[0]), + vec![ + ("other".into(), "b".into()), + ("users".into(), "name".into()) + ] + ); +} + +#[test] +fn named_lookup_through_projection_cte_and_derived_star_chain() { + let sql = "WITH base AS (SELECT * FROM users), wrapped AS (SELECT * FROM base) \ + SELECT name FROM (SELECT * FROM wrapped) derived"; + let result = analyze( + sql, + AnalyzeOptions { + catalog: Some(Box::new(SetOperationCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!(result.columns.mappings.len(), 1); + assert_eq!( + concrete_sources(&result.columns.mappings[0]), + vec![("users".into(), "name".into())] + ); +} + +#[test] +fn named_lookup_through_unknown_projection_star_is_indeterminate() { + let result = analyze( + "WITH base AS (SELECT * FROM unknown) SELECT name FROM base", + AnalyzeOptions::default(), + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!(result.columns.mappings.len(), 1); + assert!(matches!( + result.columns.mappings[0].sources.as_slice(), + [ColumnOrigin::NamedWildcard { table, column }] + if table.table == "unknown" && column == "name" + )); +} + +#[test] +fn named_lookup_through_projection_preserves_inner_transform() { + let result = analyze( + "WITH aggregated AS (SELECT SUM(amount) AS total FROM orders) \ + SELECT total FROM aggregated", + AnalyzeOptions { + catalog: Some(Box::new(SetOperationCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!(result.columns.mappings.len(), 1); + assert_eq!( + concrete_sources(&result.columns.mappings[0]), + vec![("orders".into(), "amount".into())] + ); + assert_eq!( + result.columns.mappings[0].transform, + sqllineage::TransformKind::Aggregation + ); +} + +#[test] +fn named_lookup_through_projection_preserves_inner_expression_transform() { + let result = analyze( + "WITH transformed AS (SELECT amount + 1 AS adjusted FROM orders) \ + SELECT adjusted FROM transformed", + AnalyzeOptions { + catalog: Some(Box::new(SetOperationCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!(result.columns.mappings.len(), 1); + assert_eq!( + concrete_sources(&result.columns.mappings[0]), + vec![("orders".into(), "amount".into())] + ); + assert_eq!( + result.columns.mappings[0].transform, + sqllineage::TransformKind::Expression + ); +} + +#[test] +fn catalog_known_set_arity_mismatch_is_an_analysis_error() { + let result = analyze( + "SELECT * FROM users UNION ALL SELECT a, b, c, d FROM other", + AnalyzeOptions { + catalog: Some(Box::new(SetOperationCatalog)), + ..AnalyzeOptions::default() + }, + ); + let error = match result { + Ok(_) => panic!("catalog-known arity mismatch should not be truncated"), + Err(error) => error, + }; + assert_eq!( + error.message, + "set operation arity mismatch: left has 3 columns, right has 4 columns" + ); +} diff --git a/sqllineage/tests/column_lineage.rs b/sqllineage/tests/column_lineage.rs index 3418c24..3bc41b2 100644 --- a/sqllineage/tests/column_lineage.rs +++ b/sqllineage/tests/column_lineage.rs @@ -1,7 +1,7 @@ mod common; use common::{analyze_one, concrete_sources, find_mapping, table}; -use sqllineage::TransformKind; +use sqllineage::{ColumnOrigin, TransformKind}; #[test] fn select_columns() { @@ -18,6 +18,19 @@ fn select_columns() { assert_eq!(m_b.transform, TransformKind::Direct); } +#[test] +fn unresolved_column_has_empty_ambiguous_candidates() { + let result = analyze_one("SELECT missing"); + let m = find_mapping(&result.columns.mappings, "missing"); + match &m.sources[0] { + ColumnOrigin::Ambiguous { column, candidates } => { + assert_eq!(column, "missing"); + assert!(candidates.is_empty()); + } + other => panic!("expected Ambiguous, got {other:?}"), + } +} + #[test] fn select_expression() { let result = analyze_one("SELECT a + b AS c FROM t"); @@ -68,6 +81,15 @@ fn select_aggregate() { assert_eq!(m.transform, TransformKind::Aggregation); } +#[test] +fn select_count_star_is_aggregation_without_sources() { + let result = analyze_one("SELECT COUNT(*) AS c FROM t"); + let m = find_mapping(&result.columns.mappings, "c"); + + assert!(m.sources.is_empty()); + assert_eq!(m.transform, TransformKind::Aggregation); +} + #[test] fn select_multiple_tables_qualified() { let result = analyze_one("SELECT t1.a, t2.b FROM t1 JOIN t2 ON t1.id = t2.id"); @@ -80,6 +102,46 @@ fn select_multiple_tables_qualified() { assert_eq!(concrete_sources(m_b), vec![("t2".into(), "b".into())]); } +#[test] +fn duplicate_output_names_preserve_projection_order() { + let result = analyze_one("SELECT a.id, b.id FROM a JOIN b ON a.id = b.bid"); + let sources: Vec<_> = result + .columns + .mappings + .iter() + .map(concrete_sources) + .collect(); + + assert_eq!( + sources, + vec![ + vec![("a".into(), "id".into())], + vec![("b".into(), "id".into())] + ] + ); +} + +#[test] +fn three_duplicate_output_names_preserve_projection_order() { + let result = + analyze_one("SELECT a.id, b.id, c.id FROM a JOIN b ON a.id = b.bid JOIN c ON a.id = c.cid"); + let sources: Vec<_> = result + .columns + .mappings + .iter() + .map(concrete_sources) + .collect(); + + assert_eq!( + sources, + vec![ + vec![("a".into(), "id".into())], + vec![("b".into(), "id".into())], + vec![("c".into(), "id".into())], + ] + ); +} + #[test] fn select_case_expression() { let result = analyze_one("SELECT CASE WHEN a > 0 THEN b ELSE c END AS d FROM t"); @@ -102,3 +164,113 @@ fn select_cast_passthrough() { assert_eq!(concrete_sources(m), vec![("t".into(), "a".into())]); assert_eq!(m.transform, TransformKind::Direct); } + +#[test] +fn unnest_source_free_alias_has_no_physical_sources() { + let result = analyze_one( + "SELECT item FROM UNNEST(GENERATE_DATE_ARRAY(DATE('2020-01-01'), DATE('2020-01-03'))) AS item", + ); + let mapping = find_mapping(&result.columns.mappings, "item"); + assert!(mapping.sources.is_empty()); + assert_eq!(result.tables.inputs, Vec::::new()); +} + +#[test] +fn unnest_alias_depends_on_array_column() { + let result = analyze_one("SELECT item FROM base, UNNEST(base.items_array) AS item"); + let mapping = find_mapping(&result.columns.mappings, "item"); + assert_eq!( + concrete_sources(mapping), + vec![("base".into(), "items_array".into())] + ); +} + +#[test] +fn unnest_unresolved_array_is_ambiguous_not_alias_column() { + let result = analyze_one("SELECT item FROM UNNEST(missing_array) AS item"); + let mapping = find_mapping(&result.columns.mappings, "item"); + assert!(matches!( + mapping.sources.as_slice(), + [ColumnOrigin::Ambiguous { column, candidates }] if column == "item" && candidates.is_empty() + )); +} + +#[test] +fn unnest_unqualified_array_column_keeps_prior_table_binding() { + let result = analyze_one("SELECT item FROM base, UNNEST(items_array) AS item"); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "item")), + vec![("base".into(), "items_array".into())] + ); +} + +#[test] +fn unnest_known_empty_virtual_dependency_stays_source_free() { + let result = analyze_one("SELECT item FROM UNNEST([1, 2]) AS source, UNNEST(source) AS item"); + assert!( + find_mapping(&result.columns.mappings, "item") + .sources + .is_empty() + ); +} + +#[test] +fn unqualified_identifier_does_not_capture_relation_alias() { + let result = analyze_one("SELECT a FROM table1 AS a, table2 AS b"); + let mapping = find_mapping(&result.columns.mappings, "a"); + assert!(matches!( + mapping.sources.as_slice(), + [ColumnOrigin::Ambiguous { column, candidates }] + if column == "a" && candidates.len() == 2 + )); +} + +#[test] +fn duplicate_virtual_slots_are_ambiguous_but_qualified_slots_resolve() { + let result = + analyze_one("SELECT x FROM base, UNNEST(base.first) AS u(x), UNNEST(base.second) AS v(x)"); + assert!(matches!( + find_mapping(&result.columns.mappings, "x").sources.as_slice(), + [ColumnOrigin::Ambiguous { column, candidates }] if column == "x" && candidates.is_empty() + )); + + let result = analyze_one( + "SELECT u.x, v.x FROM base, UNNEST(base.first) AS u(x), UNNEST(base.second) AS v(x)", + ); + assert_eq!( + concrete_sources(&result.columns.mappings[0]), + vec![("base".into(), "first".into())] + ); + assert_eq!( + concrete_sources(&result.columns.mappings[1]), + vec![("base".into(), "second".into())] + ); +} + +#[test] +fn unnest_alias_columns_keep_array_expression_ordinals() { + let result = analyze_one("SELECT x, y FROM base, UNNEST(base.first, base.second) AS u(x, y)"); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "x")), + vec![("base".into(), "first".into())] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "y")), + vec![("base".into(), "second".into())] + ); +} + +#[test] +fn unnest_offset_is_a_source_free_generated_slot() { + let result = analyze_one("SELECT item, off FROM UNNEST([1, 2]) AS item WITH OFFSET AS off"); + assert!( + find_mapping(&result.columns.mappings, "item") + .sources + .is_empty() + ); + assert!( + find_mapping(&result.columns.mappings, "off") + .sources + .is_empty() + ); +} diff --git a/sqllineage/tests/cte.rs b/sqllineage/tests/cte.rs index 24bdfbd..df470e6 100644 --- a/sqllineage/tests/cte.rs +++ b/sqllineage/tests/cte.rs @@ -1,7 +1,19 @@ mod common; use common::{analyze_one, concrete_sources, find_mapping, table}; -use sqllineage::{ColumnOrigin, TableRef, TransformKind}; +use sqllineage::{ColumnMapping, ColumnOrigin, TableRef, TransformKind}; + +fn has_wildcard_from(mapping: &ColumnMapping, table_name: &str) -> bool { + mapping.sources.iter().any(|source| { + match source { + ColumnOrigin::Wildcard { table } => table.table == table_name, + ColumnOrigin::Recursive { base_sources } => base_sources.iter().any(|source| { + matches!(source, ColumnOrigin::Wildcard { table } if table.table == table_name) + }), + _ => false, + } + }) +} #[test] fn single_cte() { @@ -15,6 +27,20 @@ fn single_cte() { assert_eq!(m.transform, TransformKind::Direct); } +#[test] +fn missing_column_from_cte_has_empty_ambiguous_candidates() { + let sql = "WITH cte AS (SELECT present FROM source) SELECT missing FROM cte"; + let result = analyze_one(sql); + let m = find_mapping(&result.columns.mappings, "missing"); + match &m.sources[0] { + ColumnOrigin::Ambiguous { column, candidates } => { + assert_eq!(column, "missing"); + assert!(candidates.is_empty()); + } + other => panic!("expected Ambiguous, got {other:?}"), + } +} + #[test] fn cte_chain() { let sql = "WITH a AS (SELECT x FROM t), b AS (SELECT x FROM a) SELECT x FROM b"; @@ -71,6 +97,12 @@ fn recursive_cte_base_case() { match &m.sources[0] { ColumnOrigin::Recursive { base_sources } => { assert!(!base_sources.is_empty()); + assert!(base_sources.iter().all(|source| { + !matches!( + source, + ColumnOrigin::Concrete { table, .. } if table.table == "cte" + ) + })); match &base_sources[0] { ColumnOrigin::Concrete { table, column } => { assert_eq!(table.table, "t"); @@ -151,6 +183,187 @@ fn union_all_columns() { ); } +#[test] +fn union_keeps_left_names_and_merges_explicit_columns_positionally() { + let sql = "SELECT a AS left_name, b AS second_name FROM t1 \ + UNION ALL SELECT c AS right_name, d AS other_name FROM t2"; + let result = analyze_one(sql); + assert_eq!(result.columns.mappings.len(), 2); + assert_eq!(result.columns.mappings[0].target.column, "left_name"); + assert_eq!(result.columns.mappings[1].target.column, "second_name"); + assert_eq!( + concrete_sources(&result.columns.mappings[0]), + vec![("t1".into(), "a".into()), ("t2".into(), "c".into())] + ); + assert_eq!( + concrete_sources(&result.columns.mappings[1]), + vec![("t1".into(), "b".into()), ("t2".into(), "d".into())] + ); +} + +#[test] +fn nested_union_preserves_outer_left_names_and_all_branch_sources() { + let sql = "SELECT a AS first_name FROM t1 \ + UNION ALL SELECT b AS second_name FROM t2 \ + UNION ALL SELECT c AS third_name FROM t3"; + let result = analyze_one(sql); + assert_eq!(result.columns.mappings.len(), 1); + assert_eq!(result.columns.mappings[0].target.column, "first_name"); + assert_eq!( + concrete_sources(&result.columns.mappings[0]), + vec![ + ("t1".into(), "a".into()), + ("t2".into(), "b".into()), + ("t3".into(), "c".into()) + ] + ); +} + +#[test] +fn union_transform_prefers_aggregation_across_branches() { + let sql = "SELECT SUM(a) AS value FROM t1 UNION ALL SELECT b AS other_value FROM t2"; + let result = analyze_one(sql); + assert_eq!(result.columns.mappings.len(), 1); + assert_eq!(result.columns.mappings[0].target.column, "value"); + assert_eq!( + concrete_sources(&result.columns.mappings[0]), + vec![("t1".into(), "a".into()), ("t2".into(), "b".into())] + ); + assert_eq!( + result.columns.mappings[0].transform, + TransformKind::Aggregation + ); +} + +#[test] +fn unknown_leading_star_is_preserved_without_catalog() { + let result = analyze_one("SELECT * FROM unknown_left UNION ALL SELECT a, b FROM known"); + assert_eq!(result.columns.mappings.len(), 3); + match &result.columns.mappings[0].sources[0] { + ColumnOrigin::Wildcard { table } => assert_eq!(table.table, "unknown_left"), + other => panic!("expected wildcard, got {other:?}"), + } + assert!(has_wildcard_from( + &result.columns.mappings[1], + "unknown_left" + )); + assert!(has_wildcard_from( + &result.columns.mappings[2], + "unknown_left" + )); +} + +#[test] +fn unknown_non_leading_star_does_not_drop_known_branch_columns() { + let result = analyze_one("SELECT id, * FROM unknown_left UNION ALL SELECT a, b, c FROM known"); + assert_eq!(result.columns.mappings.len(), 4); + assert!(matches!( + result.columns.mappings[0].sources.as_slice(), + [ColumnOrigin::Concrete { .. }, ColumnOrigin::Concrete { .. }] + )); + match &result.columns.mappings[1].sources[0] { + ColumnOrigin::Wildcard { table } => assert_eq!(table.table, "unknown_left"), + other => panic!("expected wildcard, got {other:?}"), + } + assert!(has_wildcard_from( + &result.columns.mappings[2], + "unknown_left" + )); + assert!(has_wildcard_from( + &result.columns.mappings[3], + "unknown_left" + )); + assert_eq!(result.columns.mappings[2].target.column, "b"); + assert_eq!(result.columns.mappings[3].target.column, "c"); +} + +#[test] +fn unknown_right_star_marks_known_left_tail_as_unresolved() { + let result = analyze_one("SELECT a, b, c FROM known UNION ALL SELECT * FROM unknown_right"); + assert_eq!(result.columns.mappings.len(), 4); + for mapping in &result.columns.mappings[..3] { + assert!(has_wildcard_from(mapping, "unknown_right")); + } +} + +#[test] +fn unknown_stars_on_both_set_branches_are_both_retained() { + let result = analyze_one("SELECT * FROM unknown_left UNION ALL SELECT * FROM unknown_right"); + assert_eq!(result.columns.mappings.len(), 2); + for (mapping, table_name) in result + .columns + .mappings + .iter() + .zip(["unknown_left", "unknown_right"]) + { + assert!(has_wildcard_from(mapping, table_name)); + assert!(has_wildcard_from( + mapping, + if table_name == "unknown_left" { + "unknown_right" + } else { + "unknown_left" + } + )); + } +} + +#[test] +fn nested_unknown_set_keeps_every_branch_mapping() { + let result = analyze_one( + "SELECT * FROM unknown_left UNION ALL SELECT a FROM known UNION ALL SELECT * FROM unknown_right", + ); + assert_eq!(result.columns.mappings.len(), 3); + assert!(matches!( + result.columns.mappings[0].sources[0], + ColumnOrigin::Wildcard { .. } + )); + assert!(matches!( + result.columns.mappings[2].sources[0], + ColumnOrigin::Wildcard { .. } + )); + assert!(has_wildcard_from( + &result.columns.mappings[1], + "unknown_left" + )); + assert!(has_wildcard_from( + &result.columns.mappings[1], + "unknown_right" + )); +} + +#[test] +fn leading_unknown_star_hides_nonleading_only_set_names() { + let result = analyze_one( + "SELECT * FROM unknown_source \ + UNION ALL SELECT id, amt AS total FROM known_table \ + UNION ALL SELECT id, fee FROM third_table", + ); + let names = result + .columns + .mappings + .iter() + .map(|mapping| mapping.target.column.as_str()) + .collect::>(); + assert_eq!(names, vec!["*", "id", "total"]); +} + +#[test] +fn exact_set_arity_mismatch_is_an_analysis_error() { + let result = sqllineage::analyze( + "SELECT a FROM t1 UNION ALL SELECT b, c FROM t2", + sqllineage::AnalyzeOptions::default(), + ); + let error = match result { + Ok(_) => panic!("exact arity mismatch should not be truncated"), + Err(error) => error, + }; + assert_eq!( + error.message, + "set operation arity mismatch: left has 1 columns, right has 2 columns" + ); +} + #[test] fn union_inside_cte() { let sql = "\ diff --git a/sqllineage/tests/dialect.rs b/sqllineage/tests/dialect.rs new file mode 100644 index 0000000..fec2f34 --- /dev/null +++ b/sqllineage/tests/dialect.rs @@ -0,0 +1,34 @@ +use sqllineage::{AnalyzeOptions, Dialect, StatementType, analyze}; + +#[test] +fn all_public_dialects_parse_basic_queries() { + let dialects = [ + Dialect::Generic, + Dialect::Ansi, + Dialect::PostgreSql, + Dialect::MySql, + Dialect::Hive, + Dialect::Databricks, + Dialect::Snowflake, + Dialect::BigQuery, + Dialect::DuckDb, + Dialect::Redshift, + Dialect::Trino, + Dialect::Spark, + Dialect::ClickHouse, + Dialect::SQLite, + Dialect::MsSql, + ]; + + for dialect in dialects { + let result = analyze( + "SELECT 1", + AnalyzeOptions { + dialect, + ..AnalyzeOptions::default() + }, + ) + .expect("basic query should parse for every public dialect"); + assert_eq!(result[0].statement_type, StatementType::Query); + } +} diff --git a/sqllineage/tests/expr_coverage.rs b/sqllineage/tests/expr_coverage.rs index 181e796..35553ac 100644 --- a/sqllineage/tests/expr_coverage.rs +++ b/sqllineage/tests/expr_coverage.rs @@ -1,7 +1,23 @@ mod common; use common::{analyze_one, concrete_sources, find_mapping}; -use sqllineage::TransformKind; +use sqllineage::{ + AnalyzeOptions, CatalogProvider, ColumnOrigin, Dialect, TableRef, TransformKind, analyze, +}; + +fn analyze_with_dialect(sql: &str, dialect: Dialect) -> sqllineage::AnalyzeResult { + analyze( + sql, + AnalyzeOptions { + dialect, + ..AnalyzeOptions::default() + }, + ) + .expect("SQL should parse") + .into_iter() + .next() + .unwrap_or_default() +} #[test] fn extract_year() { @@ -107,3 +123,765 @@ fn json_access() { let m = find_mapping(&result.columns.mappings, "val"); assert_eq!(concrete_sources(m), vec![("t".into(), "data".into())]); } + +#[test] +fn qualified_compound_field_access_uses_binding_column() { + let result = analyze_one("SELECT base.items_array[1] AS item FROM actual_table AS base"); + let m = find_mapping(&result.columns.mappings, "item"); + assert_eq!( + concrete_sources(m), + vec![("actual_table".into(), "items_array".into())] + ); +} + +#[test] +fn compound_field_access_retains_column_dependent_index() { + let result = analyze_one("SELECT base.items_array[idx] AS item FROM actual_table AS base"); + let m = find_mapping(&result.columns.mappings, "item"); + assert_eq!( + concrete_sources(m), + vec![ + ("actual_table".into(), "idx".into()), + ("actual_table".into(), "items_array".into()), + ] + ); +} + +#[test] +fn nested_qualified_compound_field_access_keeps_top_level_column() { + let result = analyze_one("SELECT base.payload.items[1] AS item FROM actual_table AS base"); + let m = find_mapping(&result.columns.mappings, "item"); + assert_eq!( + concrete_sources(m), + vec![("actual_table".into(), "payload".into())] + ); +} + +#[test] +fn cte_compound_field_access_uses_cte_binding_column() { + let result = analyze_one( + "WITH base AS (SELECT items_array FROM actual_table) SELECT base.items_array[1] AS item FROM base", + ); + let m = find_mapping(&result.columns.mappings, "item"); + assert_eq!( + concrete_sources(m), + vec![("actual_table".into(), "items_array".into())] + ); +} + +#[test] +fn unqualified_compound_field_access_uses_top_level_column() { + let result = analyze_one("SELECT payload.items[1] AS item FROM t"); + let m = find_mapping(&result.columns.mappings, "item"); + assert_eq!(concrete_sources(m), vec![("t".into(), "payload".into())]); +} + +#[test] +fn compound_identifier_without_visible_binding_keeps_qualified_relation_fallback() { + let result = analyze_one("SELECT orders.id AS id"); + let m = find_mapping(&result.columns.mappings, "id"); + match m.sources.as_slice() { + [ColumnOrigin::Concrete { table, column }] => { + assert_eq!(table.catalog, None); + assert_eq!(table.schema, None); + assert_eq!(table.table, "orders"); + assert_eq!(column, "id"); + } + other => panic!("expected structured orders.id source, got {other:?}"), + } +} + +#[test] +fn compound_identifier_without_visible_binding_preserves_relation_parts() { + for (sql, catalog, schema, table) in [ + ("SELECT raw.orders.id AS id", None, Some("raw"), "orders"), + ( + "SELECT warehouse.raw.orders.id AS id", + Some("warehouse"), + Some("raw"), + "orders", + ), + ] { + let result = analyze_one(sql); + let m = find_mapping(&result.columns.mappings, "id"); + match m.sources.as_slice() { + [ + ColumnOrigin::Concrete { + table: source_table, + column, + }, + ] => { + assert_eq!(source_table.catalog.as_deref(), catalog); + assert_eq!(source_table.schema.as_deref(), schema); + assert_eq!(source_table.table, table); + assert_eq!(column, "id"); + } + other => panic!("expected structured relation source, got {other:?}"), + } + } +} + +#[test] +fn bigquery_offset_compound_field_access_uses_binding_column() { + let result = analyze_with_dialect( + "SELECT base.items_array[OFFSET(0)] AS item FROM actual_table AS base", + Dialect::BigQuery, + ); + let m = find_mapping(&result.columns.mappings, "item"); + assert_eq!( + concrete_sources(m), + vec![("actual_table".into(), "items_array".into())] + ); +} + +#[test] +fn qualified_struct_field_access_uses_binding_column() { + let result = analyze_with_dialect( + "SELECT agg.event.qualified_field AS field FROM upstream_model AS agg", + Dialect::BigQuery, + ); + let m = find_mapping(&result.columns.mappings, "field"); + assert_eq!( + concrete_sources(m), + vec![("upstream_model".into(), "event".into())] + ); +} + +#[test] +fn bigquery_date_trunc_week_modifier_is_syntax_only() { + let result = analyze_with_dialect( + "SELECT DATE_TRUNC(event_date, WEEK(MONDAY)) AS monday_start, DATE_TRUNC(event_date, WEEK(SUNDAY)) AS sunday_start FROM events", + Dialect::BigQuery, + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "monday_start")), + vec![("events".into(), "event_date".into())] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "sunday_start")), + vec![("events".into(), "event_date".into())] + ); +} + +#[test] +fn unqualified_struct_field_access_uses_top_level_column() { + let result = analyze_with_dialect( + "SELECT event.bare_field AS field FROM upstream_model", + Dialect::BigQuery, + ); + let m = find_mapping(&result.columns.mappings, "field"); + assert_eq!( + concrete_sources(m), + vec![("upstream_model".into(), "event".into())] + ); +} + +#[test] +fn bigquery_date_diff_isoweek_keeps_only_date_values() { + let result = analyze_with_dialect( + "SELECT DATE_DIFF(event_date, other_date, ISOWEEK) AS days FROM events", + Dialect::BigQuery, + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "days")), + vec![ + ("events".into(), "event_date".into()), + ("events".into(), "other_date".into()), + ] + ); +} + +#[test] +fn cte_struct_field_access_uses_cte_binding_column() { + let result = analyze_one( + "WITH upstream AS (SELECT event FROM source) SELECT upstream.event.field AS value FROM upstream", + ); + let m = find_mapping(&result.columns.mappings, "value"); + assert_eq!(concrete_sources(m), vec![("source".into(), "event".into())]); +} + +#[test] +fn derived_struct_field_access_uses_derived_binding_column() { + let result = analyze_one( + "SELECT derived.event.field AS value FROM (SELECT event FROM source) AS derived", + ); + let m = find_mapping(&result.columns.mappings, "value"); + assert_eq!(concrete_sources(m), vec![("source".into(), "event".into())]); +} + +#[test] +fn physical_relation_prefix_struct_field_access_uses_table_parts() { + let result = + analyze_one("SELECT catalog.schema.source.event.field AS value FROM catalog.schema.source"); + let m = find_mapping(&result.columns.mappings, "value"); + assert_eq!(concrete_sources(m), vec![("source".into(), "event".into())]); + assert_eq!(result.tables.inputs[0].catalog.as_deref(), Some("catalog")); + assert_eq!(result.tables.inputs[0].schema.as_deref(), Some("schema")); + match m.sources.as_slice() { + [ColumnOrigin::Concrete { table, column }] => { + assert_eq!(table.catalog.as_deref(), Some("catalog")); + assert_eq!(table.schema.as_deref(), Some("schema")); + assert_eq!(table.table, "source"); + assert_eq!(column, "event"); + } + other => panic!("expected concrete physical relation source, got {other:?}"), + } +} + +#[test] +fn quoted_single_component_relation_name_keeps_embedded_dot() { + let result = analyze_with_dialect( + "SELECT \"orders.v2\".payload.field AS value FROM \"orders.v2\"", + Dialect::PostgreSql, + ); + let m = find_mapping(&result.columns.mappings, "value"); + match m.sources.as_slice() { + [ColumnOrigin::Concrete { table, column }] => { + assert_eq!(table.catalog, None); + assert_eq!(table.schema, None); + assert_eq!(table.table, "orders.v2"); + assert_eq!(column, "payload"); + } + other => panic!("expected quoted relation source, got {other:?}"), + } +} + +#[test] +fn qualified_struct_field_access_keeps_normal_alias_column_resolution() { + let result = analyze_one("SELECT source.user_id AS value FROM source"); + let m = find_mapping(&result.columns.mappings, "value"); + assert_eq!( + concrete_sources(m), + vec![("source".into(), "user_id".into())] + ); +} + +#[test] +fn generic_date_trunc_keeps_date_part_identifiers() { + let result = + analyze_one("SELECT DATE_TRUNC(event_date, WEEK(MONDAY)) AS week_start FROM events"); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "week_start")), + vec![ + ("events".into(), "MONDAY".into()), + ("events".into(), "event_date".into()), + ] + ); +} + +#[test] +fn unqualified_struct_field_access_keeps_binding_ambiguity() { + let result = analyze_one( + "SELECT event.field AS value FROM first_source JOIN second_source ON first_source.id = second_source.id", + ); + let m = find_mapping(&result.columns.mappings, "value"); + assert!(matches!( + m.sources.as_slice(), + [ColumnOrigin::Ambiguous { column, candidates }] if column == "event" && candidates.len() == 2 + )); +} + +struct UnqualifiedStructCatalog; + +impl CatalogProvider for UnqualifiedStructCatalog { + fn list_columns(&self, table: &TableRef) -> Option> { + match table.table.as_str() { + "first_source" => Some(vec!["event".into()]), + "second_source" => Some(vec!["other".into()]), + _ => None, + } + } + + fn resolve_column(&self, column: &str, candidates: &[TableRef]) -> Option { + (column == "event") + .then(|| { + candidates + .iter() + .find(|table| table.table == "first_source") + .cloned() + }) + .flatten() + } +} + +#[test] +fn unqualified_struct_field_access_uses_catalog_owner_for_ambiguous_root() { + let result = analyze( + "SELECT event.field AS value FROM first_source JOIN second_source ON first_source.id = second_source.id", + AnalyzeOptions { + dialect: Dialect::BigQuery, + catalog: Some(Box::new(UnqualifiedStructCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("SQL should parse") + .remove(0); + let m = find_mapping(&result.columns.mappings, "value"); + match m.sources.as_slice() { + [ColumnOrigin::Concrete { table, column }] => { + assert_eq!(table.table, "first_source"); + assert_eq!(column, "event"); + } + other => panic!("expected catalog-resolved source, got {other:?}"), + } +} + +struct RowValueCatalog; + +impl CatalogProvider for RowValueCatalog { + fn list_columns(&self, table: &TableRef) -> Option> { + (table.table == "source_table").then(|| vec!["source".into()]) + } + + fn resolve_column(&self, column: &str, candidates: &[TableRef]) -> Option { + (column == "source") + .then(|| candidates.first().cloned()) + .flatten() + } +} + +#[test] +fn bigquery_row_value_alias_prefers_catalog_column_with_same_name() { + let result = analyze( + "SELECT ARRAY_AGG(source) AS event FROM source_table AS source", + AnalyzeOptions { + dialect: Dialect::BigQuery, + catalog: Some(Box::new(RowValueCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("SQL should parse") + .remove(0); + let m = find_mapping(&result.columns.mappings, "event"); + assert_eq!( + concrete_sources(m), + vec![("source_table".into(), "source".into())] + ); +} + +#[test] +fn bigquery_row_value_alias_prefers_cte_output_column_with_same_name() { + let result = analyze_with_dialect( + "WITH source AS (SELECT source_table AS source FROM base) SELECT ARRAY_AGG(source) AS event FROM source", + Dialect::BigQuery, + ); + let m = find_mapping(&result.columns.mappings, "event"); + assert_eq!( + concrete_sources(m), + vec![("base".into(), "source_table".into())] + ); +} + +#[test] +fn bigquery_date_trunc_isoyear_and_timezone_data_are_classified() { + let result = analyze_with_dialect( + "SELECT DATE_TRUNC(event_date, ISOYEAR) AS year_start, TIMESTAMP_TRUNC(event_ts, DAY, tz_name) AS ts_day FROM events", + Dialect::BigQuery, + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "year_start")), + vec![("events".into(), "event_date".into())] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "ts_day")), + vec![ + ("events".into(), "event_ts".into()), + ("events".into(), "tz_name".into()), + ] + ); +} + +#[test] +fn bigquery_row_value_alias_prefers_derived_output_column_with_same_name() { + let result = analyze_with_dialect( + "SELECT ARRAY_AGG(source) AS event FROM (SELECT source_table AS source FROM base) AS source", + Dialect::BigQuery, + ); + let m = find_mapping(&result.columns.mappings, "event"); + assert_eq!( + concrete_sources(m), + vec![("base".into(), "source_table".into())] + ); +} + +#[test] +fn bigquery_date_trunc_three_argument_form_is_not_a_timezone_signature() { + let result = analyze_with_dialect( + "SELECT DATE_TRUNC(event_date, DAY, tz_name) AS day_start FROM events", + Dialect::BigQuery, + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "day_start")), + vec![ + ("events".into(), "DAY".into()), + ("events".into(), "event_date".into()), + ("events".into(), "tz_name".into()), + ] + ); +} + +#[test] +fn bigquery_row_value_relation_alias_is_not_a_column() { + let result = analyze_with_dialect( + "SELECT ARRAY_AGG(source) AS event FROM source_table AS source", + Dialect::BigQuery, + ); + let m = find_mapping(&result.columns.mappings, "event"); + assert!(matches!( + m.sources.as_slice(), + [ColumnOrigin::Ambiguous { column, candidates }] + if column == "source" && candidates.is_empty() + )); +} + +#[test] +fn postgresql_row_value_relation_alias_is_not_a_column() { + let result = analyze_with_dialect( + "SELECT ARRAY_AGG(source) AS event FROM source_table AS source", + Dialect::PostgreSql, + ); + let m = find_mapping(&result.columns.mappings, "event"); + assert!(matches!( + m.sources.as_slice(), + [ColumnOrigin::Ambiguous { column, candidates }] + if column == "source" && candidates.is_empty() + )); +} + +#[test] +fn generic_row_value_relation_alias_preserves_existing_behavior() { + let result = analyze_one("SELECT ARRAY_AGG(source) AS event FROM source_table AS source"); + let m = find_mapping(&result.columns.mappings, "event"); + assert_eq!( + concrete_sources(m), + vec![("source_table".into(), "source".into())] + ); +} + +#[test] +fn bigquery_date_part_position_can_still_be_a_data_expression() { + let result = analyze_with_dialect( + "SELECT DATE_TRUNC(WEEK, WEEK(MONDAY)) AS week_start FROM events", + Dialect::BigQuery, + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "week_start")), + vec![("events".into(), "WEEK".into())] + ); +} + +#[test] +fn bigquery_last_day_has_optional_static_part() { + let result = analyze_with_dialect( + "SELECT LAST_DAY(event_date, MONTH) AS month_end, LAST_DAY(event_date) AS day_end FROM events", + Dialect::BigQuery, + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "month_end")), + vec![("events".into(), "event_date".into())] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "day_end")), + vec![("events".into(), "event_date".into())] + ); +} + +#[test] +fn unknown_udf_keeps_all_arguments() { + let result = analyze_with_dialect( + "SELECT my_udf(event_date, ISOYEAR) AS udf_value FROM events", + Dialect::BigQuery, + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "udf_value")), + vec![ + ("events".into(), "ISOYEAR".into()), + ("events".into(), "event_date".into()), + ] + ); +} + +#[test] +fn snowflake_date_part_signatures_skip_only_static_parts() { + let result = analyze_with_dialect( + "SELECT DATEADD(DAY, amount, event_date) AS added, DATE_PART(YEAR, event_date) AS year_value, TRUNC(event_date, dynamic_part) AS truncated FROM events", + Dialect::Snowflake, + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "added")), + vec![ + ("events".into(), "amount".into()), + ("events".into(), "event_date".into()), + ] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "year_value")), + vec![("events".into(), "event_date".into())] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "truncated")), + vec![ + ("events".into(), "dynamic_part".into()), + ("events".into(), "event_date".into()), + ] + ); +} + +#[test] +fn mysql_and_databricks_date_part_signatures_are_dialect_scoped() { + let mysql = analyze_with_dialect( + "SELECT TIMESTAMPDIFF(DAY, start_date, end_date) AS elapsed FROM events", + Dialect::MySql, + ); + assert_eq!( + concrete_sources(find_mapping(&mysql.columns.mappings, "elapsed")), + vec![ + ("events".into(), "end_date".into()), + ("events".into(), "start_date".into()), + ] + ); + + let databricks = analyze_with_dialect( + "SELECT DATEDIFF(DAY, start_date, end_date) AS elapsed FROM events", + Dialect::Databricks, + ); + assert_eq!( + concrete_sources(find_mapping(&databricks.columns.mappings, "elapsed")), + vec![ + ("events".into(), "end_date".into()), + ("events".into(), "start_date".into()), + ] + ); +} + +#[test] +fn mysql_date_part_aliases_are_limited_to_legal_timestamp_units() { + let result = analyze_with_dialect( + "SELECT TIMESTAMPDIFF(SQL_TSI_DAY, start_date, end_date) AS aliased, TIMESTAMPDIFF(DAY_SECOND, start_date, end_date) AS composite FROM events", + Dialect::MySql, + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "aliased")), + vec![ + ("events".into(), "end_date".into()), + ("events".into(), "start_date".into()), + ] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "composite")), + vec![ + ("events".into(), "DAY_SECOND".into()), + ("events".into(), "end_date".into()), + ("events".into(), "start_date".into()), + ] + ); +} + +#[test] +fn databricks_add_and_diff_have_distinct_date_part_grammars() { + let result = analyze_with_dialect( + "SELECT DATEADD(DAYOFYEAR, amount, event_ts) AS added, DATEDIFF(DAYOFYEAR, start_ts, end_ts) AS elapsed FROM events", + Dialect::Databricks, + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "added")), + vec![ + ("events".into(), "amount".into()), + ("events".into(), "event_ts".into()), + ] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "elapsed")), + vec![ + ("events".into(), "DAYOFYEAR".into()), + ("events".into(), "end_ts".into()), + ("events".into(), "start_ts".into()), + ] + ); +} + +#[test] +fn redshift_temporal_profiles_keep_date_part_functions_and_fallbacks() { + let result = analyze_with_dialect( + "SELECT DATEADD(day, amount, event_ts) AS added, DATEDIFF(week, start_ts, end_ts) AS elapsed, DATE_PART(dow, event_ts) AS weekday_value, PGDATE_PART(dow, event_ts) AS pg_weekday_value, DATE_PART(dayofyear, event_ts) AS dayofyear_unknown, DATEADD(m, amount, event_ts) AS minute_alias, DATEADD(w, amount, event_ts) AS week_alias, DATEADD(mon, amount, event_ts) AS month_alias, DATEADD(mm, amount, event_ts) AS sqlserver_month, DATEADD(wk, amount, event_ts) AS sqlserver_week, DATEADD(dynamic_part, amount, event_ts) AS dynamic_added, DATE_TRUNC('week', event_ts) AS truncated FROM events", + Dialect::Redshift, + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "added")), + vec![ + ("events".into(), "amount".into()), + ("events".into(), "event_ts".into()), + ] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "elapsed")), + vec![ + ("events".into(), "end_ts".into()), + ("events".into(), "start_ts".into()), + ] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "weekday_value")), + vec![("events".into(), "event_ts".into())] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "pg_weekday_value")), + vec![("events".into(), "event_ts".into())] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "dayofyear_unknown")), + vec![ + ("events".into(), "dayofyear".into()), + ("events".into(), "event_ts".into()), + ] + ); + for output in ["minute_alias", "week_alias", "month_alias"] { + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, output)), + vec![ + ("events".into(), "amount".into()), + ("events".into(), "event_ts".into()), + ] + ); + } + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "sqlserver_month")), + vec![ + ("events".into(), "amount".into()), + ("events".into(), "event_ts".into()), + ("events".into(), "mm".into()), + ] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "sqlserver_week")), + vec![ + ("events".into(), "amount".into()), + ("events".into(), "event_ts".into()), + ("events".into(), "wk".into()), + ] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "dynamic_added")), + vec![ + ("events".into(), "amount".into()), + ("events".into(), "dynamic_part".into()), + ("events".into(), "event_ts".into()), + ] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "truncated")), + vec![("events".into(), "event_ts".into())] + ); +} + +#[test] +fn mssql_temporal_profiles_keep_family_specific_date_parts() { + let result = analyze_with_dialect( + "SELECT DATEADD(day, amount, event_ts) AS added, DATEDIFF(weekday, start_ts, end_ts) AS elapsed, DATEDIFF_BIG(ns, start_ts, end_ts) AS big_elapsed, DATEPART(tzoffset, event_ts) AS offset_value, DATENAME(iso_week, event_ts) AS iso_name, DATETRUNC(week, event_ts) AS truncated, DATETRUNC(weekday, event_ts) AS weekday_truncated, DATE_BUCKET(day, 7, event_ts) AS bucketed, DATE_BUCKET(day, 7, event_ts, origin_ts) AS bucketed_with_origin, DATE_BUCKET(mcs, 7, event_ts) AS bucket_microsecond, DATEADD(dynamic_part, amount, event_ts) AS dynamic_added, DATEADD(unknown_part, amount, event_ts) AS unknown_added FROM events", + Dialect::MsSql, + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "added")), + vec![ + ("events".into(), "amount".into()), + ("events".into(), "event_ts".into()), + ] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "elapsed")), + vec![ + ("events".into(), "end_ts".into()), + ("events".into(), "start_ts".into()), + ] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "big_elapsed")), + vec![ + ("events".into(), "end_ts".into()), + ("events".into(), "start_ts".into()), + ] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "offset_value")), + vec![("events".into(), "event_ts".into())] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "iso_name")), + vec![("events".into(), "event_ts".into())] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "truncated")), + vec![("events".into(), "event_ts".into())] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "weekday_truncated")), + vec![ + ("events".into(), "event_ts".into()), + ("events".into(), "weekday".into()), + ] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "bucketed")), + vec![("events".into(), "event_ts".into())] + ); + assert_eq!( + concrete_sources(find_mapping( + &result.columns.mappings, + "bucketed_with_origin" + )), + vec![ + ("events".into(), "event_ts".into()), + ("events".into(), "origin_ts".into()), + ] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "bucket_microsecond")), + vec![ + ("events".into(), "event_ts".into()), + ("events".into(), "mcs".into()), + ] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "dynamic_added")), + vec![ + ("events".into(), "amount".into()), + ("events".into(), "dynamic_part".into()), + ("events".into(), "event_ts".into()), + ] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "unknown_added")), + vec![ + ("events".into(), "amount".into()), + ("events".into(), "event_ts".into()), + ("events".into(), "unknown_part".into()), + ] + ); +} + +#[test] +fn spark_temporal_functions_use_generic_value_fallback() { + let result = analyze_with_dialect( + "SELECT DATE_TRUNC('WEEK', event_ts) AS truncated, DATEDIFF(end_ts, start_ts) AS elapsed, DATEDIFF(DAY, start_ts, end_ts) AS nonstandard FROM events", + Dialect::Spark, + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "truncated")), + vec![("events".into(), "event_ts".into())] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "elapsed")), + vec![ + ("events".into(), "end_ts".into()), + ("events".into(), "start_ts".into()), + ] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "nonstandard")), + vec![ + ("events".into(), "DAY".into()), + ("events".into(), "end_ts".into()), + ("events".into(), "start_ts".into()), + ] + ); +} diff --git a/sqllineage/tests/regressions.rs b/sqllineage/tests/regressions.rs new file mode 100644 index 0000000..a42e82d --- /dev/null +++ b/sqllineage/tests/regressions.rs @@ -0,0 +1,87 @@ +use sqllineage::{AnalyzeOptions, CatalogProvider, ColumnOrigin, TableRef, analyze}; + +fn first_result(sql: &str, catalog: Option>) -> sqllineage::AnalyzeResult { + analyze( + sql, + AnalyzeOptions { + catalog, + ..AnalyzeOptions::default() + }, + ) + .expect("analysis") + .into_iter() + .next() + .expect("one statement") +} + +struct ExtCatalog; + +impl CatalogProvider for ExtCatalog { + fn list_columns(&self, table: &TableRef) -> Option> { + (table.table == "ext_a").then(|| vec!["col_x".into(), "col_y".into()]) + } + + fn resolve_column(&self, _column: &str, _candidates: &[TableRef]) -> Option { + None + } +} + +#[test] +fn named_projection_through_unknown_star_keeps_requested_column_and_uncertainty() { + for sql in [ + "WITH src AS (SELECT * FROM some_unknown_source) SELECT id FROM src", + "SELECT id FROM (SELECT * FROM some_unknown_source) src", + ] { + let result = first_result(sql, None); + assert!(result.columns.has_unresolved_stars, "SQL: {sql}"); + assert!(matches!( + result.columns.mappings[0].sources.as_slice(), + [ColumnOrigin::NamedWildcard { table, column }] + if table.table == "some_unknown_source" && column == "id" + )); + } +} + +#[test] +fn nested_join_star_is_publicly_marked_even_when_output_is_from_base() { + let result = first_result( + "SELECT id FROM some_table JOIN (SELECT * FROM some_unknown_source) src ON 1=1", + None, + ); + assert!(result.columns.has_unresolved_stars); + assert!(matches!( + result.columns.mappings[0].sources.as_slice(), + [ColumnOrigin::Concrete { table, column }] + if table.table == "some_table" && column == "id" + )); +} + +#[test] +fn leading_unknown_star_does_not_publish_nonleading_only_set_names() { + let result = first_result( + "SELECT * FROM unknown_source UNION ALL SELECT id, amt AS total FROM known_table UNION ALL SELECT id, fee FROM third_table", + None, + ); + let names = result + .columns + .mappings + .iter() + .map(|mapping| mapping.target.column.as_str()) + .collect::>(); + assert_eq!(names, vec!["*", "id", "total"]); +} + +#[test] +fn source_free_set_branch_is_retained_as_incomplete_lineage() { + let result = first_result( + "WITH a AS (SELECT * FROM ext_a), u AS (SELECT 1 AS c1, 2 AS c2 UNION ALL SELECT a.col_x, a.col_y FROM a) SELECT c1 FROM u", + Some(Box::new(ExtCatalog)), + ); + assert!(matches!( + result.columns.mappings[0].sources.as_slice(), + [ + ColumnOrigin::Concrete { table, column }, + ColumnOrigin::SourceFree { column: marker } + ] if marker == "c1" && table.table == "ext_a" && column == "col_x" + )); +}