diff --git a/sqllineage/Cargo.toml b/sqllineage/Cargo.toml index 30a0df5..328bb07 100644 --- a/sqllineage/Cargo.toml +++ b/sqllineage/Cargo.toml @@ -16,7 +16,7 @@ name = "sqllineage" path = "src/bin/sqllineage.rs" [dependencies] -sqlparser = "0.62" +sqlparser = "0.63" serde = { version = "1", features = ["derive"] } serde_json = "1" diff --git a/sqllineage/src/build/expr.rs b/sqllineage/src/build/expr.rs index 34e4d4b..80648ce 100644 --- a/sqllineage/src/build/expr.rs +++ b/sqllineage/src/build/expr.rs @@ -41,6 +41,7 @@ impl LineageBuilder { | Expr::IsUnknown(expr) | Expr::IsNotUnknown(expr) | Expr::IsNormalized { expr, .. } + | Expr::IsJson { expr, .. } | Expr::Collate { expr, .. } | Expr::Convert { expr, .. } | Expr::Ceil { expr, .. } diff --git a/sqllineage/src/build/select.rs b/sqllineage/src/build/select.rs index 803400f..880fba5 100644 --- a/sqllineage/src/build/select.rs +++ b/sqllineage/src/build/select.rs @@ -174,6 +174,7 @@ impl LineageBuilder { | TableFactor::OpenJsonTable { .. } | TableFactor::Pivot { .. } | TableFactor::Unpivot { .. } + | TableFactor::UnpivotExpr { .. } | TableFactor::MatchRecognize { .. } | TableFactor::XmlTable { .. } | TableFactor::SemanticView { .. } => {} diff --git a/sqllineage/src/build/statement.rs b/sqllineage/src/build/statement.rs index e33de12..ac0d51c 100644 --- a/sqllineage/src/build/statement.rs +++ b/sqllineage/src/build/statement.rs @@ -96,26 +96,33 @@ impl LineageBuilder { for clause in &merge.clauses { match &clause.action { - MergeAction::Update(upd) => { - for assignment in &upd.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()); - for &anc in &ancestors { - self.graph.add_edge(anc, output, kind.clone()); + MergeAction::Update(upd) => match &upd.kind { + ast::MergeUpdateKind::Set(assignments) => { + 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()); + for &anc in &ancestors { + self.graph.add_edge(anc, output, kind.clone()); + } + self.graph.scopes.add_output_column( + self.current_scope, + ScopeColumn { + name: col_name, + node_id: output, + }, + ); } - self.graph.scopes.add_output_column( - self.current_scope, - ScopeColumn { - name: col_name, - node_id: output, - }, - ); } - } - MergeAction::Insert(ins) => { - if let ast::MergeInsertKind::Values(values) = &ins.kind { + ast::MergeUpdateKind::Wildcard => self.add_merge_source_star(), + }, + MergeAction::Insert(ins) => match &ins.kind { + ast::MergeInsertKind::Wildcard => self.add_merge_source_star(), + // `INSERT ROW` also means "every source column", but it + // predates this handling and is left as it was. + ast::MergeInsertKind::Row => {} + ast::MergeInsertKind::Values(values) => { let col_names: Vec = ins .columns .iter() @@ -148,8 +155,8 @@ impl LineageBuilder { } } } - } - MergeAction::Delete { .. } => {} + }, + MergeAction::Delete { .. } | MergeAction::DoNothing { .. } => {} } } StatementType::Merge @@ -168,6 +175,7 @@ impl LineageBuilder { | Statement::AlterSchema { .. } | Statement::AlterSession { .. } | Statement::AlterTable { .. } + | Statement::AlterTextSearch(_) | Statement::AlterType { .. } | Statement::AlterUser { .. } | Statement::AlterView { .. } @@ -188,6 +196,7 @@ impl LineageBuilder { | Statement::CreateDatabase { .. } | Statement::CreateDomain { .. } | Statement::CreateExtension { .. } + | Statement::CreateFileFormat { .. } | Statement::CreateFunction { .. } | Statement::CreateIndex(_) | Statement::CreateMacro { .. } @@ -202,11 +211,13 @@ impl LineageBuilder { | Statement::CreateSequence { .. } | Statement::CreateServer { .. } | Statement::CreateStage { .. } + | Statement::CreateTextSearch(_) | Statement::CreateTrigger { .. } | Statement::CreateType { .. } | Statement::CreateUser { .. } | Statement::CreateView { .. } | Statement::CreateVirtualTable { .. } + | Statement::CreateWarehouse { .. } | Statement::Deallocate { .. } | Statement::Declare { .. } | Statement::Deny { .. } @@ -248,6 +259,7 @@ impl LineageBuilder { | Statement::Pragma { .. } | Statement::Prepare { .. } | Statement::Print(_) + | Statement::Put { .. } | Statement::RaisError { .. } | Statement::Raise { .. } | Statement::ReleaseSavepoint { .. } @@ -288,6 +300,23 @@ impl LineageBuilder { } } + /// Record `MERGE ... UPDATE SET *` / `INSERT *`, which copy every column of + /// the source row into the target. + /// + /// The source is already bound in the current scope, so this is the same + /// `Star` node `SELECT *` produces — unexpanded without a catalog, and + /// expanded against the source's columns with one. + fn add_merge_source_star(&mut self) { + let star = self.graph.add_star(None, self.current_scope); + self.graph.scopes.add_output_column( + self.current_scope, + ScopeColumn { + name: "*".to_string(), + node_id: star, + }, + ); + } + pub(crate) fn scan_expr_for_tables(&mut self, expr: &ast::Expr) { match expr { ast::Expr::Subquery(query) => { diff --git a/sqllineage/tests/merge.rs b/sqllineage/tests/merge.rs index 9d0ff8a..7184dfd 100644 --- a/sqllineage/tests/merge.rs +++ b/sqllineage/tests/merge.rs @@ -1,7 +1,7 @@ mod common; use common::{analyze_one, concrete_sources, find_mapping, table}; -use sqllineage::TransformKind; +use sqllineage::{AnalyzeOptions, CatalogProvider, ColumnOrigin, TableRef, TransformKind, analyze}; #[test] fn merge_when_matched_update_set() { @@ -66,3 +66,62 @@ fn merge_both_clauses() { .collect(); assert!(!id_mappings.is_empty(), "expected id from INSERT clause"); } + +/// `UPDATE SET *` copies every source column; without a catalog the shape stays +/// an unexpanded star rather than a guess at the column list. +#[test] +fn merge_update_set_wildcard_is_an_unexpanded_star() { + let sql = "\ + MERGE INTO target t \ + USING source s ON t.id = s.id \ + WHEN MATCHED THEN UPDATE SET *"; + let result = analyze_one(sql); + assert_eq!(result.tables.output, Some(table("target"))); + assert_eq!(result.tables.inputs, vec![table("source")]); + + let m = find_mapping(&result.columns.mappings, "*"); + assert!( + matches!(&m.sources[..], [ColumnOrigin::Wildcard { table }] if table.table == "source"), + "expected a wildcard on source, got {:?}", + m.sources + ); +} + +struct SourceCatalog; + +impl CatalogProvider for SourceCatalog { + fn list_columns(&self, table: &TableRef) -> Option> { + (table.table == "source").then(|| vec!["id".into(), "val".into()]) + } + + fn resolve_column(&self, _column: &str, _candidates: &[TableRef]) -> Option { + None + } +} + +/// `INSERT *` resolves to the source's real columns once a catalog can name them. +#[test] +fn merge_insert_wildcard_expands_from_catalog() { + let sql = "\ + MERGE INTO target t \ + USING source s ON t.id = s.id \ + WHEN NOT MATCHED THEN INSERT *"; + let result = analyze( + sql, + AnalyzeOptions { + catalog: Some(Box::new(SourceCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("MERGE should parse") + .remove(0); + + let m_id = find_mapping(&result.columns.mappings, "id"); + assert_eq!(concrete_sources(m_id), vec![("source".into(), "id".into())]); + + let m_val = find_mapping(&result.columns.mappings, "val"); + assert_eq!( + concrete_sources(m_val), + vec![("source".into(), "val".into())] + ); +}