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/select.rs b/sqllineage/src/build/select.rs index 803400f..4ed3489 100644 --- a/sqllineage/src/build/select.rs +++ b/sqllineage/src/build/select.rs @@ -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()); } diff --git a/sqllineage/src/build/statement.rs b/sqllineage/src/build/statement.rs index 684f052..0c2980d 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,7 @@ 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 +280,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..1e2134d 100644 --- a/sqllineage/src/graph/mod.rs +++ b/sqllineage/src/graph/mod.rs @@ -31,8 +31,11 @@ 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 { diff --git a/sqllineage/src/graph/node.rs b/sqllineage/src/graph/node.rs index 3227495..ef7eafa 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, + /// 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, 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..c74fa31 100644 --- a/sqllineage/src/resolve/mod.rs +++ b/sqllineage/src/resolve/mod.rs @@ -1,7 +1,7 @@ mod catalog; mod topo; -use std::collections::{HashMap, HashSet}; +use std::collections::HashSet; use crate::graph::RawGraph; use crate::graph::edge::EdgeKind; @@ -53,18 +53,17 @@ pub(crate) fn resolve( 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 { + for col in &ordered_cols { + let node_id = col.node_id; 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); + let transform = derive_transform(&graph.nodes[node_id], &edge_kinds); if has_back { mappings.push(ColumnMapping { @@ -96,17 +95,6 @@ pub(crate) fn resolve( } } - 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)); - if let Some(cat) = catalog { catalog::apply_catalog(&mut mappings, cat); } @@ -195,7 +183,7 @@ fn expand_scope_columns( 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 transform = derive_transform(&graph.nodes[col.node_id], &edge_kinds); mappings.push(ColumnMapping { target: ColumnRef { table: output_table.cloned(), @@ -397,9 +385,9 @@ fn resolve_from_bindings( } } } 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(); @@ -456,14 +444,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)) { diff --git a/sqllineage/src/types.rs b/sqllineage/src/types.rs index fbdc051..5f392b9 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, @@ -225,6 +228,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..72f0ca9 100644 --- a/sqllineage/tests/catalog.rs +++ b/sqllineage/tests/catalog.rs @@ -23,6 +23,18 @@ 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")) + } +} + fn opts_with_catalog() -> AnalyzeOptions { AnalyzeOptions { catalog: Some(Box::new(MockCatalog)), @@ -110,6 +122,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 = diff --git a/sqllineage/tests/column_lineage.rs b/sqllineage/tests/column_lineage.rs index 3418c24..eadc99a 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"); diff --git a/sqllineage/tests/cte.rs b/sqllineage/tests/cte.rs index 24bdfbd..cd2e6c1 100644 --- a/sqllineage/tests/cte.rs +++ b/sqllineage/tests/cte.rs @@ -15,6 +15,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";