diff --git a/sqllineage/src/build/expr.rs b/sqllineage/src/build/expr.rs index 80648ce..8b92be8 100644 --- a/sqllineage/src/build/expr.rs +++ b/sqllineage/src/build/expr.rs @@ -303,7 +303,11 @@ impl LineageBuilder { } } -pub(crate) fn determine_edge_kind(expr: &Expr) -> EdgeKind { +/// Classify what an expression does to the values it reads. +/// +/// The answer is stamped on the output column it defines and on every edge it +/// draws to an ancestor, so it survives an expression that reads no column. +pub(crate) fn classify_expr(expr: &Expr) -> EdgeKind { match expr { Expr::Identifier(_) | Expr::CompoundIdentifier(_) | Expr::Value(_) => EdgeKind::Direct, Expr::Function(f) => { @@ -317,7 +321,7 @@ pub(crate) fn determine_edge_kind(expr: &Expr) -> EdgeKind { } Expr::Case { .. } => EdgeKind::ViaConditional, Expr::Cast { expr, .. } | Expr::Nested(expr) | Expr::Collate { expr, .. } => { - determine_edge_kind(expr) + classify_expr(expr) } _ => EdgeKind::ViaExpression, } diff --git a/sqllineage/src/build/select.rs b/sqllineage/src/build/select.rs index 880fba5..3fb58b5 100644 --- a/sqllineage/src/build/select.rs +++ b/sqllineage/src/build/select.rs @@ -3,7 +3,7 @@ use sqlparser::ast::{ }; use crate::build::LineageBuilder; -use crate::build::expr::determine_edge_kind; +use crate::build::expr::classify_expr; use crate::graph::scope::{Binding, ScopeColumn, ScopeKind}; impl LineageBuilder { @@ -22,9 +22,9 @@ impl LineageBuilder { match item { SelectItem::UnnamedExpr(expr) => { let ancestors = self.collect_ancestors(expr); - let kind = determine_edge_kind(expr); + let kind = classify_expr(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()); } @@ -38,9 +38,9 @@ impl LineageBuilder { } SelectItem::ExprWithAlias { expr, alias } => { let ancestors = self.collect_ancestors(expr); - let kind = determine_edge_kind(expr); + let kind = classify_expr(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()); } @@ -54,10 +54,10 @@ impl LineageBuilder { } SelectItem::ExprWithAliases { expr, aliases } => { let ancestors = self.collect_ancestors(expr); - let kind = determine_edge_kind(expr); + let kind = classify_expr(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()); } diff --git a/sqllineage/src/build/statement.rs b/sqllineage/src/build/statement.rs index ac0d51c..aeea223 100644 --- a/sqllineage/src/build/statement.rs +++ b/sqllineage/src/build/statement.rs @@ -1,7 +1,7 @@ use sqlparser::ast::{self, AssignmentTarget, FunctionArguments, Ident, MergeAction, Statement}; use crate::build::LineageBuilder; -use crate::build::expr::determine_edge_kind; +use crate::build::expr::classify_expr; use crate::graph::scope::ScopeColumn; use crate::types::{StatementType, TableRef}; @@ -49,8 +49,8 @@ impl LineageBuilder { for assignment in &update.assignments { 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 kind = classify_expr(&assignment.value); + let output = self.graph.add_output(col_name.clone(), kind.clone()); for &anc in &ancestors { self.graph.add_edge(anc, output, kind.clone()); } @@ -101,8 +101,9 @@ impl LineageBuilder { for assignment in assignments { 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 kind = classify_expr(&assignment.value); + let output = + self.graph.add_output(col_name.clone(), kind.clone()); for &anc in &ancestors { self.graph.add_edge(anc, output, kind.clone()); } @@ -140,8 +141,9 @@ impl LineageBuilder { .cloned() .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 kind = classify_expr(expr); + let output = + self.graph.add_output(col_name.clone(), kind.clone()); for &anc in &ancestors { self.graph.add_edge(anc, output, kind.clone()); } diff --git a/sqllineage/src/graph/mod.rs b/sqllineage/src/graph/mod.rs index 6cc33cd..002cbbc 100644 --- a/sqllineage/src/graph/mod.rs +++ b/sqllineage/src/graph/mod.rs @@ -31,8 +31,8 @@ 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, kind: EdgeKind) -> NodeId { + self.add_node(RawNode::Output { name, kind }) } pub fn add_ref(&mut self, name: String, qualifier: Option, scope: ScopeId) -> NodeId { diff --git a/sqllineage/src/graph/node.rs b/sqllineage/src/graph/node.rs index 3227495..284f6cb 100644 --- a/sqllineage/src/graph/node.rs +++ b/sqllineage/src/graph/node.rs @@ -1,3 +1,4 @@ +use crate::graph::edge::EdgeKind; use crate::graph::scope::ScopeId; use crate::types::TableRef; @@ -6,7 +7,14 @@ 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, + /// What the defining expression does, independent of what it draws + /// from. An expression that reaches no column still has one — + /// `COUNT(*)` has no column ancestor but is still an aggregate — so + /// this is where the classification survives when no edge carries it. + kind: EdgeKind, + }, /// Named reference — alias, CTE reference, derived table column. Ref { name: String, diff --git a/sqllineage/src/resolve/mod.rs b/sqllineage/src/resolve/mod.rs index 9e2a121..038e2dd 100644 --- a/sqllineage/src/resolve/mod.rs +++ b/sqllineage/src/resolve/mod.rs @@ -59,11 +59,11 @@ pub(crate) fn resolve( for col in &ordered_cols { let node_id = col.node_id; match &graph.nodes[node_id] { - RawNode::Output { name, .. } => { + RawNode::Output { name, kind } => { 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); + let transform = derive_transform(kind, &edge_kinds); if has_back { mappings.push(ColumnMapping { @@ -209,8 +209,8 @@ fn expand_scope_columns( return; } for col in graph.scopes.output_columns(scope_id) { - if let RawNode::Star { table, scope } = &graph.nodes[col.node_id] { - expand_star( + match &graph.nodes[col.node_id] { + RawNode::Star { table, scope } => expand_star( table.as_ref(), *scope, graph, @@ -219,20 +219,22 @@ fn expand_scope_columns( output_table, mappings, visited_scopes, - ); - } else { - let mut visited = HashSet::new(); - let (sources, edge_kinds, _) = - collect_output_sources(col.node_id, graph, resolved, incoming, &mut visited); - let transform = derive_transform(&edge_kinds); - mappings.push(ColumnMapping { - target: ColumnRef { - table: output_table.cloned(), - column: col.name.clone(), - }, - sources, - transform, - }); + ), + RawNode::Output { kind, .. } => { + let mut visited = HashSet::new(); + let (sources, edge_kinds, _) = + collect_output_sources(col.node_id, graph, resolved, incoming, &mut visited); + let transform = derive_transform(kind, &edge_kinds); + mappings.push(ColumnMapping { + target: ColumnRef { + table: output_table.cloned(), + column: col.name.clone(), + }, + sources, + transform, + }); + } + _ => {} } } } @@ -511,12 +513,22 @@ fn resolve_through_scope( } } -fn derive_transform(kinds: &[EdgeKind]) -> TransformKind { - if kinds.iter().any(|k| matches!(k, EdgeKind::ViaAggregation)) { +/// Classify a column from its own kind together with the kinds of the edges +/// that reached a source. +/// +/// The two normally say the same thing — a projection stamps one kind on +/// itself and on every edge it draws. They differ at a set operation, where +/// the edges redirected from the other branch carry that branch's kind, and +/// wherever a branch reached no source at all and so left no edge behind. +/// Both have to count, so this is a union and not a preference. +fn derive_transform(own_kind: &EdgeKind, edge_kinds: &[EdgeKind]) -> TransformKind { + let has = |probe: fn(&EdgeKind) -> bool| probe(own_kind) || edge_kinds.iter().any(probe); + + if has(|k| matches!(k, EdgeKind::ViaAggregation)) { TransformKind::Aggregation - } else if kinds.iter().any(|k| matches!(k, EdgeKind::ViaConditional)) { + } else if has(|k| matches!(k, EdgeKind::ViaConditional)) { TransformKind::Conditional - } else if kinds.iter().any(|k| matches!(k, EdgeKind::ViaExpression)) { + } else if has(|k| matches!(k, EdgeKind::ViaExpression)) { TransformKind::Expression } else { TransformKind::Direct diff --git a/sqllineage/tests/column_lineage.rs b/sqllineage/tests/column_lineage.rs index 4062579..ba3573e 100644 --- a/sqllineage/tests/column_lineage.rs +++ b/sqllineage/tests/column_lineage.rs @@ -68,6 +68,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"); @@ -153,3 +162,56 @@ fn proven_and_unresolved_sources_coexist_in_one_mapping() { m.sources ); } + +/// A source-free projection is classified by what the expression is, not by +/// the fact that it has no ancestors. A literal really is a direct value; a +/// function call or an operator is not. +#[test] +fn source_free_projections_are_classified_by_their_own_kind() { + for (sql, expected) in [ + ("SELECT 1 AS c FROM t", TransformKind::Direct), + ("SELECT NULL AS c FROM t", TransformKind::Direct), + ("SELECT CAST(1 AS INT) AS c FROM t", TransformKind::Direct), + ("SELECT 1 + 2 AS c FROM t", TransformKind::Expression), + ("SELECT NOW() AS c FROM t", TransformKind::Expression), + ("SELECT COUNT(1) AS c FROM t", TransformKind::Aggregation), + ( + "SELECT CASE WHEN 1 = 1 THEN 2 ELSE 3 END AS c FROM t", + TransformKind::Conditional, + ), + ] { + let result = analyze_one(sql); + let m = find_mapping(&result.columns.mappings, "c"); + assert!(m.sources.is_empty(), "{sql}: expected no sources"); + assert_eq!(m.transform, expected, "{sql}"); + } +} + +/// A set operation classifies the column from every branch, including one that +/// reached no source and so left no edge behind — the branch's own kind still +/// counts. +#[test] +fn a_set_operation_is_classified_by_every_branch() { + for (sql, expected) in [ + ( + "SELECT COUNT(*) AS c FROM t UNION ALL SELECT a FROM u", + TransformKind::Aggregation, + ), + ( + "SELECT a AS c FROM t UNION ALL SELECT SUM(b) FROM u", + TransformKind::Aggregation, + ), + ( + "SELECT a AS c FROM t UNION ALL SELECT b + 1 FROM u", + TransformKind::Expression, + ), + ( + "SELECT a AS c FROM t UNION ALL SELECT b FROM u", + TransformKind::Direct, + ), + ] { + let result = analyze_one(sql); + let m = find_mapping(&result.columns.mappings, "c"); + assert_eq!(m.transform, expected, "{sql}"); + } +} diff --git a/sqllineage/tests/cte.rs b/sqllineage/tests/cte.rs index aedfc8a..7b4656c 100644 --- a/sqllineage/tests/cte.rs +++ b/sqllineage/tests/cte.rs @@ -314,3 +314,34 @@ fn column_behind_an_unexpanded_star_is_unresolved() { ); } } + +/// Known limitation: only the last hop classifies the column. +/// +/// `collect_output_sources` records the kind of the *immediate* incoming edge +/// and drops every kind met deeper in the walk, so an aggregate below a CTE or +/// a derived table reads as `Direct`. The column's own kind does not rescue +/// this — it describes the outer projection, which really is a plain +/// reference; the aggregation is a hop further down and never travels. +/// +/// Pinned so the change is visible when that lands. +#[test] +fn only_the_last_hop_classifies_the_column() { + for sql in [ + "WITH x AS (SELECT COUNT(*) AS c FROM t) SELECT c FROM x", + "SELECT c FROM (SELECT COUNT(*) AS c FROM t) d", + // Not a question of missing sources: `t.x` survives the hop, the + // aggregation does not. + "SELECT c FROM (SELECT SUM(x) AS c FROM t) d", + ] { + let result = analyze_one(sql); + let m = find_mapping(&result.columns.mappings, "c"); + assert_eq!(m.transform, TransformKind::Direct, "{sql}"); + } + + // Expanding a star takes no second hop — it classifies each inner column + // from that column's own node — so the same query disagrees with itself + // depending on how the column is selected. + let result = analyze_one("SELECT * FROM (SELECT COUNT(*) AS c FROM t) d"); + let m = find_mapping(&result.columns.mappings, "c"); + assert_eq!(m.transform, TransformKind::Aggregation); +}