From 96d89cf9f515b7acbd12558fe801cceeeaa88d96 Mon Sep 17 00:00:00 2001 From: Yang Xiufeng Date: Mon, 27 Jul 2026 14:12:19 +0800 Subject: [PATCH 1/2] feat(query): expose column ids in system columns --- src/query/storages/system/src/columns_table.rs | 15 +++++++++++++++ .../base/01_system/01_0006_system_columns.test | 10 ++++++++++ 2 files changed, 25 insertions(+) diff --git a/src/query/storages/system/src/columns_table.rs b/src/query/storages/system/src/columns_table.rs index fec4ce493373b..32c4198977272 100644 --- a/src/query/storages/system/src/columns_table.rs +++ b/src/query/storages/system/src/columns_table.rs @@ -25,7 +25,9 @@ use databend_common_expression::TableDataType; use databend_common_expression::TableField; use databend_common_expression::TableSchemaRefExt; use databend_common_expression::infer_table_schema; +use databend_common_expression::types::NumberDataType; use databend_common_expression::types::StringType; +use databend_common_expression::types::UInt64Type; use databend_common_expression::utils::FromData; use databend_common_functions::BUILTIN_FUNCTIONS; use databend_common_meta_app::schema::CatalogInfo; @@ -84,10 +86,12 @@ impl AsyncSystemTable for ColumnsTable { let mut default_exprs: Vec = Vec::with_capacity(rows.len()); let mut is_nullables: Vec = Vec::with_capacity(rows.len()); let mut comments: Vec = Vec::with_capacity(rows.len()); + let mut column_ids: Vec> = Vec::with_capacity(rows.len()); for column_info in rows.into_iter() { names.push(column_info.column.name().clone()); tables.push(column_info.table_name); databases.push(column_info.database_name); + column_ids.push(column_info.column_id); types.push(column_info.column.data_type().wrapped_display()); let data_type = column_info .column @@ -123,6 +127,7 @@ impl AsyncSystemTable for ColumnsTable { StringType::from_data(default_exprs), StringType::from_data(is_nullables), StringType::from_data(comments), + UInt64Type::from_opt_data(column_ids), ])) } } @@ -133,6 +138,7 @@ struct TableColumnInfo { column: TableField, column_comment: String, + column_id: Option, } impl ColumnsTable { @@ -149,6 +155,10 @@ impl ColumnsTable { TableField::new("default_expression", TableDataType::String), TableField::new("is_nullable", TableDataType::String), TableField::new("comment", TableDataType::String), + TableField::new( + "column_id", + TableDataType::Nullable(Box::new(TableDataType::Number(NumberDataType::UInt64))), + ), ]); let table_info = TableInfo { @@ -210,6 +220,9 @@ impl ColumnsTable { table_name: table.name().into(), column: field.clone(), column_comment: "".to_string(), + // View output columns are inferred from the query plan here; their + // ids are not stable table column ids. + column_id: None, }) } } @@ -223,6 +236,7 @@ impl ColumnsTable { table_name: table.name().into(), column: field.clone(), column_comment: "".to_string(), + column_id: Some(field.column_id() as u64), }) } } @@ -253,6 +267,7 @@ impl ColumnsTable { table_name: table.name().into(), column: field.clone(), column_comment: comment, + column_id: Some(field.column_id() as u64), }) } } diff --git a/tests/sqllogictests/suites/base/01_system/01_0006_system_columns.test b/tests/sqllogictests/suites/base/01_system/01_0006_system_columns.test index 994f059f01afd..0e786c2853ec1 100644 --- a/tests/sqllogictests/suites/base/01_system/01_0006_system_columns.test +++ b/tests/sqllogictests/suites/base/01_system/01_0006_system_columns.test @@ -43,12 +43,22 @@ columntest tid2 columntest tid3 columntest tid4 +query TI +SELECT lower(name), column_id FROM system.columns WHERE database = 'default' AND table = 't' ORDER BY name +---- +id 0 + statement ok drop view if exists default.test_v_t; statement ok create view default.test_v_t as select * from default.t; +query TB +SELECT lower(name), column_id IS NULL FROM system.columns WHERE database = 'default' AND table = 'test_v_t' ORDER BY name +---- +id true + statement ok drop table default.t; From 85fb6a9c4d339f5dbd0f2accf96b860e769ee0e8 Mon Sep 17 00:00:00 2001 From: Yang Xiufeng Date: Wed, 22 Jul 2026 21:25:01 +0800 Subject: [PATCH 2/2] feat(query): extract query lineage --- src/meta/app/src/schema/catalog.rs | 2 +- src/query/catalog/src/table.rs | 9 + .../interpreters/interpreter_table_create.rs | 6 + .../binder/bind_table_reference/bind_table.rs | 88 +- src/query/sql/src/planner/binder/ddl/view.rs | 1 + src/query/sql/src/planner/binder/insert.rs | 4 + src/query/sql/src/planner/lineage.rs | 2020 +++++++++++++++++ .../sql/src/planner/metadata/metadata.rs | 52 + src/query/sql/src/planner/mod.rs | 2 + .../sql/src/planner/optimizer/optimizer.rs | 7 + src/query/sql/src/planner/plans/ddl/view.rs | 19 +- src/query/sql/src/planner/plans/insert.rs | 7 + src/query/sql/test-support/src/lib.rs | 1 + .../sql/test-support/src/lite_context.rs | 104 +- src/query/sql/tests/it/planner.rs | 2 + src/query/sql/tests/it/planner/lineage.rs | 577 +++++ src/query/storages/stream/src/stream_table.rs | 6 + 17 files changed, 2885 insertions(+), 22 deletions(-) create mode 100644 src/query/sql/src/planner/lineage.rs create mode 100644 src/query/sql/tests/it/planner/lineage.rs diff --git a/src/meta/app/src/schema/catalog.rs b/src/meta/app/src/schema/catalog.rs index 3be7e94c3a319..9d44d56f9fbb4 100644 --- a/src/meta/app/src/schema/catalog.rs +++ b/src/meta/app/src/schema/catalog.rs @@ -26,7 +26,7 @@ use crate::schema::catalog_id_ident; use crate::storage::StorageParams; use crate::tenant::Tenant; -#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)] +#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug, Hash)] pub enum CatalogType { Default = 1, Hive = 2, diff --git a/src/query/catalog/src/table.rs b/src/query/catalog/src/table.rs index 0c86b6c641375..18b52224cf2e7 100644 --- a/src/query/catalog/src/table.rs +++ b/src/query/catalog/src/table.rs @@ -109,6 +109,15 @@ pub trait Table: Sync + Send { fn get_table_info(&self) -> &TableInfo; + /// Returns the source table whose data columns a stream exposes. + /// + /// Lineage intentionally passes through a stream to its source table. + /// Views are lineage boundaries and must not use this relation-level hook; + /// their output columns are annotated separately by the planner. + fn stream_source_table_info(&self) -> Option<&TableInfo> { + None + } + fn get_data_source_info(&self) -> DataSourceInfo { DataSourceInfo::TableSource(self.get_table_info().clone()) } diff --git a/src/query/service/src/interpreters/interpreter_table_create.rs b/src/query/service/src/interpreters/interpreter_table_create.rs index d0bb2dbdaef84..a88149cc1e9b6 100644 --- a/src/query/service/src/interpreters/interpreter_table_create.rs +++ b/src/query/service/src/interpreters/interpreter_table_create.rs @@ -251,6 +251,12 @@ impl CreateTableInterpreter { overwrite: false, source: InsertInputSource::SelectPlan(select_plan), table_info: Some(table_info), + lineage_target_table_id: None, + lineage_target_catalog_type: if self.plan.engine == Engine::Iceberg { + databend_common_meta_app::schema::CatalogType::Iceberg + } else { + databend_common_meta_app::schema::CatalogType::Default + }, }; let mut pipeline = InsertInterpreter::try_create(self.ctx.clone(), insert_plan)? diff --git a/src/query/sql/src/planner/binder/bind_table_reference/bind_table.rs b/src/query/sql/src/planner/binder/bind_table_reference/bind_table.rs index 438d1f2a0a62f..b56fa12bf2570 100644 --- a/src/query/sql/src/planner/binder/bind_table_reference/bind_table.rs +++ b/src/query/sql/src/planner/binder/bind_table_reference/bind_table.rs @@ -27,10 +27,13 @@ use databend_common_catalog::table_with_options::get_with_opt_consume; use databend_common_catalog::table_with_options::get_with_opt_max_batch_size; use databend_common_exception::ErrorCode; use databend_common_exception::Result; +use databend_common_expression::ColumnId; use databend_common_storages_basic::view_table::QUERY; use databend_storages_common_table_meta::table::get_change_type; use crate::BindContext; +use crate::LineageSourceRelation; +use crate::ViewLineageSourceColumn; use crate::binder::Binder; use crate::binder::ViewIdent; use crate::binder::util::TableIdentifier; @@ -189,17 +192,25 @@ impl Binder { { let change_type = get_change_type(&table_name_alias); if change_type.is_some() { - let table_index = self.metadata.write().add_table( - catalog, - database.clone(), - table_meta.clone(), - branch_name, - table_name_alias, - !bind_context.binding_views.is_empty(), - bind_context.planning_agg_index, - false, - cte_suffix_name, - ); + let stream_lineage_source = stream_lineage_source_relation(&table_meta); + let table_index = { + let mut metadata = self.metadata.write(); + let table_index = metadata.add_table( + catalog, + database.clone(), + table_meta.clone(), + branch_name, + table_name_alias, + !bind_context.binding_views.is_empty(), + bind_context.planning_agg_index, + false, + cte_suffix_name, + ); + if let Some(stream_lineage_source) = stream_lineage_source { + metadata.set_stream_lineage_source(table_index, stream_lineage_source); + } + table_index + }; let (s_expr, mut bind_context) = self.bind_base_table( bind_context, database.as_str(), @@ -280,11 +291,11 @@ impl Binder { new_bind_context.binding_views.insert(view_ident); if let Statement::Query(query) = &stmt { self.metadata.write().add_table( - catalog, + catalog.clone(), database.clone(), - table_meta, - branch_name, - table_name_alias, + table_meta.clone(), + branch_name.clone(), + table_name_alias.clone(), false, false, false, @@ -302,6 +313,13 @@ impl Binder { column.table_name = Some(self.normalize_identifier(table).name); } } + self.add_view_lineage_source_columns( + &new_bind_context, + catalog.as_str(), + database.as_str(), + table_name.as_str(), + &table_meta, + ); // Restore binding_views to the outer scope's value so the // current view does not leak into sibling/parent contexts. new_bind_context.binding_views = bind_context.binding_views.clone(); @@ -343,4 +361,44 @@ impl Binder { } } } + + fn add_view_lineage_source_columns( + &mut self, + bind_context: &BindContext, + catalog: &str, + database: &str, + view_name: &str, + view: &std::sync::Arc, + ) { + let relation = LineageSourceRelation { + catalog: catalog.to_string(), + database: database.to_string(), + name: view_name.to_string(), + id: view.get_table_info().ident.table_id, + }; + let mut metadata = self.metadata.write(); + for (idx, column) in bind_context.columns.iter().enumerate() { + metadata.add_view_lineage_source_column(column.index, ViewLineageSourceColumn { + relation: relation.clone(), + // View TableMeta has no persisted schema. The bound view query (including an + // explicit view column list) is the source of truth for output column names. + name: column.column_name.clone(), + id: idx as ColumnId, + }); + } + } +} + +fn stream_lineage_source_relation( + table: &std::sync::Arc, +) -> Option { + table.stream_source_table_info().and_then(|table_info| { + let database = table_info.database_name().ok()?.to_string(); + Some(LineageSourceRelation { + catalog: table_info.catalog().to_string(), + database, + name: table_info.name.clone(), + id: table_info.ident.table_id, + }) + }) } diff --git a/src/query/sql/src/planner/binder/ddl/view.rs b/src/query/sql/src/planner/binder/ddl/view.rs index 5b49183a267f5..720ceb6c1a805 100644 --- a/src/query/sql/src/planner/binder/ddl/view.rs +++ b/src/query/sql/src/planner/binder/ddl/view.rs @@ -74,6 +74,7 @@ impl Binder { view_name, column_names, subquery, + query_plan: None, }; Ok(Plan::CreateView(plan.into())) } diff --git a/src/query/sql/src/planner/binder/insert.rs b/src/query/sql/src/planner/binder/insert.rs index 655bb96a6096f..158b5e46c40c7 100644 --- a/src/query/sql/src/planner/binder/insert.rs +++ b/src/query/sql/src/planner/binder/insert.rs @@ -311,6 +311,10 @@ impl Binder { overwrite: *overwrite, source: input_source?, table_info: None, + lineage_target_table_id: (table.get_table_info().catalog_info.catalog_type() + == databend_common_meta_app::schema::CatalogType::Default) + .then_some(table.get_table_info().ident.table_id), + lineage_target_catalog_type: table.get_table_info().catalog_info.catalog_type(), }; Ok(Plan::Insert(Box::new(plan))) diff --git a/src/query/sql/src/planner/lineage.rs b/src/query/sql/src/planner/lineage.rs new file mode 100644 index 0000000000000..aff1fd8314988 --- /dev/null +++ b/src/query/sql/src/planner/lineage.rs @@ -0,0 +1,2020 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::BTreeMap; +use std::collections::BTreeSet; +use std::collections::HashMap; + +use databend_common_catalog::plan::DataSourceInfo; +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; +use databend_common_expression::ColumnId; +use databend_common_expression::FieldIndex; +use databend_common_expression::TableField; +use databend_common_meta_app::principal::StageInfo; +use databend_common_meta_app::principal::StageType; +use databend_common_meta_app::schema::CatalogType; + +use crate::BindContext; +use crate::ColumnEntry; +use crate::Metadata; +use crate::MetadataRef; +use crate::ScalarExpr; +use crate::Symbol; +use crate::optimizer::ir::SExpr; +use crate::plans::BoundColumnRef; +use crate::plans::InsertInputSource; +use crate::plans::Plan; +use crate::plans::RelOperator; +use crate::plans::ScalarItem; +use crate::plans::SubqueryExpr; +use crate::plans::Visitor; + +#[derive(Clone, Debug, Eq, PartialEq)] +struct RelationLineage { + relations: Vec, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +struct TableLineage { + target: QueryLineageRelation, + columns: Vec, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +struct ColumnLineage { + target_column: QueryLineageColumn, + source_tables: Vec, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +struct SourceTableColumns { + table: QueryLineageRelation, + columns: Vec, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct QueryLineage { + pub kind: QueryLineageKind, + /// Objects written or defined by the query. + pub targets: Vec, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum QueryLineageKind { + Ctas, + Dml, + CreateView, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct LineageTarget { + pub relation: QueryLineageRelation, + /// Objects whose data flows into this target. + pub sources: Vec, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct LineageSource { + pub relation: QueryLineageRelation, + pub columns: Vec, +} + +#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)] +pub struct QueryLineageColumnEdge { + /// The source column whose value contributes to the target column. + pub source: QueryLineageColumn, + pub target: QueryLineageColumn, +} + +#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)] +pub struct QueryLineageRelation { + pub catalog: String, + pub database: String, + pub name: String, + pub id: Option, + pub catalog_type: Option, + pub kind: QueryLineageRelationKind, +} + +#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)] +pub enum QueryLineageRelationKind { + Table, + View, + Stage, +} + +#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)] +pub struct QueryLineageColumn { + pub name: String, + pub id: ColumnId, +} + +#[derive(Clone, Debug)] +struct TargetColumnBinding { + target_relation: QueryLineageRelation, + target_column: QueryLineageColumn, + value: TargetValue, +} + +#[derive(Clone, Debug)] +enum TargetValue { + QueryOutput { output_column_index: Symbol }, + Expr { scalar: Box }, +} + +#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)] +struct SourceColumn { + table: QueryLineageRelation, + column: QueryLineageColumn, +} + +#[derive(Default)] +struct LineageResolver { + definitions: HashMap>, + active_columns: BTreeSet, +} + +#[derive(Clone, Debug)] +enum SourceExpr { + Symbol(Symbol), + Scalar(Box), + Base(SourceColumn), +} + +impl Plan { + pub fn query_lineage(&self) -> Result> { + RelationExtractor::new(self).extract_query_lineage() + } +} + +impl RelationLineage { + fn from_query_plan( + query_plan: &Plan, + target_bindings: Vec, + ) -> Result { + let (s_expr, metadata, _) = query_parts(query_plan)?; + RelationResolver::resolve(s_expr, metadata, target_bindings) + } + + #[cfg(test)] + fn from_query_outputs( + query_plan: &Plan, + target_relation: QueryLineageRelation, + target_columns: Vec, + ) -> Result { + let (_, _, bind_context) = query_parts(query_plan)?; + let mut target_bindings = Vec::with_capacity(target_columns.len()); + for (idx, target_column) in target_columns.into_iter().enumerate() { + let output = bind_context.columns.get(idx).ok_or_else(|| { + ErrorCode::Internal(format!( + "Missing query output column for target column {}", + target_column.name + )) + })?; + target_bindings.push(TargetColumnBinding { + target_relation: target_relation.clone(), + target_column, + value: TargetValue::QueryOutput { + output_column_index: output.index, + }, + }); + } + + Self::from_query_plan(query_plan, target_bindings) + } +} + +impl QueryLineage { + fn from_relation_lineage(kind: QueryLineageKind, lineage: RelationLineage) -> Self { + let targets = lineage + .relations + .into_iter() + .map(|target| { + let mut sources: BTreeMap> = + BTreeMap::new(); + for column in target.columns { + for source_table in column.source_tables { + let edges = sources.entry(source_table.table).or_default(); + edges.extend(source_table.columns.into_iter().map(|source_column| { + QueryLineageColumnEdge { + source: source_column, + target: column.target_column.clone(), + } + })); + } + } + + LineageTarget { + relation: target.target, + sources: sources + .into_iter() + .map(|(relation, mut columns)| { + columns.sort(); + columns.dedup(); + LineageSource { relation, columns } + }) + .collect(), + } + }) + .collect(); + + QueryLineage { kind, targets } + } +} + +struct RelationExtractor<'a> { + plan: &'a Plan, +} + +impl<'a> RelationExtractor<'a> { + fn new(plan: &'a Plan) -> Self { + Self { plan } + } + + fn extract_query_lineage(&self) -> Result> { + self.extract_with_kind().map(|lineage| { + lineage.map(|(kind, relation_lineage)| { + QueryLineage::from_relation_lineage(kind, relation_lineage) + }) + }) + } + + fn extract_with_kind(&self) -> Result> { + let (kind, target_bindings) = match self.plan { + Plan::CreateTable(plan) => { + let Some(select_plan) = plan.as_select.as_deref() else { + return Ok(None); + }; + let target_bindings = self.query_output_targets( + select_plan, + QueryLineageRelation { + catalog: plan.catalog.clone(), + database: plan.database.clone(), + name: plan.table.clone(), + id: None, + catalog_type: Some( + if plan.engine == databend_common_ast::ast::Engine::Iceberg { + CatalogType::Iceberg + } else { + CatalogType::Default + }, + ), + kind: QueryLineageRelationKind::Table, + }, + plan.schema.fields(), + )?; + (QueryLineageKind::Ctas, target_bindings) + } + Plan::CreateView(plan) => { + let Some(query_plan) = plan.query_plan.as_deref() else { + return Ok(None); + }; + let target_bindings = self.query_output_columns_targets( + query_plan, + QueryLineageRelation { + catalog: plan.catalog.clone(), + database: plan.database.clone(), + name: plan.view_name.clone(), + id: None, + catalog_type: Some(CatalogType::Default), + kind: QueryLineageRelationKind::View, + }, + self.view_target_columns(query_plan, &plan.column_names)?, + )?; + (QueryLineageKind::CreateView, target_bindings) + } + Plan::Insert(plan) => { + let InsertInputSource::SelectPlan(select_plan) = &plan.source else { + return Ok(None); + }; + let target_bindings = self.query_output_targets( + select_plan, + QueryLineageRelation { + catalog: plan.catalog.clone(), + database: plan.database.clone(), + name: plan.table.clone(), + id: plan.lineage_target_table_id, + catalog_type: Some(plan.lineage_target_catalog_type), + kind: QueryLineageRelationKind::Table, + }, + plan.schema.fields(), + )?; + (QueryLineageKind::Dml, target_bindings) + } + Plan::Replace(plan) => { + let InsertInputSource::SelectPlan(select_plan) = &plan.source else { + return Ok(None); + }; + let target_bindings = self.query_output_targets( + select_plan, + QueryLineageRelation { + catalog: plan.catalog.clone(), + database: plan.database.clone(), + name: plan.table.clone(), + id: Some(plan.table_id), + catalog_type: Some(CatalogType::Default), + kind: QueryLineageRelationKind::Table, + }, + plan.schema.fields(), + )?; + (QueryLineageKind::Dml, target_bindings) + } + Plan::CopyIntoTable(plan) => { + if plan.no_file_to_copy { + return Ok(None); + } + let Some(source) = stage_relation(&plan.stage_table_info.stage_info) else { + return Ok(None); + }; + return Ok(Some((QueryLineageKind::Dml, RelationLineage { + relations: vec![TableLineage { + target: QueryLineageRelation { + catalog: plan.catalog_info.catalog_name().to_string(), + database: plan.database_name.clone(), + name: plan.table_name.clone(), + id: None, + catalog_type: Some(plan.catalog_info.catalog_type()), + kind: QueryLineageRelationKind::Table, + }, + columns: vec![ColumnLineage { + target_column: QueryLineageColumn { + name: String::new(), + id: 0, + }, + source_tables: vec![SourceTableColumns { + table: source, + columns: vec![], + }], + }], + }], + }))); + } + Plan::CopyIntoLocation(plan) => { + let Some(target) = stage_relation(&plan.info.stage) else { + return Ok(None); + }; + let (_, _, bind_context) = query_parts(&plan.from)?; + let target_columns = bind_context + .columns + .iter() + .enumerate() + .map(|(idx, output)| QueryLineageColumn { + name: output.column_name.clone(), + id: idx as ColumnId, + }) + .collect(); + let target_bindings = + self.query_output_columns_targets(&plan.from, target, target_columns)?; + let lineage = RelationLineage::from_query_plan(&plan.from, target_bindings)?; + return Ok(Some((QueryLineageKind::Dml, lineage))); + } + Plan::InsertMultiTable(plan) => { + let (s_expr, metadata, bind_context) = query_parts(&plan.input_source)?; + let output_columns = bind_context + .columns + .iter() + .map(|column| column.index) + .collect::>(); + let mut bindings = Vec::new(); + for into in self.multi_insert_intos(plan) { + let relation = QueryLineageRelation { + catalog: into.catalog.clone(), + database: into.database.clone(), + name: into.table.clone(), + id: None, + catalog_type: Some(CatalogType::Default), + kind: QueryLineageRelationKind::Table, + }; + for (idx, field) in into.casted_schema.fields().iter().enumerate() { + let value = if let Some(source_exprs) = &into.source_scalar_exprs { + let scalar = source_exprs.get(idx).ok_or_else(|| { + ErrorCode::Internal(format!( + "Missing multi insert source expr for column {}", + field.name() + )) + })?; + TargetValue::Expr { + scalar: Box::new(scalar.clone()), + } + } else { + let output_column_index = + *output_columns.get(idx).ok_or_else(|| { + ErrorCode::Internal(format!( + "Missing query output column for multi insert column {}", + field.name() + )) + })?; + TargetValue::QueryOutput { + output_column_index, + } + }; + bindings.push(TargetColumnBinding { + target_relation: relation.clone(), + target_column: column_info_from_data_field(field.name(), idx), + value, + }); + } + } + let lineage = RelationResolver::resolve(s_expr, metadata, bindings)?; + return Ok(Some((QueryLineageKind::Dml, lineage))); + } + Plan::DataMutation { + s_expr, metadata, .. + } => { + let Some(mutation) = find_mutation(s_expr) else { + return Ok(None); + }; + let target_metadata = mutation.metadata.read(); + let (target_table_id, target_catalog_type) = target_metadata + .tables() + .get(mutation.target_table_index) + .map(|table| { + let table = table.table(); + let table_info = table.get_table_info(); + let catalog_type = table_info.catalog_info.catalog_type(); + let table_id = (catalog_type == CatalogType::Default) + .then_some(table_info.ident.table_id); + (table_id, catalog_type) + }) + .unwrap_or((None, CatalogType::Default)); + let relation = QueryLineageRelation { + catalog: mutation.catalog_name.clone(), + database: mutation.database_name.clone(), + name: mutation.table_name.clone(), + id: target_table_id, + catalog_type: Some(target_catalog_type), + kind: QueryLineageRelationKind::Table, + }; + let mut bindings = Vec::new(); + for matched in &mutation.matched_evaluators { + let Some(update) = &matched.update else { + continue; + }; + for (field_index, scalar) in update { + let Some(column) = target_column_from_field_index( + &target_metadata, + mutation, + *field_index, + ) else { + continue; + }; + bindings.push(TargetColumnBinding { + target_relation: relation.clone(), + target_column: column, + value: TargetValue::Expr { + scalar: Box::new(scalar.clone()), + }, + }); + } + } + drop(target_metadata); + for unmatched in &mutation.unmatched_evaluators { + for (idx, scalar) in unmatched.values.iter().enumerate() { + let Some(field) = unmatched.source_schema.field(idx).ok() else { + continue; + }; + bindings.push(TargetColumnBinding { + target_relation: relation.clone(), + target_column: column_info_from_data_field(field.name(), idx), + value: TargetValue::Expr { + scalar: Box::new(scalar.clone()), + }, + }); + } + } + let lineage = RelationResolver::resolve(s_expr, metadata, bindings)?; + return Ok(Some((QueryLineageKind::Dml, lineage))); + } + _ => return Ok(None), + }; + + RelationLineage::from_query_plan(query_plan(self.plan)?, target_bindings) + .map(|lineage| Some((kind, lineage))) + } + + fn query_output_targets( + &self, + select_plan: &'a Plan, + relation: QueryLineageRelation, + target_fields: &'a [TableField], + ) -> Result> { + let target_columns = target_fields + .iter() + .map(column_info_from_table_field) + .collect::>(); + self.query_output_columns_targets(select_plan, relation, target_columns) + } + + fn query_output_columns_targets( + &self, + select_plan: &'a Plan, + relation: QueryLineageRelation, + target_columns: Vec, + ) -> Result> { + let (_, _, bind_context) = query_parts(select_plan)?; + let mut bindings = Vec::with_capacity(target_columns.len()); + for (idx, target_column) in target_columns.into_iter().enumerate() { + let output = bind_context.columns.get(idx).ok_or_else(|| { + ErrorCode::Internal(format!( + "Missing query output column for target column {}", + target_column.name + )) + })?; + bindings.push(TargetColumnBinding { + target_relation: relation.clone(), + target_column, + value: TargetValue::QueryOutput { + output_column_index: output.index, + }, + }); + } + Ok(bindings) + } + + fn view_target_columns( + &self, + query_plan: &'a Plan, + column_names: &[String], + ) -> Result> { + let (_, _, bind_context) = query_parts(query_plan)?; + Ok(bind_context + .columns + .iter() + .enumerate() + .map(|(idx, output)| QueryLineageColumn { + name: column_names + .get(idx) + .cloned() + .unwrap_or_else(|| output.column_name.clone()), + id: idx as ColumnId, + }) + .collect()) + } + + fn multi_insert_intos( + &self, + plan: &'a crate::plans::InsertMultiTable, + ) -> Vec<&'a crate::plans::Into> { + let mut intos = Vec::new(); + for when in &plan.whens { + intos.extend(when.intos.iter()); + } + if let Some(else_clause) = &plan.opt_else { + intos.extend(else_clause.intos.iter()); + } + intos.extend(plan.intos.iter()); + intos + } +} + +fn stage_relation(stage: &StageInfo) -> Option { + if stage.is_temporary + || stage.stage_type == StageType::User + || stage.stage_name.trim().is_empty() + { + return None; + } + + Some(QueryLineageRelation { + catalog: String::new(), + database: String::new(), + name: stage.stage_name.clone(), + id: None, + catalog_type: None, + kind: QueryLineageRelationKind::Stage, + }) +} + +struct RelationResolver; + +impl RelationResolver { + fn resolve( + s_expr: &SExpr, + metadata: &MetadataRef, + targets: Vec, + ) -> Result { + let mut resolver = LineageResolver::default(); + resolver.collect_s_expr(s_expr, &metadata.read())?; + + let metadata_guard = metadata.read(); + let mut relation_columns: BTreeMap> = + BTreeMap::new(); + for target in targets { + let sources = resolver.resolve_target(&target.value, &metadata_guard)?; + relation_columns + .entry(target.target_relation) + .or_default() + .push(ColumnLineage { + target_column: target.target_column, + source_tables: group_sources_by_table(sources), + }); + } + + Ok(RelationLineage { + relations: relation_columns + .into_iter() + .map(|(target, columns)| TableLineage { target, columns }) + .collect(), + }) + } +} + +impl LineageResolver { + fn collect_s_expr(&mut self, s_expr: &SExpr, metadata: &Metadata) -> Result<()> { + self.collect_base_columns(metadata, s_expr.plan())?; + + match s_expr.plan() { + RelOperator::EvalScalar(eval_scalar) => { + self.collect_scalar_items(&eval_scalar.items); + } + RelOperator::Aggregate(aggregate) => { + self.collect_scalar_items(&aggregate.group_items); + self.collect_scalar_items(&aggregate.aggregate_functions); + } + RelOperator::Window(window) => { + self.definitions.insert( + window.index, + window + .arguments + .iter() + .map(|item| SourceExpr::Scalar(Box::new(item.scalar.clone()))) + .chain( + window + .partition_by + .iter() + .map(|item| SourceExpr::Scalar(Box::new(item.scalar.clone()))), + ) + .chain(window.order_by.iter().map(|item| { + SourceExpr::Scalar(Box::new(item.order_by_item.scalar.clone())) + })) + .collect(), + ); + } + RelOperator::WindowGroup(window_group) => { + self.collect_scalar_items(&window_group.scalar_items); + for window in &window_group.windows { + self.definitions.insert( + window.index, + window + .arguments + .iter() + .map(|item| SourceExpr::Scalar(Box::new(item.scalar.clone()))) + .chain( + window + .partition_by + .iter() + .map(|item| SourceExpr::Scalar(Box::new(item.scalar.clone()))), + ) + .chain(window.order_by.iter().map(|item| { + SourceExpr::Scalar(Box::new(item.order_by_item.scalar.clone())) + })) + .collect(), + ); + } + } + RelOperator::ProjectSet(project_set) => { + self.collect_scalar_items(&project_set.srfs); + } + RelOperator::Udf(udf) => { + self.collect_scalar_items(&udf.items); + } + RelOperator::AsyncFunction(async_function) => { + self.collect_scalar_items(&async_function.items); + } + RelOperator::UnionAll(union_all) => { + for (idx, (left, right)) in union_all.output_indexes.iter().zip( + union_all + .left_outputs + .iter() + .zip(union_all.right_outputs.iter()), + ) { + let mut sources = Vec::with_capacity(2); + sources.push(SourceExpr::Symbol(left.0)); + if let Some(cast) = &left.1 { + sources.push(SourceExpr::Scalar(Box::new(cast.clone()))); + } + sources.push(SourceExpr::Symbol(right.0)); + if let Some(cast) = &right.1 { + sources.push(SourceExpr::Scalar(Box::new(cast.clone()))); + } + self.definitions.insert(*idx, sources); + } + } + RelOperator::ExpressionScan(expression_scan) => { + for (idx, column_index) in expression_scan.column_indexes.iter().enumerate() { + let mut sources = Vec::new(); + for row in &expression_scan.values { + if let Some(expr) = row.get(idx) { + sources.push(SourceExpr::Scalar(Box::new(expr.clone()))); + } + } + self.definitions.insert(*column_index, sources); + } + } + RelOperator::MaterializedCTERef(cte_ref) => { + self.collect_s_expr(&cte_ref.def, metadata)?; + for (consumer, producer) in &cte_ref.column_mapping { + self.definitions + .insert(*consumer, vec![SourceExpr::Symbol(*producer)]); + } + } + _ => {} + } + + for child in s_expr.children() { + self.collect_s_expr(child, metadata)?; + } + Ok(()) + } + + fn collect_base_columns(&mut self, metadata: &Metadata, operator: &RelOperator) -> Result<()> { + let RelOperator::Scan(scan) = operator else { + return Ok(()); + }; + for column in &scan.columns { + let Some(source) = source_column_from_symbol(metadata, *column)? else { + continue; + }; + self.definitions + .insert(*column, vec![SourceExpr::Base(source)]); + } + Ok(()) + } + + fn collect_scalar_items(&mut self, items: &[ScalarItem]) { + for item in items { + self.definitions + .insert(item.index, vec![SourceExpr::Scalar(Box::new( + item.scalar.clone(), + ))]); + } + } + + fn resolve_target( + &mut self, + target: &TargetValue, + metadata: &Metadata, + ) -> Result> { + match target { + TargetValue::QueryOutput { + output_column_index, + .. + } => self.resolve_symbol(*output_column_index, metadata), + TargetValue::Expr { scalar } => self.resolve_scalar(scalar, metadata), + } + } + + fn resolve_symbol( + &mut self, + symbol: Symbol, + metadata: &Metadata, + ) -> Result> { + if !self.active_columns.insert(symbol) { + return Ok(BTreeSet::new()); + } + + let result = if let Some(source) = source_column_from_symbol(metadata, symbol)? { + BTreeSet::from([source]) + } else if let Some(definitions) = self.definitions.get(&symbol).cloned() { + let mut sources = BTreeSet::new(); + for definition in definitions { + sources.extend(self.resolve_source_expr(&definition, metadata)?); + } + sources + } else { + BTreeSet::new() + }; + + self.active_columns.remove(&symbol); + Ok(result) + } + + fn resolve_source_expr( + &mut self, + expr: &SourceExpr, + metadata: &Metadata, + ) -> Result> { + match expr { + SourceExpr::Symbol(symbol) => self.resolve_symbol(*symbol, metadata), + SourceExpr::Scalar(scalar) => self.resolve_scalar(scalar, metadata), + SourceExpr::Base(source) => Ok(BTreeSet::from([source.clone()])), + } + } + + fn resolve_scalar( + &mut self, + scalar: &ScalarExpr, + metadata: &Metadata, + ) -> Result> { + let mut visitor = SourceColumnVisitor { + resolver: self, + metadata, + columns: BTreeSet::new(), + }; + visitor.visit(scalar)?; + Ok(visitor.columns) + } +} + +struct SourceColumnVisitor<'a, 'b> { + resolver: &'a mut LineageResolver, + metadata: &'b Metadata, + columns: BTreeSet, +} + +impl<'a, 'b, 'c> Visitor<'c> for SourceColumnVisitor<'a, 'b> { + fn visit_bound_column_ref(&mut self, col: &'c BoundColumnRef) -> Result<()> { + self.columns.extend( + self.resolver + .resolve_symbol(col.column.index, self.metadata)?, + ); + Ok(()) + } + + fn visit_subquery(&mut self, subquery: &'c SubqueryExpr) -> Result<()> { + if let Some(child_expr) = subquery.child_expr.as_ref() { + self.visit(child_expr)?; + } + self.resolver + .collect_s_expr(&subquery.subquery, self.metadata)?; + self.columns.extend( + self.resolver + .resolve_symbol(subquery.output_column.index, self.metadata)?, + ); + Ok(()) + } +} + +fn source_column_from_symbol(metadata: &Metadata, symbol: Symbol) -> Result> { + if symbol.is_dummy_column() || symbol.as_usize() >= metadata.columns().len() { + return Ok(None); + } + + // A view is a lineage boundary even though its query is expanded in the + // plan. Resolve annotated output symbols to the view, not its base tables. + if let Some(source) = metadata.view_lineage_source_column(symbol) { + return Ok(Some(SourceColumn { + table: QueryLineageRelation { + catalog: source.relation.catalog.clone(), + database: source.relation.database.clone(), + name: source.relation.name.clone(), + id: Some(source.relation.id), + catalog_type: Some(CatalogType::Default), + kind: QueryLineageRelationKind::View, + }, + column: QueryLineageColumn { + name: source.name.clone(), + id: source.id, + }, + })); + } + + let column = metadata.column(symbol); + // Stream metadata describes the change process and is not source data. + // Regular stream data columns continue below and resolve through the + // stream's backing table. + if column.is_stream_column() { + return Ok(None); + } + let Some(table_index) = column.table_index() else { + return Ok(None); + }; + if table_index >= metadata.tables().len() { + return Ok(None); + } + let table = &metadata.tables()[table_index]; + if table.is_source_of_stage() { + let stage_info = match table.table().get_data_source_info() { + DataSourceInfo::StageSource(info) => Some(info.stage_info), + DataSourceInfo::ParquetSource(info) => Some(info.stage_info), + DataSourceInfo::ORCSource(info) => Some(info.stage_table_info.stage_info), + _ => None, + }; + let Some(relation) = stage_info.as_ref().and_then(stage_relation) else { + return Ok(None); + }; + return Ok(Some(SourceColumn { + table: relation, + column: QueryLineageColumn { + name: column.name(), + id: column_id(column), + }, + })); + } + let Some(relation) = + relation_info_from_table_index(metadata, table_index, QueryLineageRelationKind::Table)? + else { + return Ok(None); + }; + Ok(Some(SourceColumn { + table: relation, + column: QueryLineageColumn { + name: column.name(), + id: column_id(column), + }, + })) +} + +fn column_id(column: &ColumnEntry) -> ColumnId { + match column { + ColumnEntry::BaseTableColumn(base) => base.column_id, + ColumnEntry::InternalColumn(internal) => internal.internal_column.column_id(), + ColumnEntry::VirtualColumn(virtual_column) => virtual_column.column_id, + ColumnEntry::DerivedColumn(derived) => derived.column_index.as_usize() as ColumnId, + } +} + +fn relation_info_from_table_index( + metadata: &Metadata, + table_index: usize, + kind: QueryLineageRelationKind, +) -> Result> { + let Some(table) = metadata.tables().get(table_index) else { + return Ok(None); + }; + // Streams are transparent for lineage: attribute their data columns to the + // backing table. Views use the per-output-symbol boundary above instead. + if let Some(lineage_source) = table.stream_lineage_source() { + return Ok(Some(QueryLineageRelation { + catalog: lineage_source.catalog.clone(), + database: lineage_source.database.clone(), + name: lineage_source.name.clone(), + id: Some(lineage_source.id), + catalog_type: Some(CatalogType::Default), + kind, + })); + } + let table_object = table.table(); + let table_info = table_object.get_table_info(); + if table_object.is_temp() + || matches!( + table_info.engine().to_ascii_uppercase().as_str(), + "MEMORY" | "DELTA" + ) + { + return Ok(None); + } + let catalog_type = table_info.catalog_info.catalog_type(); + // External catalog ids are runtime identities rather than stable Databend metadata ids. + // Name-address their objects and columns across history batches. + let table_id = match catalog_type { + CatalogType::Default => { + (table_info.ident.table_id != 0).then_some(table_info.ident.table_id) + } + CatalogType::Hive | CatalogType::Iceberg | CatalogType::Paimon => None, + }; + Ok(Some(QueryLineageRelation { + catalog: table.catalog().to_string(), + database: table.database().to_string(), + name: table.name().to_string(), + id: table_id, + catalog_type: Some(catalog_type), + kind, + })) +} + +fn group_sources_by_table(sources: BTreeSet) -> Vec { + let mut tables: BTreeMap> = BTreeMap::new(); + for source in sources { + tables + .entry(source.table) + .or_default() + .insert(source.column); + } + tables + .into_iter() + .map(|(table, columns)| SourceTableColumns { + table, + columns: columns.into_iter().collect(), + }) + .collect() +} + +fn query_parts(plan: &Plan) -> Result<(&SExpr, &MetadataRef, &BindContext)> { + match plan { + Plan::Query { + s_expr, + metadata, + bind_context, + .. + } => Ok((s_expr, metadata, bind_context)), + _ => Err(ErrorCode::Internal( + "Lineage extraction expects a query plan".to_string(), + )), + } +} + +fn query_plan(plan: &Plan) -> Result<&Plan> { + match plan { + Plan::CreateTable(plan) => plan.as_select.as_deref().ok_or_else(|| { + ErrorCode::Internal("CTAS lineage extraction expects as_select".to_string()) + }), + Plan::CreateView(plan) => plan.query_plan.as_deref().ok_or_else(|| { + ErrorCode::Internal("Create view lineage extraction expects query plan".to_string()) + }), + Plan::Insert(plan) => match &plan.source { + InsertInputSource::SelectPlan(query) => Ok(query), + _ => Err(ErrorCode::Internal( + "Insert lineage extraction expects select source".to_string(), + )), + }, + Plan::Replace(plan) => match &plan.source { + InsertInputSource::SelectPlan(query) => Ok(query), + _ => Err(ErrorCode::Internal( + "Replace lineage extraction expects select source".to_string(), + )), + }, + _ => Err(ErrorCode::Internal( + "Unsupported relation lineage plan".to_string(), + )), + } +} + +fn column_info_from_table_field(field: &TableField) -> QueryLineageColumn { + QueryLineageColumn { + name: field.name().to_string(), + id: field.column_id(), + } +} + +fn column_info_from_data_field(name: &str, ordinal: usize) -> QueryLineageColumn { + QueryLineageColumn { + name: name.to_string(), + id: ordinal as ColumnId, + } +} + +fn find_mutation(s_expr: &SExpr) -> Option<&crate::plans::Mutation> { + match s_expr.plan() { + RelOperator::Mutation(mutation) => Some(mutation), + _ => s_expr.children().find_map(find_mutation), + } +} + +fn target_column_from_field_index( + metadata: &Metadata, + mutation: &crate::plans::Mutation, + field_index: FieldIndex, +) -> Option { + let column_entries = metadata.columns_by_table_index(mutation.target_table_index); + if let Some(column_index) = mutation.field_index_map.get(&field_index) + && let Some(column_entry) = column_entries + .iter() + .find(|entry| entry.index().to_string() == *column_index) + { + return Some(QueryLineageColumn { + name: column_entry.name(), + id: column_id(column_entry), + }); + } + + column_entries + .get(field_index) + .map(|entry| QueryLineageColumn { + name: entry.name(), + id: column_id(entry), + }) +} + +trait SourceSchemaExt { + fn field(&self, index: usize) -> Result<&databend_common_expression::DataField>; +} + +impl SourceSchemaExt for databend_common_expression::DataSchemaRef { + fn field(&self, index: usize) -> Result<&databend_common_expression::DataField> { + self.fields().get(index).ok_or_else(|| { + ErrorCode::Internal(format!("Data schema field index {} out of bounds", index)) + }) + } +} + +#[cfg(test)] +mod tests { + use std::any::Any; + use std::collections::BTreeMap; + use std::collections::HashMap; + use std::sync::Arc; + + use databend_common_ast::ast::Engine; + use databend_common_catalog::plan::StreamColumn; + use databend_common_catalog::plan::StreamColumnType; + use databend_common_catalog::table::Table; + use databend_common_expression::ORIGIN_VERSION_COL_NAME; + use databend_common_expression::TableDataType; + use databend_common_expression::TableField; + use databend_common_expression::TableSchema; + use databend_common_expression::types::DataType; + use databend_common_expression::types::NumberDataType; + use databend_common_meta_app::schema::CatalogInfo; + use databend_common_meta_app::schema::CatalogOption; + use databend_common_meta_app::schema::CreateOption; + use databend_common_meta_app::schema::DatabaseType; + use databend_common_meta_app::schema::HiveCatalogOption; + use databend_common_meta_app::schema::IcebergCatalogOption; + use databend_common_meta_app::schema::IcebergRestCatalogOption; + use databend_common_meta_app::schema::PaimonCatalogOption; + use databend_common_meta_app::schema::TableIdent; + use databend_common_meta_app::schema::TableInfo; + use databend_common_meta_app::schema::TableMeta; + use databend_common_meta_app::tenant::Tenant; + use parking_lot::RwLock; + + use super::*; + use crate::ColumnBindingBuilder; + use crate::LineageSourceRelation; + use crate::ViewLineageSourceColumn; + use crate::Visibility; + use crate::plans::CreateTablePlan; + use crate::plans::EvalScalar; + use crate::plans::Filter; + use crate::plans::FunctionCall; + use crate::plans::MaterializedCTERef; + use crate::plans::Scan; + + #[derive(Debug)] + struct FakeTable { + table_info: TableInfo, + stream_source_table_info: Option, + stream_columns: Vec, + data_source_info: Option, + } + + #[async_trait::async_trait] + impl Table for FakeTable { + fn as_any(&self) -> &dyn Any { + self + } + + fn get_table_info(&self) -> &TableInfo { + &self.table_info + } + + fn stream_source_table_info(&self) -> Option<&TableInfo> { + self.stream_source_table_info.as_ref() + } + + fn stream_columns(&self) -> Vec { + self.stream_columns.clone() + } + + fn get_data_source_info(&self) -> DataSourceInfo { + self.data_source_info + .clone() + .unwrap_or_else(|| DataSourceInfo::TableSource(self.table_info.clone())) + } + } + + #[test] + fn test_query_output_lineage_excludes_filter_columns() -> Result<()> { + // Simulates: + // INSERT INTO dst SELECT a + b AS x FROM src WHERE c + // The target column `dst.x` depends on `src.a` and `src.b`; `src.c` is filter-only. + let metadata = MetadataRef::new(RwLock::new(Metadata::default())); + let table_index = add_fake_table(&metadata, 10, "src", &["a", "b", "c"]); + let a = column_index(&metadata, table_index, "a"); + let b = column_index(&metadata, table_index, "b"); + let c = column_index(&metadata, table_index, "c"); + let x = metadata + .write() + .add_derived_column("x".to_string(), int_data_type()); + + let scan = scan_expr(&metadata, table_index); + let filter = scan.build_unary(Filter { + predicates: vec![bound_column(c, "c", Some(table_index))], + }); + let s_expr = filter.build_unary(EvalScalar { + items: vec![ScalarItem { + scalar: plus( + bound_column(a, "a", Some(table_index)), + bound_column(b, "b", Some(table_index)), + ), + index: x, + }], + }); + let query = query_plan(metadata.clone(), s_expr, vec![binding(x, "x", None)]); + + let lineage = RelationLineage::from_query_outputs( + &query, + relation( + "default", + "default", + "dst", + Some(20), + QueryLineageRelationKind::Table, + ), + vec![QueryLineageColumn { + name: "x".to_string(), + id: 0, + }], + )?; + + assert_source_columns(&lineage, "dst", "x", &["a", "b"]); + Ok(()) + } + + #[test] + fn test_expr_lineage_resolves_target_binding() -> Result<()> { + // Simulates: + // UPDATE dst SET x = src.a + src.b + // DML adapters pass SET/VALUES expressions as TargetValue::Expr. + let metadata = MetadataRef::new(RwLock::new(Metadata::default())); + let table_index = add_fake_table(&metadata, 10, "src", &["a", "b"]); + let a = column_index(&metadata, table_index, "a"); + let b = column_index(&metadata, table_index, "b"); + let query = query_plan(metadata.clone(), scan_expr(&metadata, table_index), vec![]); + + let lineage = RelationLineage::from_query_plan(&query, vec![TargetColumnBinding { + target_relation: relation( + "default", + "default", + "dst", + Some(20), + QueryLineageRelationKind::Table, + ), + target_column: QueryLineageColumn { + name: "x".to_string(), + id: 0, + }, + value: TargetValue::Expr { + scalar: Box::new(plus( + bound_column(a, "a", Some(table_index)), + bound_column(b, "b", Some(table_index)), + )), + }, + }])?; + + assert_source_columns(&lineage, "dst", "x", &["a", "b"]); + Ok(()) + } + + #[test] + fn test_stream_lineage_uses_base_table_and_skips_stream_columns() -> Result<()> { + // Simulates: + // INSERT INTO dst SELECT a + _origin_version AS x FROM stream_src + // Stream data columns are attributed to the base table; stream metadata columns are ignored. + let metadata = MetadataRef::new(RwLock::new(Metadata::default())); + let table_index = add_fake_stream_table( + &metadata, + 10, + "stream_src", + relation( + "default", + "default", + "base_src", + Some(30), + QueryLineageRelationKind::Table, + ), + ); + let a = column_index(&metadata, table_index, "a"); + let origin_version = column_index(&metadata, table_index, ORIGIN_VERSION_COL_NAME); + let x = metadata + .write() + .add_derived_column("x".to_string(), int_data_type()); + + let s_expr = scan_expr(&metadata, table_index).build_unary(EvalScalar { + items: vec![ScalarItem { + scalar: plus( + bound_column(a, "a", Some(table_index)), + bound_column(origin_version, ORIGIN_VERSION_COL_NAME, Some(table_index)), + ), + index: x, + }], + }); + let query = query_plan(metadata.clone(), s_expr, vec![binding(x, "x", None)]); + + let lineage = RelationLineage::from_query_outputs( + &query, + relation( + "default", + "default", + "dst", + Some(20), + QueryLineageRelationKind::Table, + ), + vec![QueryLineageColumn { + name: "x".to_string(), + id: 0, + }], + )?; + + assert_source_columns_qualified(&lineage, "dst", "x", &["base_src.a"]); + Ok(()) + } + + #[test] + fn test_catalog_type_lineage_addressing() -> Result<()> { + for (catalog_type, engine, expected_id) in [ + (CatalogType::Default, "FUSE", Some(10)), + (CatalogType::Hive, "HIVE", None), + (CatalogType::Iceberg, "ICEBERG", None), + (CatalogType::Paimon, "PAIMON", None), + ] { + let metadata = MetadataRef::new(RwLock::new(Metadata::default())); + let table_index = add_fake_table_with_catalog_type( + &metadata, + 10, + "src", + &["a"], + engine, + catalog_type, + ); + let relation = relation_info_from_table_index( + &metadata.read(), + table_index, + QueryLineageRelationKind::Table, + )? + .expect("supported engine should produce a lineage relation"); + assert_eq!(relation.id, expected_id, "catalog_type={catalog_type:?}"); + assert_eq!(relation.catalog_type, Some(catalog_type)); + } + + for engine in ["MEMORY", "DELTA"] { + let metadata = MetadataRef::new(RwLock::new(Metadata::default())); + let table_index = add_fake_table_with_engine(&metadata, 10, "src", &["a"], engine); + assert_eq!( + relation_info_from_table_index( + &metadata.read(), + table_index, + QueryLineageRelationKind::Table, + )?, + None, + "engine={engine}" + ); + } + Ok(()) + } + + #[test] + fn test_regular_table_has_no_stream_lineage_source() { + let metadata = MetadataRef::new(RwLock::new(Metadata::default())); + let table_index = add_fake_table(&metadata, 10, "src", &["a"]); + + assert!( + metadata + .read() + .table(table_index) + .stream_lineage_source() + .is_none() + ); + } + + #[test] + fn test_stage_scan_lineage_uses_named_stage() -> Result<()> { + // Simulates: INSERT INTO dst SELECT a FROM @named_stage. + // Runtime Stage tables use placeholder TableInfo, so lineage must use StageTableInfo. + let metadata = MetadataRef::new(RwLock::new(Metadata::default())); + let table_index = add_fake_stage_table(&metadata, StageInfo { + stage_name: "named_stage".to_string(), + stage_type: StageType::Internal, + ..Default::default() + }); + let a = column_index(&metadata, table_index, "a"); + let query = query_plan(metadata.clone(), scan_expr(&metadata, table_index), vec![ + binding(a, "a", Some(table_index)), + ]); + + let lineage = RelationLineage::from_query_outputs( + &query, + relation( + "default", + "default", + "dst", + Some(20), + QueryLineageRelationKind::Table, + ), + vec![QueryLineageColumn { + name: "x".to_string(), + id: 0, + }], + )?; + + let source = &lineage.relations[0].columns[0].source_tables[0].table; + assert_eq!(source.kind, QueryLineageRelationKind::Stage); + assert_eq!(source.name, "named_stage"); + assert_eq!(source.id, None); + Ok(()) + } + + #[test] + fn test_unstable_stages_are_not_lineage_sources() { + let temporary = StageInfo { + stage_name: "s3://bucket/path".to_string(), + is_temporary: true, + ..Default::default() + }; + let user = StageInfo { + stage_name: "user_name".to_string(), + stage_type: StageType::User, + ..Default::default() + }; + + assert_eq!(stage_relation(&temporary), None); + assert_eq!(stage_relation(&user), None); + } + + #[test] + fn test_view_lineage_stops_at_view_boundary() -> Result<()> { + // Simulates: + // CREATE VIEW v(vx) AS SELECT a + b FROM src; + // INSERT INTO dst SELECT vx FROM v; + // The execution plan expands the view, but lineage records `v.vx` instead of `src.a/src.b`. + let metadata = MetadataRef::new(RwLock::new(Metadata::default())); + let table_index = add_fake_table(&metadata, 10, "src", &["a", "b"]); + let a = column_index(&metadata, table_index, "a"); + let b = column_index(&metadata, table_index, "b"); + let vx = metadata + .write() + .add_derived_column("vx".to_string(), int_data_type()); + metadata + .write() + .add_view_lineage_source_column(vx, ViewLineageSourceColumn { + relation: LineageSourceRelation { + catalog: "default".to_string(), + database: "default".to_string(), + name: "v".to_string(), + id: 30, + }, + name: "vx".to_string(), + id: 0, + }); + + let s_expr = scan_expr(&metadata, table_index).build_unary(EvalScalar { + items: vec![ScalarItem { + scalar: plus( + bound_column(a, "a", Some(table_index)), + bound_column(b, "b", Some(table_index)), + ), + index: vx, + }], + }); + let query = query_plan(metadata.clone(), s_expr, vec![binding(vx, "vx", None)]); + + let lineage = RelationLineage::from_query_outputs( + &query, + relation( + "default", + "default", + "dst", + Some(20), + QueryLineageRelationKind::Table, + ), + vec![QueryLineageColumn { + name: "x".to_string(), + id: 0, + }], + )?; + + assert_source_columns_qualified(&lineage, "dst", "x", &["v.vx"]); + assert_eq!( + lineage.relations[0].columns[0].source_tables[0].table.kind, + QueryLineageRelationKind::View + ); + Ok(()) + } + + #[test] + fn test_ctas_query_lineage_from_plan() -> Result<()> { + // Simulates: + // CREATE TABLE dst AS SELECT a + b AS x FROM src WHERE c + // CTAS gets its target columns from the create-table schema and query outputs. + let metadata = MetadataRef::new(RwLock::new(Metadata::default())); + let table_index = add_fake_table(&metadata, 10, "src", &["a", "b", "c"]); + let a = column_index(&metadata, table_index, "a"); + let b = column_index(&metadata, table_index, "b"); + let c = column_index(&metadata, table_index, "c"); + let x = metadata + .write() + .add_derived_column("x".to_string(), int_data_type()); + + let scan = scan_expr(&metadata, table_index); + let filter = scan.build_unary(Filter { + predicates: vec![bound_column(c, "c", Some(table_index))], + }); + let s_expr = filter.build_unary(EvalScalar { + items: vec![ScalarItem { + scalar: plus( + bound_column(a, "a", Some(table_index)), + bound_column(b, "b", Some(table_index)), + ), + index: x, + }], + }); + let query = query_plan(metadata.clone(), s_expr, vec![binding(x, "x", None)]); + let plan = Plan::CreateTable(Box::new(create_table_plan("dst", &["x"], Some(query)))); + + let lineage = plan.query_lineage()?.expect("CTAS should have lineage"); + + assert_eq!(lineage.kind, QueryLineageKind::Ctas); + assert_query_source_columns(&lineage, "dst", "x", &["src.a", "src.b"]); + Ok(()) + } + + #[test] + fn test_materialized_cte_ref_lineage_from_plan() -> Result<()> { + // Simulates the optimized shape for: + // INSERT INTO dst WITH q AS MATERIALIZED (SELECT a + b AS x FROM src) SELECT x FROM q + // The MaterializedCTERef maps each consumer output symbol to its producer symbol in `def`. + let metadata = MetadataRef::new(RwLock::new(Metadata::default())); + let table_index = add_fake_table(&metadata, 10, "src", &["a", "b"]); + let a = column_index(&metadata, table_index, "a"); + let b = column_index(&metadata, table_index, "b"); + let producer_x = metadata + .write() + .add_derived_column("producer_x".to_string(), int_data_type()); + let consumer_x = metadata + .write() + .add_derived_column("consumer_x".to_string(), int_data_type()); + + let def = scan_expr(&metadata, table_index).build_unary(EvalScalar { + items: vec![ScalarItem { + scalar: plus( + bound_column(a, "a", Some(table_index)), + bound_column(b, "b", Some(table_index)), + ), + index: producer_x, + }], + }); + let s_expr = SExpr::create_leaf(RelOperator::MaterializedCTERef(MaterializedCTERef { + cte_name: "q".to_string(), + output_columns: vec![consumer_x], + def, + column_mapping: HashMap::from([(consumer_x, producer_x)]), + stat_info: None, + })); + let query = query_plan(metadata.clone(), s_expr, vec![binding( + consumer_x, "x", None, + )]); + + let lineage = RelationLineage::from_query_outputs( + &query, + relation( + "default", + "default", + "dst", + Some(20), + QueryLineageRelationKind::Table, + ), + vec![QueryLineageColumn { + name: "x".to_string(), + id: 0, + }], + )?; + + assert_source_columns_qualified(&lineage, "dst", "x", &["src.a", "src.b"]); + Ok(()) + } + + fn add_fake_table( + metadata: &MetadataRef, + table_id: u64, + table_name: &str, + columns: &[&str], + ) -> usize { + add_fake_table_with_engine(metadata, table_id, table_name, columns, "FUSE") + } + + fn add_fake_table_with_engine( + metadata: &MetadataRef, + table_id: u64, + table_name: &str, + columns: &[&str], + engine: &str, + ) -> usize { + add_fake_table_with_catalog_type( + metadata, + table_id, + table_name, + columns, + engine, + CatalogType::Default, + ) + } + + fn add_fake_table_with_catalog_type( + metadata: &MetadataRef, + table_id: u64, + table_name: &str, + columns: &[&str], + engine: &str, + catalog_type: CatalogType, + ) -> usize { + metadata.write().add_table( + catalog_name(catalog_type), + "default".to_string(), + fake_table_with_catalog_type(table_id, table_name, columns, engine, catalog_type), + None, + None, + false, + false, + false, + None, + ) + } + + fn add_fake_stream_table( + metadata: &MetadataRef, + table_id: u64, + table_name: &str, + lineage_source: QueryLineageRelation, + ) -> usize { + let mut metadata = metadata.write(); + let table_index = metadata.add_table( + "default".to_string(), + "default".to_string(), + fake_stream_table(table_id, table_name, lineage_source.clone()), + None, + None, + false, + false, + false, + None, + ); + metadata.set_stream_lineage_source(table_index, LineageSourceRelation { + catalog: lineage_source.catalog, + database: lineage_source.database, + name: lineage_source.name, + id: lineage_source.id.expect("base source table id must be set"), + }); + table_index + } + + fn fake_table_with_catalog_type( + table_id: u64, + table_name: &str, + columns: &[&str], + engine: &str, + catalog_type: CatalogType, + ) -> Arc { + Arc::new(FakeTable { + table_info: table_info_with_catalog_type( + table_id, + table_name, + columns, + engine, + catalog_type, + ), + stream_source_table_info: None, + stream_columns: vec![], + data_source_info: None, + }) + } + + fn add_fake_stage_table(metadata: &MetadataRef, stage_info: StageInfo) -> usize { + let schema = Arc::new(TableSchema::new(vec![TableField::new( + "a", + TableDataType::Number(NumberDataType::Int32), + )])); + let table_info = TableInfo { + name: "stage".to_string(), + meta: TableMeta { + schema: schema.clone(), + engine: "STAGE".to_string(), + ..Default::default() + }, + ..Default::default() + }; + metadata.write().add_table( + "default".to_string(), + "system".to_string(), + Arc::new(FakeTable { + table_info, + stream_source_table_info: None, + stream_columns: vec![], + data_source_info: Some(DataSourceInfo::StageSource( + databend_common_catalog::plan::StageTableInfo { + stage_info, + schema, + ..Default::default() + }, + )), + }), + None, + None, + false, + false, + true, + None, + ) + } + + fn create_table_plan( + table_name: &str, + columns: &[&str], + as_select: Option, + ) -> CreateTablePlan { + CreateTablePlan { + create_option: CreateOption::Create, + tenant: Tenant::new_literal("default"), + catalog: "default".to_string(), + database: "default".to_string(), + table: table_name.to_string(), + schema: Arc::new(TableSchema::new( + columns + .iter() + .map(|column| { + TableField::new(column, TableDataType::Number(NumberDataType::Int32)) + }) + .collect(), + )), + engine: Engine::Fuse, + engine_options: BTreeMap::new(), + storage_params: None, + options: BTreeMap::new(), + table_properties: None, + table_partition: None, + field_comments: vec![], + field_stats_truncate_len: vec![], + cluster_key: None, + as_select: as_select.map(Box::new), + table_indexes: None, + table_constraints: None, + attached_columns: None, + } + } + + fn fake_stream_table( + table_id: u64, + table_name: &str, + lineage_source: QueryLineageRelation, + ) -> Arc { + let fields = vec![TableField::new( + "a", + TableDataType::Number(NumberDataType::Int32), + )]; + Arc::new(FakeTable { + table_info: TableInfo { + ident: TableIdent::new(table_id, 0), + desc: format!("'default'.'{table_name}'"), + name: table_name.to_string(), + meta: TableMeta { + schema: Arc::new(TableSchema::new(fields)), + engine: "STREAM".to_string(), + ..Default::default() + }, + catalog_info: Arc::new(CatalogInfo::default()), + db_type: DatabaseType::NormalDB, + }, + stream_source_table_info: Some(table_info( + lineage_source.id.expect("base source table id must be set"), + &lineage_source.name, + &["a"], + "FUSE", + )), + stream_columns: vec![StreamColumn::new( + ORIGIN_VERSION_COL_NAME, + StreamColumnType::OriginVersion, + )], + data_source_info: None, + }) + } + + fn table_info(table_id: u64, table_name: &str, columns: &[&str], engine: &str) -> TableInfo { + table_info_with_catalog_type(table_id, table_name, columns, engine, CatalogType::Default) + } + + fn table_info_with_catalog_type( + table_id: u64, + table_name: &str, + columns: &[&str], + engine: &str, + catalog_type: CatalogType, + ) -> TableInfo { + TableInfo { + ident: TableIdent::new(table_id, 0), + desc: format!("'default'.'{table_name}'"), + name: table_name.to_string(), + meta: TableMeta { + schema: Arc::new(TableSchema::new( + columns + .iter() + .map(|column| { + TableField::new(column, TableDataType::Number(NumberDataType::Int32)) + }) + .collect(), + )), + engine: engine.to_string(), + ..Default::default() + }, + catalog_info: catalog_info(catalog_type), + db_type: DatabaseType::NormalDB, + } + } + + fn catalog_name(catalog_type: CatalogType) -> String { + match catalog_type { + CatalogType::Default => "default", + CatalogType::Hive => "hive", + CatalogType::Iceberg => "iceberg", + CatalogType::Paimon => "paimon", + } + .to_string() + } + + fn catalog_info(catalog_type: CatalogType) -> Arc { + let mut info = CatalogInfo::default(); + info.name_ident.catalog_name = catalog_name(catalog_type); + info.meta.catalog_option = match catalog_type { + CatalogType::Default => CatalogOption::Default, + CatalogType::Hive => CatalogOption::Hive(HiveCatalogOption { + address: String::new(), + storage_params: None, + }), + CatalogType::Iceberg => { + CatalogOption::Iceberg(IcebergCatalogOption::Rest(IcebergRestCatalogOption { + uri: String::new(), + warehouse: String::new(), + props: HashMap::new(), + })) + } + CatalogType::Paimon => CatalogOption::Paimon(PaimonCatalogOption { + options: HashMap::new(), + }), + }; + Arc::new(info) + } + + fn scan_expr(metadata: &MetadataRef, table_index: usize) -> SExpr { + let columns = metadata + .read() + .columns_by_table_index(table_index) + .into_iter() + .map(|column| column.index()) + .collect(); + SExpr::create_leaf(Arc::new(RelOperator::Scan(Scan { + table_index, + columns, + ..Default::default() + }))) + } + + fn query_plan( + metadata: MetadataRef, + s_expr: SExpr, + columns: Vec, + ) -> Plan { + let bind_context = BindContext { + columns, + ..Default::default() + }; + Plan::Query { + s_expr: Box::new(s_expr), + metadata, + bind_context: Box::new(bind_context), + rewrite_kind: None, + formatted_ast: None, + ignore_result: false, + } + } + + fn column_index(metadata: &MetadataRef, table_index: usize, name: &str) -> Symbol { + metadata + .read() + .columns_by_table_index(table_index) + .into_iter() + .find(|column| column.name() == name) + .unwrap_or_else(|| panic!("missing column {name}")) + .index() + } + + fn bound_column(index: Symbol, name: &str, table_index: Option) -> ScalarExpr { + BoundColumnRef { + span: None, + column: binding(index, name, table_index), + } + .into() + } + + fn binding(index: Symbol, name: &str, table_index: Option) -> crate::ColumnBinding { + ColumnBindingBuilder::new( + name.to_string(), + index, + Box::new(int_data_type()), + Visibility::Visible, + ) + .table_index(table_index) + .build() + } + + fn plus(left: ScalarExpr, right: ScalarExpr) -> ScalarExpr { + FunctionCall { + span: None, + func_name: "plus".to_string(), + params: vec![], + arguments: vec![left, right], + } + .into() + } + + fn int_data_type() -> DataType { + DataType::Number(NumberDataType::Int32) + } + + fn relation( + catalog: &str, + database: &str, + name: &str, + id: Option, + kind: QueryLineageRelationKind, + ) -> QueryLineageRelation { + QueryLineageRelation { + catalog: catalog.to_string(), + database: database.to_string(), + name: name.to_string(), + id, + catalog_type: Some(CatalogType::Default), + kind, + } + } + + fn assert_source_columns( + lineage: &RelationLineage, + target_table: &str, + target_column: &str, + expected_sources: &[&str], + ) { + let table = lineage + .relations + .iter() + .find(|table| table.target.name == target_table) + .unwrap_or_else(|| panic!("missing target table {target_table}: {lineage:?}")); + let column = table + .columns + .iter() + .find(|column| column.target_column.name == target_column) + .unwrap_or_else(|| panic!("missing target column {target_column}: {lineage:?}")); + let mut sources = column + .source_tables + .iter() + .flat_map(|table| table.columns.iter().map(|column| column.name.as_str())) + .collect::>(); + sources.sort_unstable(); + assert_eq!(sources, expected_sources, "unexpected lineage: {lineage:?}"); + } + + fn assert_source_columns_qualified( + lineage: &RelationLineage, + target_table: &str, + target_column: &str, + expected_sources: &[&str], + ) { + let table = lineage + .relations + .iter() + .find(|table| table.target.name == target_table) + .unwrap_or_else(|| panic!("missing target table {target_table}: {lineage:?}")); + let column = table + .columns + .iter() + .find(|column| column.target_column.name == target_column) + .unwrap_or_else(|| panic!("missing target column {target_column}: {lineage:?}")); + let mut sources = column + .source_tables + .iter() + .flat_map(|table| { + table + .columns + .iter() + .map(|column| format!("{}.{}", table.table.name, column.name)) + }) + .collect::>(); + sources.sort_unstable(); + assert_eq!(sources, expected_sources, "unexpected lineage: {lineage:?}"); + } + + fn assert_query_source_columns( + lineage: &QueryLineage, + target_table: &str, + target_column: &str, + expected_sources: &[&str], + ) { + let mut sources = lineage + .targets + .iter() + .find(|table| table.relation.name == target_table) + .unwrap_or_else(|| panic!("missing target table {target_table}: {lineage:?}")) + .sources + .iter() + .flat_map(|table| { + table + .columns + .iter() + .filter(|column| column.target.name == target_column) + .map(|column| format!("{}.{}", table.relation.name, column.source.name)) + }) + .collect::>(); + sources.sort_unstable(); + assert_eq!(sources, expected_sources, "unexpected lineage: {lineage:?}"); + } +} diff --git a/src/query/sql/src/planner/metadata/metadata.rs b/src/query/sql/src/planner/metadata/metadata.rs index 2eb43cf50afd7..887d106a09288 100644 --- a/src/query/sql/src/planner/metadata/metadata.rs +++ b/src/query/sql/src/planner/metadata/metadata.rs @@ -83,6 +83,9 @@ pub struct Metadata { next_scan_id: usize, /// Mappings from base column index to scan id. base_column_scan_id: HashMap, + /// View output symbols that should remain at the view boundary after the + /// view query is expanded into the surrounding plan. + view_lineage_source_columns: HashMap, next_runtime_filter_id: usize, next_logical_recursive_cte_id: u32, next_materialized_cte_id: usize, @@ -408,6 +411,7 @@ impl Metadata { source_of_view, source_of_index, source_of_stage, + stream_lineage_source: None, }; self.tables.push(table_entry); let table_schema = table_meta.schema_with_stream(); @@ -535,6 +539,30 @@ impl Metadata { self.base_column_scan_id.get(&column_index).cloned() } + pub(crate) fn add_view_lineage_source_column( + &mut self, + column_index: Symbol, + source_column: ViewLineageSourceColumn, + ) { + self.view_lineage_source_columns + .insert(column_index, source_column); + } + + pub(crate) fn view_lineage_source_column( + &self, + column_index: Symbol, + ) -> Option<&ViewLineageSourceColumn> { + self.view_lineage_source_columns.get(&column_index) + } + + pub(crate) fn set_stream_lineage_source( + &mut self, + table_index: IndexType, + relation: LineageSourceRelation, + ) { + self.tables[table_index].stream_lineage_source = Some(relation); + } + pub fn replace_all_tables(&mut self, table: Arc) { for entry in self.tables.iter_mut() { entry.table = table.clone(); @@ -556,9 +584,29 @@ pub struct TableEntry { source_of_index: bool, source_of_stage: bool, + /// Source relation for a transparent stream scan. Stream data columns are + /// attributed to this relation; stream metadata columns are excluded. + stream_lineage_source: Option, table: Arc, } +/// Relation identity shared by the explicit Stream relation and View column +/// lineage annotations. This is a value object, not a generic extension point. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct LineageSourceRelation { + pub(crate) catalog: String, + pub(crate) database: String, + pub(crate) name: String, + pub(crate) id: u64, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct ViewLineageSourceColumn { + pub(crate) relation: LineageSourceRelation, + pub(crate) name: String, + pub(crate) id: ColumnId, +} + impl Debug for TableEntry { fn fmt(&self, f: &mut Formatter) -> std::fmt::Result { f.debug_struct("TableEntry") @@ -620,6 +668,10 @@ impl TableEntry { self.source_of_index } + pub(crate) fn stream_lineage_source(&self) -> Option<&LineageSourceRelation> { + self.stream_lineage_source.as_ref() + } + pub fn update_table_index(&mut self, table_index: IndexType) { self.index = table_index; } diff --git a/src/query/sql/src/planner/mod.rs b/src/query/sql/src/planner/mod.rs index ddff0f75a140a..610cea9aac85c 100644 --- a/src/query/sql/src/planner/mod.rs +++ b/src/query/sql/src/planner/mod.rs @@ -13,6 +13,7 @@ // limitations under the License. mod format; +mod lineage; mod metadata; #[allow(clippy::module_inception)] mod planner; @@ -38,6 +39,7 @@ pub use binder::parse_result_scan_args; pub use execution::*; pub use expression::*; pub use format::*; +pub use lineage::*; pub use metadata::*; pub use optimizer::optimize; pub use planner::PlanExtras; diff --git a/src/query/sql/src/planner/optimizer/optimizer.rs b/src/query/sql/src/planner/optimizer/optimizer.rs index 26eb554d9c476..a4379a81d2e51 100644 --- a/src/query/sql/src/planner/optimizer/optimizer.rs +++ b/src/query/sql/src/planner/optimizer/optimizer.rs @@ -220,7 +220,14 @@ pub async fn optimize(opt_ctx: Arc, plan: Plan) -> Result { + if let Some(p) = &plan.query_plan { + let optimized_plan = optimize(opt_ctx.clone(), *p.clone()).await?; + plan.query_plan = Some(Box::new(optimized_plan)); + } + Ok(Plan::CreateView(plan)) + } Plan::Set(mut plan) => { if let SetScalarsOrQuery::Query(q) = plan.values { let optimized_plan = optimize(opt_ctx.clone(), *q.clone()).await?; diff --git a/src/query/sql/src/planner/plans/ddl/view.rs b/src/query/sql/src/planner/plans/ddl/view.rs index 18620e3f64590..965b0bf801828 100644 --- a/src/query/sql/src/planner/plans/ddl/view.rs +++ b/src/query/sql/src/planner/plans/ddl/view.rs @@ -16,7 +16,9 @@ use databend_common_expression::DataSchemaRef; use databend_common_meta_app::schema::CreateOption; use databend_common_meta_app::tenant::Tenant; -#[derive(Clone, Debug, PartialEq, Eq)] +use crate::plans::Plan; + +#[derive(Clone, Debug)] pub struct CreateViewPlan { pub create_option: CreateOption, pub tenant: Tenant, @@ -25,8 +27,23 @@ pub struct CreateViewPlan { pub view_name: String, pub column_names: Vec, pub subquery: String, + pub query_plan: Option>, +} + +impl PartialEq for CreateViewPlan { + fn eq(&self, other: &Self) -> bool { + self.create_option == other.create_option + && self.tenant == other.tenant + && self.catalog == other.catalog + && self.database == other.database + && self.view_name == other.view_name + && self.column_names == other.column_names + && self.subquery == other.subquery + } } +impl Eq for CreateViewPlan {} + #[derive(Clone, Debug, PartialEq, Eq)] pub struct AlterViewPlan { pub tenant: Tenant, diff --git a/src/query/sql/src/planner/plans/insert.rs b/src/query/sql/src/planner/plans/insert.rs index a73c4493cd130..1b2cf50e31e56 100644 --- a/src/query/sql/src/planner/plans/insert.rs +++ b/src/query/sql/src/planner/plans/insert.rs @@ -29,6 +29,7 @@ use databend_common_expression::types::DataType; use databend_common_expression::types::NumberDataType; use databend_common_expression::types::StringType; use databend_common_meta_app::principal::FileFormatParams; +use databend_common_meta_app::schema::CatalogType; use databend_common_meta_app::schema::TableInfo; use enum_as_inner::EnumAsInner; use parking_lot::Mutex; @@ -89,6 +90,10 @@ pub struct Insert { // it should be provided as some `table_info`. // otherwise, the table being inserted will be resolved by using `catalog`.`database`.`table` pub table_info: Option, + /// Target table id captured for lineage extraction only. Execution must + /// continue to resolve ordinary INSERT targets by name. + pub lineage_target_table_id: Option, + pub lineage_target_catalog_type: CatalogType, } impl PartialEq for Insert { @@ -126,6 +131,8 @@ impl Insert { overwrite, // table_info only used create table as select. table_info: _, + lineage_target_table_id: _, + lineage_target_catalog_type: _, source, } = self; diff --git a/src/query/sql/test-support/src/lib.rs b/src/query/sql/test-support/src/lib.rs index e171ece5e0519..58f650f40c188 100644 --- a/src/query/sql/test-support/src/lib.rs +++ b/src/query/sql/test-support/src/lib.rs @@ -18,4 +18,5 @@ pub mod optimizer; pub use lite_context::FrequencyStatsMap; pub use lite_context::LiteTableContext; pub use lite_context::init_testing_globals; +pub use lite_context::init_testing_globals_with_config; pub use optimizer::*; diff --git a/src/query/sql/test-support/src/lite_context.rs b/src/query/sql/test-support/src/lite_context.rs index bc94246ba0c72..641485737f8a7 100644 --- a/src/query/sql/test-support/src/lite_context.rs +++ b/src/query/sql/test-support/src/lite_context.rs @@ -146,14 +146,17 @@ thread_local! { } pub fn init_testing_globals() { + init_testing_globals_with_config(InnerConfig::default()); +} + +pub fn init_testing_globals_with_config(config: InnerConfig) { #[cfg(debug_assertions)] { INIT_TESTING_GLOBALS.with(|init| { init.call_once(|| { let thread_name = std::thread::current().name().unwrap().to_string(); GlobalInstance::init_testing(&thread_name); - GlobalConfig::init(&InnerConfig::default(), &TEST_BUILD_INFO) - .expect("init global config"); + GlobalConfig::init(&config, &TEST_BUILD_INFO).expect("init global config"); LiteLicenseManager::init("default".to_string()).expect("init lite license manager"); SecurityPolicyCacheManager::init().unwrap(); }); @@ -165,8 +168,7 @@ pub fn init_testing_globals() { static INIT_GLOBALS: std::sync::Once = std::sync::Once::new(); INIT_GLOBALS.call_once(|| { GlobalInstance::init_production(); - GlobalConfig::init(&InnerConfig::default(), &TEST_BUILD_INFO) - .expect("init global config"); + GlobalConfig::init(&config, &TEST_BUILD_INFO).expect("init global config"); LiteLicenseManager::init("default".to_string()).expect("init lite license manager"); SecurityPolicyCacheManager::init().expect("init security policy cache manager"); }); @@ -262,6 +264,7 @@ impl DummyCatalog { #[derive(Debug, Clone)] struct FakeTable { table_info: TableInfo, + stream_source_table_info: Option, warehouse_distribution: bool, table_stats: Option, column_stats: HashMap, @@ -325,6 +328,10 @@ impl Table for FakeTable { &self.table_info } + fn stream_source_table_info(&self) -> Option<&TableInfo> { + self.stream_source_table_info.as_ref() + } + fn distribution_level(&self) -> DistributionLevel { if self.warehouse_distribution { DistributionLevel::Cluster @@ -346,6 +353,25 @@ impl Table for FakeTable { true } + async fn generate_changes_query( + &self, + ctx: Arc, + database_name: &str, + table_name: &str, + _with_options: &str, + ) -> Result { + if self.stream_source_table_info.is_none() { + return Err(ErrorCode::Unimplemented(format!( + "change query is not supported for test table {database_name}.{table_name}" + ))); + } + + let quote = ctx.get_settings().get_sql_dialect()?.default_ident_quote(); + Ok(format!( + "SELECT * FROM {quote}{database_name}{quote}.{quote}{table_name}{quote} AS _change_append$ffffffff" + )) + } + async fn table_statistics( &self, _ctx: Arc, @@ -878,6 +904,7 @@ impl LiteTableContext { catalog_info: self.default_catalog.info(), db_type: DatabaseType::NormalDB, }, + stream_source_table_info: None, warehouse_distribution, table_stats, column_stats, @@ -1094,6 +1121,72 @@ impl LiteTableContext { } fn register_replay_view(&self, view: &ReplayView) -> Result<()> { + self.register_view(view, Arc::new(TableSchema::new(vec![]))) + } + + pub async fn register_view_sql( + self: &Arc, + database: &str, + view_name: &str, + query: &str, + ) -> Result<()> { + let plan = self.bind_sql(query).await?; + let Plan::Query { bind_context, .. } = plan else { + return unsupported("lite sql harness view registration from non-query SQL"); + }; + if bind_context.columns.is_empty() { + return unsupported("lite sql harness view registration from empty query output"); + } + + self.register_view( + &ReplayView { + catalog: self.current_catalog.clone(), + database: database.to_string(), + view: view_name.to_string(), + query: query.to_string(), + }, + // Production View TableMeta stores the query but no schema. Keep the harness aligned + // so lineage tests cannot accidentally depend on persisted view fields. + Arc::new(TableSchema::default()), + ) + } + + pub async fn register_lineage_stream( + self: &Arc, + database: &str, + stream_name: &str, + source_table_name: &str, + ) -> Result<()> { + let source = self + .get_table(&self.current_catalog, database, source_table_name) + .await?; + let table_id = self.next_table_id.fetch_add(1, Ordering::Relaxed); + let table = Arc::new(FakeTable { + table_info: TableInfo { + ident: TableIdent::new(table_id, 0), + desc: format!("'{database}'.'{stream_name}'"), + name: stream_name.to_string(), + meta: TableMeta { + schema: source.schema(), + engine: "STREAM".to_string(), + ..Default::default() + }, + catalog_info: self.default_catalog.info(), + db_type: DatabaseType::NormalDB, + }, + stream_source_table_info: Some(source.get_table_info().clone()), + warehouse_distribution: false, + table_stats: None, + column_stats: HashMap::new(), + histograms: HashMap::new(), + top_n: HashMap::new(), + count_min_sketch: HashMap::new(), + }); + self.default_catalog.insert_table(database, table); + Ok(()) + } + + fn register_view(&self, view: &ReplayView, schema: Arc) -> Result<()> { let mut options = BTreeMap::new(); options.insert(LITE_VIEW_QUERY_KEY.to_string(), view.query.clone()); @@ -1104,7 +1197,7 @@ impl LiteTableContext { desc: format!("'{}'.'{}'", view.database, view.view), name: view.view.clone(), meta: TableMeta { - schema: Arc::new(TableSchema::new(vec![])), + schema, engine: LITE_VIEW_ENGINE.to_string(), options, ..Default::default() @@ -1112,6 +1205,7 @@ impl LiteTableContext { catalog_info: self.default_catalog.info(), db_type: DatabaseType::NormalDB, }, + stream_source_table_info: None, warehouse_distribution: false, table_stats: None, column_stats: HashMap::new(), diff --git a/src/query/sql/tests/it/planner.rs b/src/query/sql/tests/it/planner.rs index 66f5917b4d31b..a0f534da68755 100644 --- a/src/query/sql/tests/it/planner.rs +++ b/src/query/sql/tests/it/planner.rs @@ -40,6 +40,8 @@ use crate::framework::LiteTableContext; use crate::framework::golden::open_golden_file; use crate::framework::golden::write_case_title; +mod lineage; + struct LiteRunner(Arc); struct LiteReplayCaseSpec { diff --git a/src/query/sql/tests/it/planner/lineage.rs b/src/query/sql/tests/it/planner/lineage.rs new file mode 100644 index 0000000000000..24d0e47b93b35 --- /dev/null +++ b/src/query/sql/tests/it/planner/lineage.rs @@ -0,0 +1,577 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use databend_common_catalog::table_context::TableContextSettings; +use databend_common_catalog::table_context::TableContextTableAccess; +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; +use databend_common_expression::ColumnId; +use databend_common_meta_app::schema::CatalogType; +use databend_common_sql::LineageSource; +use databend_common_sql::LineageTarget; +use databend_common_sql::Planner; +use databend_common_sql::QueryLineage; +use databend_common_sql::QueryLineageColumn; +use databend_common_sql::QueryLineageColumnEdge; +use databend_common_sql::QueryLineageKind; +use databend_common_sql::QueryLineageRelation; +use databend_common_sql::QueryLineageRelationKind; +use databend_common_sql::plans::Plan; + +use crate::framework::LiteTableContext; + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_query_lineage_insert_select_from_sql() -> Result<()> { + let ctx = lineage_test_context().await?; + ctx.register_setup_sql("CREATE TABLE src(a INT, b INT, c INT)") + .await?; + ctx.register_setup_sql("CREATE TABLE dst(x INT)").await?; + + // INSERT INTO dst SELECT a + b AS x FROM src WHERE c > 0 + let sql = "INSERT INTO dst SELECT a + b AS x FROM src WHERE c > 0"; + let lineage = query_lineage_from_sql(&ctx, sql).await?; + let expected = + expected_table_query_lineage(QueryLineageKind::Dml, &ctx, "dst", "src", "x", &["a", "b"]) + .await?; + + assert_eq!(lineage, expected); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_query_lineage_stream_passes_through_to_source_table() -> Result<()> { + let ctx = lineage_test_context().await?; + ctx.register_setup_sql("CREATE TABLE src(a INT)").await?; + ctx.register_setup_sql("CREATE TABLE dst(x INT)").await?; + ctx.register_lineage_stream("default", "src_stream", "src") + .await?; + + let lineage = + query_lineage_from_sql(&ctx, "INSERT INTO dst SELECT a + 1 FROM src_stream").await?; + let expected = + expected_table_query_lineage(QueryLineageKind::Dml, &ctx, "dst", "src", "x", &["a"]) + .await?; + + assert_eq!(lineage, expected); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_query_lineage_aggregate_arguments_from_sql() -> Result<()> { + let ctx = lineage_test_context().await?; + ctx.register_setup_sql("CREATE TABLE src(a INT)").await?; + ctx.register_setup_sql("CREATE TABLE dst(all_count UInt64, sum_a Int64)") + .await?; + + let lineage = + query_lineage_from_sql(&ctx, "INSERT INTO dst SELECT count(*), sum(a) FROM src").await?; + + assert_lineage_sources(&lineage, QueryLineageKind::Dml, "dst", "all_count", &[]); + assert_lineage_sources(&lineage, QueryLineageKind::Dml, "dst", "sum_a", &["src.a"]); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_query_lineage_scalar_subquery_from_sql() -> Result<()> { + let ctx = lineage_test_context().await?; + ctx.register_setup_sql("CREATE TABLE left_src(a INT)") + .await?; + ctx.register_setup_sql("CREATE TABLE right_src(b INT, filter_flag INT)") + .await?; + ctx.register_setup_sql("CREATE TABLE dst(x INT)").await?; + + let lineage = query_lineage_from_sql( + &ctx, + "INSERT INTO dst SELECT a + (SELECT max(b) FROM right_src WHERE filter_flag > 0) FROM left_src", + ) + .await?; + + assert_lineage_sources(&lineage, QueryLineageKind::Dml, "dst", "x", &[ + "left_src.a", + "right_src.b", + ]); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_insert_lineage_does_not_pin_target_table_info() -> Result<()> { + let ctx = lineage_test_context().await?; + ctx.register_setup_sql("CREATE TABLE src(a INT)").await?; + ctx.register_setup_sql("CREATE TABLE dst(a INT)").await?; + + let plan = ctx.bind_sql("INSERT INTO dst SELECT a FROM src").await?; + let Plan::Insert(plan) = plan else { + return Err(ErrorCode::Internal("expected insert plan")); + }; + // Lineage capture needs the target table id, but ordinary INSERT must keep + // resolving its target by name at execution time. Populating `table_info` + // would pin execution to the table object seen during binding, so lineage + // stores the id in its dedicated field instead. + assert!(plan.table_info.is_none()); + assert!(plan.lineage_target_table_id.is_some()); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_query_lineage_insert_select_with_cte_from_sql() -> Result<()> { + let ctx = lineage_test_context().await?; + ctx.register_setup_sql("CREATE TABLE src(a INT, b INT, c INT)") + .await?; + ctx.register_setup_sql("CREATE TABLE dst(x INT)").await?; + + // INSERT INTO dst WITH q AS (SELECT a + b AS x, c FROM src) SELECT x FROM q WHERE c > 0 + let sql = + "INSERT INTO dst WITH q AS (SELECT a + b AS x, c FROM src) SELECT x FROM q WHERE c > 0"; + let lineage = query_lineage_from_sql(&ctx, sql).await?; + let expected = + expected_table_query_lineage(QueryLineageKind::Dml, &ctx, "dst", "src", "x", &["a", "b"]) + .await?; + + assert_eq!(lineage, expected); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_query_lineage_insert_select_with_auto_materialized_cte_from_sql() -> Result<()> { + let ctx = lineage_test_context().await?; + ctx.get_settings().set_enable_auto_materialize_cte(1)?; + ctx.register_setup_sql("CREATE TABLE src(a INT, b INT, c INT)") + .await?; + ctx.register_setup_sql("CREATE TABLE dst(x INT)").await?; + + // INSERT INTO dst WITH q AS (SELECT a, b, c FROM src) SELECT q1.a + q2.b FROM q q1 JOIN q q2 ON q1.c = q2.c + let sql = "INSERT INTO dst WITH q AS (SELECT a, b, c FROM src) SELECT q1.a + q2.b FROM q q1 JOIN q q2 ON q1.c = q2.c"; + let lineage = query_lineage_from_sql(&ctx, sql).await?; + + assert_lineage_sources(&lineage, QueryLineageKind::Dml, "dst", "x", &[ + "src.a", "src.b", + ]); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_query_lineage_join_and_exists_filter_columns_are_excluded_from_sql() -> Result<()> { + let ctx = lineage_test_context().await?; + ctx.register_setup_sql("CREATE TABLE left_src(id INT, k INT, a INT)") + .await?; + ctx.register_setup_sql("CREATE TABLE right_src(id INT, b INT)") + .await?; + ctx.register_setup_sql("CREATE TABLE filter_src(k INT, flag INT)") + .await?; + ctx.register_setup_sql("CREATE TABLE dst(x INT)").await?; + + // INSERT INTO dst SELECT l.a + r.b FROM left_src l JOIN right_src r ON l.id = r.id + // WHERE EXISTS (SELECT 1 FROM filter_src f WHERE f.k = l.k AND f.flag > 0) + let sql = "INSERT INTO dst SELECT l.a + r.b FROM left_src l JOIN right_src r ON l.id = r.id WHERE EXISTS (SELECT 1 FROM filter_src f WHERE f.k = l.k AND f.flag > 0)"; + let lineage = query_lineage_from_sql(&ctx, sql).await?; + + assert_lineage_sources(&lineage, QueryLineageKind::Dml, "dst", "x", &[ + "left_src.a", + "right_src.b", + ]); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_query_lineage_join_using_prefers_left_column_from_sql() -> Result<()> { + let ctx = lineage_test_context().await?; + ctx.register_setup_sql("CREATE TABLE left_src(id INT, a INT)") + .await?; + ctx.register_setup_sql("CREATE TABLE right_src(id INT, b INT)") + .await?; + ctx.register_setup_sql("CREATE TABLE dst(x INT)").await?; + + // INSERT INTO dst SELECT id FROM left_src JOIN right_src USING(id) + let sql = "INSERT INTO dst SELECT id FROM left_src JOIN right_src USING(id)"; + let lineage = query_lineage_from_sql(&ctx, sql).await?; + + assert_lineage_sources(&lineage, QueryLineageKind::Dml, "dst", "x", &[ + "left_src.id", + ]); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_query_lineage_create_view_from_sql() -> Result<()> { + let ctx = lineage_test_context().await?; + ctx.register_setup_sql("CREATE TABLE src(a INT, b INT, c INT)") + .await?; + + // CREATE VIEW v(vx) AS SELECT a + b FROM src WHERE c > 0 + let sql = "CREATE VIEW v(vx) AS SELECT a + b FROM src WHERE c > 0"; + let mut plan = ctx.bind_sql(sql).await?; + let query_plan = ctx.bind_sql("SELECT a + b FROM src WHERE c > 0").await?; + let Plan::CreateView(create_view) = &mut plan else { + return Err(ErrorCode::Internal("expected create view plan")); + }; + create_view.query_plan = Some(Box::new(query_plan)); + let lineage = plan + .query_lineage()? + .ok_or_else(|| ErrorCode::Internal("missing create view lineage"))?; + let expected = expected_query_lineage( + QueryLineageKind::CreateView, + relation("v", QueryLineageRelationKind::View, None), + table_relation(&ctx, "src").await?, + column("vx", 0), + vec![ + table_column(&ctx, "src", "a").await?, + table_column(&ctx, "src", "b").await?, + ], + ); + + assert_eq!(lineage, expected); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_create_view_plan_omits_lineage_query_by_default() -> Result<()> { + let ctx = LiteTableContext::create().await?; + let plan = ctx + .bind_sql("CREATE VIEW v AS SELECT * FROM missing_table") + .await?; + + let Plan::CreateView(plan) = plan else { + return Err(ErrorCode::Internal("expected create view plan")); + }; + assert!(plan.query_plan.is_none()); + Ok(()) +} + +#[test] +fn test_query_lineage_insert_select_from_view_stops_at_view_from_sql() -> Result<()> { + std::thread::Builder::new() + .name("lineage_view_boundary_sql".to_string()) + .spawn(|| { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .enable_all() + .build() + .map_err(|err| ErrorCode::Internal(err.to_string()))? + .block_on(async { + let ctx = lineage_test_context().await?; + ctx.register_setup_sql("CREATE TABLE src(a INT, b INT)") + .await?; + ctx.register_setup_sql("CREATE TABLE dst(x INT)").await?; + ctx.register_view_sql("default", "v", "SELECT a + b AS vx FROM src") + .await?; + + // Keep an expression above the view output so optimized-plan + // extraction must preserve the view lineage boundary. + let lineage = + query_lineage_from_sql(&ctx, "INSERT INTO dst SELECT vx + 1 FROM v") + .await?; + let expected = expected_query_lineage( + QueryLineageKind::Dml, + table_relation(&ctx, "dst").await?, + view_relation(&ctx, "v").await?, + table_column(&ctx, "dst", "x").await?, + vec![column("vx", 0)], + ); + + assert_eq!(lineage, expected); + Ok(()) + }) + }) + .map_err(|err| ErrorCode::Internal(err.to_string()))? + .join() + .map_err(|_| ErrorCode::Internal("lineage view boundary test panicked"))? +} + +#[test] +fn test_query_lineage_ctas_from_view_stops_at_view_from_sql() -> Result<()> { + std::thread::Builder::new() + .name("lineage_ctas_view_boundary_sql".to_string()) + .spawn(|| { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .enable_all() + .build() + .map_err(|err| ErrorCode::Internal(err.to_string()))? + .block_on(async { + let ctx = lineage_test_context().await?; + ctx.register_setup_sql("CREATE TABLE src(a INT, b INT)") + .await?; + ctx.register_view_sql("default", "v", "SELECT a + b AS vx FROM src") + .await?; + + let lineage = query_lineage_from_sql( + &ctx, + "CREATE TABLE dst ENGINE=NULL AS SELECT vx AS x FROM v", + ) + .await?; + let expected = expected_query_lineage( + QueryLineageKind::Ctas, + relation("dst", QueryLineageRelationKind::Table, None), + view_relation(&ctx, "v").await?, + column("x", 0), + vec![column("vx", 0)], + ); + + assert_eq!(lineage, expected); + Ok(()) + }) + }) + .map_err(|err| ErrorCode::Internal(err.to_string()))? + .join() + .map_err(|_| ErrorCode::Internal("CTAS view boundary test panicked"))? +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_query_lineage_replace_into_from_sql() -> Result<()> { + let ctx = lineage_test_context().await?; + ctx.register_setup_sql("CREATE TABLE src(id INT, a INT, b INT, c INT)") + .await?; + ctx.register_setup_sql("CREATE TABLE dst(id INT, x INT)") + .await?; + + // REPLACE INTO dst(id, x) ON(id) SELECT id, a + b FROM src WHERE c > 0 + let sql = "REPLACE INTO dst(id, x) ON(id) SELECT id, a + b FROM src WHERE c > 0"; + let lineage = query_lineage_from_sql(&ctx, sql).await?; + + assert_lineage_sources(&lineage, QueryLineageKind::Dml, "dst", "id", &["src.id"]); + assert_lineage_sources(&lineage, QueryLineageKind::Dml, "dst", "x", &[ + "src.a", "src.b", + ]); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_query_lineage_multi_insert_from_sql() -> Result<()> { + let ctx = lineage_test_context().await?; + ctx.register_setup_sql("CREATE TABLE src(a INT, b INT, c INT)") + .await?; + ctx.register_setup_sql("CREATE TABLE dst1(x INT, y INT)") + .await?; + ctx.register_setup_sql("CREATE TABLE dst2(x INT, y INT)") + .await?; + + // INSERT ALL INTO dst1 VALUES(a, b) INTO dst2(x, y) VALUES(b, c) SELECT a, b, c FROM src + let sql = + "INSERT ALL INTO dst1 VALUES(a, b) INTO dst2(x, y) VALUES(b, c) SELECT a, b, c FROM src"; + let lineage = query_lineage_from_sql(&ctx, sql).await?; + + assert_lineage_sources(&lineage, QueryLineageKind::Dml, "dst1", "x", &["src.a"]); + assert_lineage_sources(&lineage, QueryLineageKind::Dml, "dst1", "y", &["src.b"]); + assert_lineage_sources(&lineage, QueryLineageKind::Dml, "dst2", "x", &["src.b"]); + assert_lineage_sources(&lineage, QueryLineageKind::Dml, "dst2", "y", &["src.c"]); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_query_lineage_update_from_sql() -> Result<()> { + let ctx = lineage_test_context().await?; + ctx.register_setup_sql("CREATE TABLE src(id INT, a INT, b INT, c INT)") + .await?; + ctx.register_setup_sql("CREATE TABLE dst(id INT, x INT)") + .await?; + + // UPDATE dst SET x = src.a + src.b FROM src WHERE dst.id = src.id AND src.c > 0 + let sql = "UPDATE dst SET x = src.a + src.b FROM src WHERE dst.id = src.id AND src.c > 0"; + let lineage = query_lineage_from_bound_sql(&ctx, sql).await?; + + assert_lineage_sources(&lineage, QueryLineageKind::Dml, "dst", "x", &[ + "src.a", "src.b", + ]); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_query_lineage_merge_multiple_when_from_sql() -> Result<()> { + let ctx = lineage_test_context().await?; + ctx.register_setup_sql("CREATE TABLE src(id INT, a INT, b INT, c INT, d INT)") + .await?; + ctx.register_setup_sql("CREATE TABLE dst(id INT, x INT, y INT)") + .await?; + + // MERGE INTO dst USING src ON dst.id = src.id + // WHEN MATCHED AND src.c > 0 THEN UPDATE SET x = src.a + // WHEN MATCHED AND src.d > 0 THEN UPDATE SET y = src.b + // WHEN NOT MATCHED THEN INSERT (id, x, y) VALUES (src.id, src.a, src.b) + let sql = "MERGE INTO dst USING src ON dst.id = src.id WHEN MATCHED AND src.c > 0 THEN UPDATE SET x = src.a WHEN MATCHED AND src.d > 0 THEN UPDATE SET y = src.b WHEN NOT MATCHED THEN INSERT (id, x, y) VALUES (src.id, src.a, src.b)"; + let lineage = query_lineage_from_bound_sql(&ctx, sql).await?; + + assert_lineage_sources(&lineage, QueryLineageKind::Dml, "dst", "id", &["src.id"]); + assert_lineage_sources(&lineage, QueryLineageKind::Dml, "dst", "x", &["src.a"]); + assert_lineage_sources(&lineage, QueryLineageKind::Dml, "dst", "y", &["src.b"]); + Ok(()) +} + +async fn lineage_test_context() -> Result> { + LiteTableContext::create().await +} + +async fn query_lineage_from_sql(ctx: &Arc, sql: &str) -> Result { + let mut planner = Planner::new(ctx.clone()); + let (plan, _) = planner.plan_sql(sql).await?; + plan.query_lineage()? + .ok_or_else(|| ErrorCode::Internal(format!("missing query lineage for SQL: {sql}"))) +} + +async fn query_lineage_from_bound_sql( + ctx: &Arc, + sql: &str, +) -> Result { + let plan = ctx.bind_sql(sql).await?; + plan.query_lineage()? + .ok_or_else(|| ErrorCode::Internal(format!("missing query lineage for SQL: {sql}"))) +} + +fn assert_lineage_sources( + lineage: &QueryLineage, + kind: QueryLineageKind, + target_table: &str, + target_column: &str, + expected: &[&str], +) { + assert_eq!(lineage.kind, kind, "unexpected lineage kind: {lineage:?}"); + + let mut actual = lineage + .targets + .iter() + .find(|target| target.relation.name == target_table) + .unwrap_or_else(|| panic!("missing target table {target_table}: {lineage:?}")) + .sources + .iter() + .flat_map(|from_relation| { + from_relation + .columns + .iter() + .filter(|edge| edge.target.name == target_column) + .map(|edge| format!("{}.{}", from_relation.relation.name, edge.source.name)) + }) + .collect::>(); + actual.sort(); + actual.dedup(); + + let mut expected = expected + .iter() + .map(|source| source.to_string()) + .collect::>(); + expected.sort(); + + assert_eq!( + actual, expected, + "unexpected lineage for {target_table}.{target_column}: {lineage:?}" + ); +} + +async fn table_id(ctx: &Arc, table: &str) -> Result { + Ok(ctx + .get_table("default", "default", table) + .await? + .get_table_info() + .ident + .table_id) +} + +async fn column_id(ctx: &Arc, table: &str, column: &str) -> Result { + ctx.get_table("default", "default", table) + .await? + .schema() + .column_id_of(column) +} + +async fn expected_table_query_lineage( + kind: QueryLineageKind, + ctx: &Arc, + to_table: &str, + from_table: &str, + to_column: &str, + from_columns: &[&str], +) -> Result { + let mut sources = Vec::with_capacity(from_columns.len()); + for from_column in from_columns { + sources.push(table_column(ctx, from_table, from_column).await?); + } + + Ok(expected_query_lineage( + kind, + table_relation(ctx, to_table).await?, + table_relation(ctx, from_table).await?, + table_column(ctx, to_table, to_column).await?, + sources, + )) +} + +async fn table_relation(ctx: &Arc, table: &str) -> Result { + Ok(relation( + table, + QueryLineageRelationKind::Table, + Some(table_id(ctx, table).await?), + )) +} + +async fn view_relation(ctx: &Arc, view: &str) -> Result { + Ok(relation( + view, + QueryLineageRelationKind::View, + Some(table_id(ctx, view).await?), + )) +} + +async fn table_column( + ctx: &Arc, + table: &str, + column_name: &str, +) -> Result { + Ok(column( + column_name, + column_id(ctx, table, column_name).await?, + )) +} + +fn expected_query_lineage( + kind: QueryLineageKind, + target: QueryLineageRelation, + source: QueryLineageRelation, + target_column: QueryLineageColumn, + source_columns: Vec, +) -> QueryLineage { + QueryLineage { + kind, + targets: vec![LineageTarget { + relation: target, + sources: vec![LineageSource { + relation: source, + columns: source_columns + .into_iter() + .map(|source| QueryLineageColumnEdge { + source, + target: target_column.clone(), + }) + .collect(), + }], + }], + } +} + +fn relation(name: &str, kind: QueryLineageRelationKind, id: Option) -> QueryLineageRelation { + QueryLineageRelation { + catalog: "default".to_string(), + database: "default".to_string(), + name: name.to_string(), + id, + catalog_type: Some(CatalogType::Default), + kind, + } +} + +fn column(name: &str, id: ColumnId) -> QueryLineageColumn { + QueryLineageColumn { + name: name.to_string(), + id, + } +} diff --git a/src/query/storages/stream/src/stream_table.rs b/src/query/storages/stream/src/stream_table.rs index 46edaa93e44d7..3eed40417c461 100644 --- a/src/query/storages/stream/src/stream_table.rs +++ b/src/query/storages/stream/src/stream_table.rs @@ -395,6 +395,12 @@ impl Table for StreamTable { &self.info } + fn stream_source_table_info(&self) -> Option<&TableInfo> { + self.source_table + .as_ref() + .map(|table| table.get_table_info()) + } + fn supported_internal_column(&self, column_id: ColumnId) -> bool { (BASE_BLOCK_IDS_COLUMN_ID..=BASE_ROW_ID_COLUMN_ID).contains(&column_id) }