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/src/build/expr.rs b/sqllineage/src/build/expr.rs index 2b2cb3e..6dcbabf 100644 --- a/sqllineage/src/build/expr.rs +++ b/sqllineage/src/build/expr.rs @@ -1,30 +1,49 @@ -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::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 +68,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 +86,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 +102,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 +124,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 +167,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,7 +214,9 @@ 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) => { @@ -192,7 +250,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 +292,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 +316,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/mod.rs b/sqllineage/src/build/mod.rs index 605b29c..b0fd16b 100644 --- a/sqllineage/src/build/mod.rs +++ b/sqllineage/src/build/mod.rs @@ -5,7 +5,7 @@ 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 +15,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 +29,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..69e0f0e 100644 --- a/sqllineage/src/build/select.rs +++ b/sqllineage/src/build/select.rs @@ -1,10 +1,10 @@ use sqlparser::ast::{ - Expr, Ident, Select, SelectItem, SelectItemQualifiedWildcardKind, TableFactor, TableWithJoins, + Expr, Select, SelectItem, SelectItemQualifiedWildcardKind, TableFactor, TableWithJoins, }; use crate::build::LineageBuilder; use crate::build::expr::determine_edge_kind; -use crate::graph::scope::{Binding, ScopeColumn, ScopeKind}; +use crate::graph::scope::{Binding, ScopeColumn, ScopeKind, VirtualColumn, VirtualColumnState}; impl LineageBuilder { /// Process a SELECT — FROM first, then projection. @@ -24,7 +24,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 +40,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 +57,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()); } @@ -167,9 +167,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 +253,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..587ce93 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 { .. } diff --git a/sqllineage/src/graph/mod.rs b/sqllineage/src/graph/mod.rs index 6cc33cd..bedf2ee 100644 --- a/sqllineage/src/graph/mod.rs +++ b/sqllineage/src/graph/mod.rs @@ -31,20 +31,52 @@ 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 { diff --git a/sqllineage/src/graph/node.rs b/sqllineage/src/graph/node.rs index 3227495..6061105 100644 --- a/sqllineage/src/graph/node.rs +++ b/sqllineage/src/graph/node.rs @@ -1,4 +1,5 @@ -use crate::graph::scope::ScopeId; +use crate::graph::edge::EdgeKind; +use crate::graph::scope::{Binding, ScopeId}; use crate::types::TableRef; pub(crate) type NodeId = usize; @@ -6,12 +7,23 @@ pub(crate) type NodeId = usize; #[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 { @@ -19,5 +31,21 @@ pub(crate) enum RawNode { 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..c617df8 100644 --- a/sqllineage/src/resolve/mod.rs +++ b/sqllineage/src/resolve/mod.rs @@ -1,34 +1,38 @@ 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::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(), warnings, - }; + }); } if topo::topological_sort(&graph.nodes, &graph.edges).is_err() { @@ -36,14 +40,16 @@ pub(crate) fn resolve( kind: WarningKind::UnexpectedCycle, location: None, }); - return AnalyzeResult { + return Ok(AnalyzeResult { statement_type, tables: graph.tables, columns: ColumnLineage::default(), warnings, - }; + }); } + validate_set_arities(&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,71 +58,147 @@ 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 }, + warnings, + }) +} + +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 { + table, + scope: star_scope, + } = &graph.nodes[col.node_id] + { + match star_arity(table.as_ref(), *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; } } - RawNode::Star { table, scope } => { - expand_star( - table.as_ref(), - *scope, - &graph, - &mut resolved, - &incoming, - output_table.as_ref(), - &mut mappings, - &mut HashSet::new(), - ); + 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, } - _ => {} } - } - - 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)); + }; + active.remove(&scope); + Ok(result) +} - if let Some(cat) = catalog { - catalog::apply_catalog(&mut mappings, cat); +fn star_arity( + table: Option<&TableRef>, + scope: usize, + graph: &RawGraph, + catalog: Option<&dyn CatalogProvider>, + active: &mut HashSet, +) -> Result, ParseError> { + if let Some(table) = table { + let binding = graph.scopes.lookup(scope, &table.table).cloned(); + return match binding { + Some(Binding::Cte(child) | Binding::DerivedTable(child)) => { + scope_arity(child, graph, catalog, active) + } + Some(Binding::Table(actual_table)) => Ok(catalog + .and_then(|catalog| catalog.list_columns(&actual_table)) + .map(|columns| columns.len())), + _ => Ok(catalog + .and_then(|catalog| catalog.list_columns(table)) + .map(|columns| columns.len())), + }; } - AnalyzeResult { - statement_type, - tables: graph.tables, - columns: ColumnLineage { mappings }, - warnings, + let bindings = effective_bindings(scope, graph); + let mut width = 0; + for (_, binding) in bindings { + let binding_width = match binding { + Binding::Table(table) => catalog + .and_then(|catalog| catalog.list_columns(&table)) + .map(|columns| columns.len()), + Binding::Cte(child) | Binding::DerivedTable(child) => { + scope_arity(child, graph, catalog, active)? + } + Binding::VirtualSource(source) => { + Some(graph.scopes.virtual_source(source).columns.len()) + } + }; + match binding_width { + Some(binding_width) => width += binding_width, + None => return Ok(None), + } } + for &child in graph.scopes.anonymous_derived(scope) { + match scope_arity(child, graph, catalog, active)? { + Some(child_width) => width += child_width, + None => return Ok(None), + } + } + Ok(Some(width)) } fn effective_bindings(scope: usize, graph: &RawGraph) -> Vec<(String, Binding)> { @@ -141,6 +223,371 @@ fn wildcard_mapping(output_table: Option<&TableRef>, source_table: TableRef) -> } } +#[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 { table, scope } => expand_star( + table.as_ref(), + *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() { + let mut sources = left_mapping.sources; + let mut transform = left_mapping.transform.clone(); + if let Some(right_mapping) = right_mappings.get(idx) { + sources.extend(right_mapping.sources.clone()); + transform = merge_transform(&transform, &right_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| { + mapping.sources.iter().any(|source| match source { + ColumnOrigin::Wildcard { .. } => true, + ColumnOrigin::Recursive { base_sources } => base_sources + .iter() + .any(|source| matches!(source, ColumnOrigin::Wildcard { .. })), + _ => 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 mut sources = left_mapping.sources.clone(); + sources.extend(right_mapping.sources.clone()); + merged.push(ColumnMapping { + target: left_mapping.target.clone(), + sources, + transform: merge_transform(&left_mapping.transform, &right_mapping.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 +} + +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::Recursive { base_sources } => sources.extend( + base_sources + .iter() + .filter(|source| matches!(source, ColumnOrigin::Wildcard { .. })) + .cloned(), + ), + _ => {} + } + } + } + sources +} + +fn append_unknown_sources( + mut mapping: ColumnMapping, + unknown_sources: &[ColumnOrigin], +) -> ColumnMapping { + 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( @@ -149,56 +596,133 @@ fn expand_star( graph: &RawGraph, resolved: &mut Vec>, incoming: &[Vec], - output_table: Option<&TableRef>, + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, 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); + expand_scope_columns( + s, + graph, + resolved, + incoming, + catalog, + mapping_cache, + mappings, + visited_scopes, + ); + } else if let Some(Binding::Table(actual_table)) = binding { + mappings.push(wildcard_mapping(None, actual_table)); + } else if let Some(Binding::VirtualSource(source)) = binding { + expand_virtual_source( + source, + graph, + resolved, + incoming, + catalog, + mapping_cache, + mappings, + ); } else { - mappings.push(wildcard_mapping(output_table, t.clone())); + mappings.push(wildcard_mapping(None, t.clone())); } } else { for (_, binding) in effective_bindings(scope, graph) { match binding { - Binding::Table(tref) => mappings.push(wildcard_mapping(output_table, tref)), + Binding::Table(tref) => mappings.push(wildcard_mapping(None, tref)), Binding::Cte(s) | Binding::DerivedTable(s) => { - expand_scope_columns(s, graph, resolved, incoming, output_table, mappings, visited_scopes); + expand_scope_columns( + s, + graph, + resolved, + incoming, + catalog, + mapping_cache, + mappings, + visited_scopes, + ); } + Binding::VirtualSource(source) => expand_virtual_source( + source, + graph, + resolved, + incoming, + catalog, + mapping_cache, + mappings, + ), } } for &child in graph.scopes.anonymous_derived(scope) { - expand_scope_columns(child, graph, resolved, incoming, output_table, mappings, visited_scopes); + expand_scope_columns( + child, + graph, + resolved, + incoming, + catalog, + mapping_cache, + mappings, + visited_scopes, + ); } } } /// 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], - output_table: Option<&TableRef>, + 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 { table, scope } = &graph.nodes[col.node_id] { - expand_star(table.as_ref(), *scope, graph, resolved, incoming, output_table, mappings, visited_scopes); + expand_star( + table.as_ref(), + *scope, + graph, + resolved, + incoming, + catalog, + mapping_cache, + 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); + 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: output_table.cloned(), + table: None, column: col.name.clone(), }, sources, @@ -208,20 +732,53 @@ fn expand_scope_columns( } } +/// 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, -) -> (Vec, Vec, bool) { + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, +) -> (Vec, Vec, bool, TransformKind) { if !visited.insert(node_id) { - return (vec![], vec![], false); + 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]; @@ -229,16 +786,24 @@ fn collect_output_sources( has_back = true; continue; } - let (sub_sources, sub_back) = - collect_leaf_origins(edge.from, graph, resolved, incoming, visited); + 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) + (sources, kinds, has_back, inherited_transform) } fn collect_leaf_origins( @@ -247,24 +812,306 @@ fn collect_leaf_origins( 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); + 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) = - collect_output_sources(node_id, graph, resolved, incoming, visited); - (sources, has_back) + 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); + let origin = resolve_node( + node_id, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ); match origin { - Some(o) => (vec![o], false), - None => (vec![], false), + 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())); + } + // A catalog-less star has no individual named mapping. Returning its + // wildcard origin is safer than falling through to a fabricated concrete + // source for a named reference through the CTE/derived scope. + let wildcard_sources = mappings + .iter() + .flat_map(|mapping| mapping.sources.iter()) + .filter_map(|source| match source { + ColumnOrigin::Wildcard { .. } => Some(source.clone()), + ColumnOrigin::Recursive { base_sources } => base_sources + .iter() + .find(|source| matches!(source, ColumnOrigin::Wildcard { .. })) + .cloned(), + _ => None, + }) + .collect::>(); + (!wildcard_sources.is_empty()).then_some((wildcard_sources, false, TransformKind::Direct)) +} + +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()), } } + (flattened, has_back) } fn resolve_node( @@ -273,6 +1120,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,31 +1132,95 @@ 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::RowValueCandidate { + name, + scope, + binding, + } => resolve_row_value_candidate( + name, + *scope, + binding.clone(), + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ), + RawNode::Star { table, .. } => table .as_ref() .map(|t| ColumnOrigin::Wildcard { table: t.clone() }), @@ -319,18 +1232,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 +1355,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 +1386,7 @@ fn find_single_binding(scope: usize, graph: &RawGraph) -> Option { } } +#[allow(clippy::too_many_arguments)] fn resolve_unqualified( name: &str, scope: usize, @@ -373,10 +1394,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 +1417,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 +1427,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 +1485,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 +1518,63 @@ 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() + }; + } + if let Some(source) = mappings + .iter() + .flat_map(|mapping| mapping.sources.iter()) + .find_map(|source| match source { + ColumnOrigin::Wildcard { .. } => Some(source.clone()), + ColumnOrigin::Recursive { base_sources } => base_sources + .iter() + .find(|source| matches!(source, ColumnOrigin::Wildcard { .. })) + .cloned(), + _ => None, + }) + { + 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 +1588,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 +1615,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..2504ee4 100644 --- a/sqllineage/src/types.rs +++ b/sqllineage/src/types.rs @@ -122,7 +122,10 @@ pub struct ColumnMapping { 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, @@ -207,7 +210,20 @@ pub enum Dialect { BigQuery, } -/// Error returned when SQL parsing fails. +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 or semantic validation fails. #[derive(Debug, Clone)] pub struct ParseError { pub message: String, @@ -225,6 +241,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..c7b12ad 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,30 @@ 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 + } +} + fn opts_with_catalog() -> AnalyzeOptions { AnalyzeOptions { catalog: Some(Box::new(MockCatalog)), @@ -66,6 +90,68 @@ fn select_star_with_catalog_expands() { ); } +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 +196,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 +237,268 @@ 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_eq!( + concrete_sources(&result.columns.mappings[0]), + vec![("ext_a".into(), "col_x".into())] + ); +} + +#[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::Wildcard { table }] if table.table == "unknown" + )); +} + +#[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/expr_coverage.rs b/sqllineage/tests/expr_coverage.rs index 181e796..3a97f47 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,358 @@ 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 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 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 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_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_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())] + ); +}