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/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..297bfe8 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,12 +53,11 @@ 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(); @@ -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); } @@ -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,9 +444,9 @@ 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(), }) } } 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..5388a4a 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"); @@ -80,6 +93,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";