diff --git a/.github/actions/test_logs/action.yml b/.github/actions/test_logs/action.yml index dd1843af8f3f5..eacc5cea28055 100644 --- a/.github/actions/test_logs/action.yml +++ b/.github/actions/test_logs/action.yml @@ -4,6 +4,8 @@ runs: using: "composite" steps: - uses: ./.github/actions/setup_test + with: + artifacts: sqllogictests,meta,query - name: Install lsof shell: bash diff --git a/src/common/tracing/src/config.rs b/src/common/tracing/src/config.rs index df1588b0c62fd..d0e50257b8d9a 100644 --- a/src/common/tracing/src/config.rs +++ b/src/common/tracing/src/config.rs @@ -398,7 +398,8 @@ pub struct HistoryConfig { #[derive(Clone, Debug, PartialEq, Eq, serde::Serialize)] pub struct HistoryTableConfig { pub table_name: String, - pub retention: usize, + /// `None` is resolved to the table-specific policy by `init_history_tables`. + pub retention: Option, pub invisible: bool, } @@ -415,7 +416,12 @@ impl Display for HistoryConfig { self.retention_interval, self.tables .iter() - .map(|f| format!("{}({} hours)", f.table_name.clone(), f.retention)) + .map(|f| match f.retention { + Some(retention) => { + format!("{}({} hours)", f.table_name, retention) + } + None => format!("{}(default)", f.table_name), + }) .join(", "), self.storage_params .as_ref() @@ -445,13 +451,21 @@ impl Default for HistoryTableConfig { fn default() -> Self { Self { table_name: "".to_string(), - retention: 168, + retention: None, invisible: false, } } } impl HistoryConfig { + pub fn is_table_enabled(&self, table_name: &str) -> bool { + self.on + && self + .tables + .iter() + .any(|table| table.table_name.eq_ignore_ascii_case(table_name)) + } + pub fn is_invisible(&self, table_name: &str) -> bool { self.tables .iter() diff --git a/src/common/tracing/src/predefined_tables/history_tables.rs b/src/common/tracing/src/predefined_tables/history_tables.rs index 72cf7e529ebe8..cd5d6ed5d192c 100644 --- a/src/common/tracing/src/predefined_tables/history_tables.rs +++ b/src/common/tracing/src/predefined_tables/history_tables.rs @@ -22,39 +22,56 @@ use databend_common_exception::Result; use crate::HistoryConfig; const TABLES_TOML: &str = include_str!("./history_tables.toml"); +const DEFAULT_RETENTION_HOURS: u64 = 24 * 7; +const LINEAGE_UNRESOLVED_TABLE: &str = "lineage_unresolved"; #[derive(Debug)] pub struct HistoryTable { pub name: String, pub create: String, - pub transform: String, - pub delete: String, + pub transforms: Vec, + pub delete: Option, } impl HistoryTable { - pub fn create(predefined: PredefinedTable, retention: u64) -> Self { + pub fn create(predefined: PredefinedTable, retention: Option) -> Self { + let transforms = predefined.transforms(); HistoryTable { name: predefined.name, create: predefined.create, - transform: predefined.transform, - delete: predefined - .delete - .replace("{retention_hours}", &retention.to_string()), + transforms, + delete: retention.map(|retention| { + predefined + .delete + .replace("{retention_hours}", &retention.to_string()) + }), } } - pub fn assemble_log_history_transform(&self, stage_name: &str, batch_number: u64) -> String { - let mut transform = self.transform.clone(); - transform = transform.replace("{stage_name}", stage_name); - transform = transform.replace("{batch_number}", &batch_number.to_string()); - transform + pub fn assemble_log_history_transforms( + &self, + stage_name: &str, + batch_number: u64, + ) -> Vec { + self.transforms + .iter() + .map(|transform| { + transform + .replace("{stage_name}", stage_name) + .replace("{batch_number}", &batch_number.to_string()) + }) + .collect() } - pub fn assemble_normal_transform(&self, begin: u64, end: u64) -> String { - let mut transform = self.transform.clone(); - transform = transform.replace("{batch_begin}", &begin.to_string()); - transform = transform.replace("{batch_end}", &end.to_string()); - transform + pub fn assemble_normal_transforms(&self, begin: u64, end: u64) -> Vec { + self.transforms + .iter() + .map(|transform| { + transform + .replace("{batch_begin}", &begin.to_string()) + .replace("{batch_end}", &end.to_string()) + }) + .collect() } } @@ -69,9 +86,19 @@ pub struct PredefinedTable { pub target: String, pub create: String, pub transform: String, + #[serde(default)] + pub additional_transforms: Vec, pub delete: String, } +impl PredefinedTable { + fn transforms(&self) -> Vec { + std::iter::once(self.transform.clone()) + .chain(self.additional_transforms.iter().cloned()) + .collect() + } +} + pub fn init_history_tables(cfg: &HistoryConfig) -> Result>> { let predefined_tables: PredefinedTables = toml::from_str(TABLES_TOML).expect("Failed to parse toml"); @@ -93,11 +120,12 @@ pub fn init_history_tables(cfg: &HistoryConfig) -> Result> user_defined_log_history = true; } if let Some(predefined_table) = predefined_map.remove(&enable_table.table_name) { - let retention = enable_table.retention; - history_tables.push(Arc::new(HistoryTable::create( - predefined_table, - retention as u64, - ))); + let retention = match enable_table.retention { + Some(retention) => Some(retention as u64), + None if enable_table.table_name == LINEAGE_UNRESOLVED_TABLE => None, + None => Some(DEFAULT_RETENTION_HOURS), + }; + history_tables.push(Arc::new(HistoryTable::create(predefined_table, retention))); } else { return Err(ErrorCode::InvalidConfig(format!( "Invalid history table name {}", @@ -108,7 +136,7 @@ pub fn init_history_tables(cfg: &HistoryConfig) -> Result> if !user_defined_log_history { history_tables.push(Arc::new(HistoryTable::create( predefined_map.remove("log_history").unwrap(), - 24 * 7, + Some(DEFAULT_RETENTION_HOURS), ))); } Ok(history_tables) @@ -135,3 +163,131 @@ pub fn get_all_history_table_names() -> Vec { .map(|t| t.name) .collect() } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_lineage_unresolved_uses_hash_pattern_merge() { + let tables: PredefinedTables = toml::from_str(TABLES_TOML).unwrap(); + let lineage = tables + .tables + .iter() + .find(|table| table.name == "lineage_unresolved") + .unwrap(); + + assert!( + lineage + .create + .contains("column_lineage_hash STRING NOT NULL") + ); + assert!(lineage.create.contains("source_to_target_columns MAP")); + assert!(lineage.create.contains("target_to_source_columns MAP")); + assert!( + lineage + .transform + .contains("MERGE INTO system_history.lineage_unresolved") + ); + assert!(lineage.create.contains("source_catalog_type STRING")); + assert!(lineage.create.contains("target_catalog_type STRING")); + assert!(lineage.transform.contains("AS source_catalog_type")); + assert!(lineage.delete.contains("lineage_kind = 'DML'")); + assert!( + lineage + .transform + .contains("target.column_lineage_hash = source.column_lineage_hash") + ); + assert!(lineage.transform.contains( + "PARTITION BY m['source']['lineage_key']::STRING, m['target']['lineage_key']::STRING, m['lineage_kind']::STRING, m['column_lineage_hash']::STRING" + )); + assert!(lineage.transform.contains( + "target.source_lineage_key = source.source_lineage_key AND target.target_lineage_key = source.target_lineage_key AND target.lineage_kind = source.lineage_kind" + )); + assert!( + lineage + .transform + .contains("WHEN MATCHED THEN UPDATE * WHEN NOT MATCHED THEN INSERT *") + ); + assert!( + lineage + .transform + .contains("coalesce(m['operation']::STRING, 'UPSERT_EDGE') = 'UPSERT_EDGE'") + ); + assert_eq!(lineage.additional_transforms.len(), 1); + assert!(lineage.additional_transforms[0].contains("DELETE_OBJECT")); + assert!(lineage.additional_transforms[0].contains("source_id IN")); + assert!(lineage.additional_transforms[0].contains("target_id IN")); + } + + #[test] + fn test_lineage_same_batch_tombstone_replay_order() { + let tables: PredefinedTables = toml::from_str(TABLES_TOML).unwrap(); + let lineage = tables + .tables + .into_iter() + .find(|table| table.name == LINEAGE_UNRESOLVED_TABLE) + .unwrap(); + let history = HistoryTable::create(lineage, None); + + // A same-batch tombstone must run after edge upserts. Replaying a batch repeats the same + // order, so the MERGE cannot resurrect an edge removed by the tombstone phase. + let first_attempt = history.assemble_normal_transforms(17, 23); + let replay = history.assemble_normal_transforms(17, 23); + assert_eq!(first_attempt, replay); + assert_eq!(first_attempt.len(), 2); + assert!(first_attempt[0].contains("MERGE INTO system_history.lineage_unresolved")); + assert!(first_attempt[0].contains("= 'UPSERT_EDGE'")); + assert!(first_attempt[1].starts_with("DELETE FROM system_history.lineage_unresolved")); + assert!(first_attempt[1].contains("= 'DELETE_OBJECT'")); + for phase in first_attempt { + assert!(phase.contains("batch_number >= 17")); + assert!(phase.contains("batch_number < 23")); + } + } + + #[test] + fn test_lineage_retention_is_optional() { + let config = |table_name: &str, retention| HistoryConfig { + tables: vec![crate::HistoryTableConfig { + table_name: table_name.to_string(), + retention, + invisible: false, + }], + ..Default::default() + }; + + let lineage = init_history_tables(&config(LINEAGE_UNRESOLVED_TABLE, None)).unwrap(); + assert!( + lineage + .iter() + .find(|table| table.name == LINEAGE_UNRESOLVED_TABLE) + .unwrap() + .delete + .is_none() + ); + + let lineage = init_history_tables(&config(LINEAGE_UNRESOLVED_TABLE, Some(24))).unwrap(); + let delete = lineage + .iter() + .find(|table| table.name == LINEAGE_UNRESOLVED_TABLE) + .unwrap() + .delete + .as_ref() + .unwrap(); + assert!(delete.contains("lineage_kind = 'DML'")); + assert!(delete.contains("subtract_hours(NOW(), 24)")); + + let query = init_history_tables(&config("query_history", None)).unwrap(); + assert!( + query + .iter() + .find(|table| table.name == "query_history") + .unwrap() + .delete + .as_ref() + .unwrap() + .contains("subtract_hours(NOW(), 168)") + ); + } +} diff --git a/src/common/tracing/src/predefined_tables/history_tables.toml b/src/common/tracing/src/predefined_tables/history_tables.toml index 927fac7cd0fcb..8415e5e8c3a12 100644 --- a/src/common/tracing/src/predefined_tables/history_tables.toml +++ b/src/common/tracing/src/predefined_tables/history_tables.toml @@ -33,3 +33,11 @@ target = "databend::log::access" create = "CREATE TABLE IF NOT EXISTS system_history.access_history (query_id VARCHAR NULL, query_start TIMESTAMP NULL, user_name STRING NULL, base_objects_accessed VARIANT NULL, direct_objects_accessed VARIANT NULL, objects_modified VARIANT NULL, object_modified_by_ddl VARIANT NULL)" transform = "settings (timezone='Etc/UTC') INSERT INTO system_history.access_history FROM (SELECT m['query_id'], m['query_start'], m['user_name'], m['base_objects_accessed'], m['direct_objects_accessed'], m['objects_modified'], m['object_modified_by_ddl'] FROM (SELECT parse_json(message) as m FROM system_history.log_history WHERE target='databend::log::access' AND batch_number >= {batch_begin} AND batch_number < {batch_end}))" delete = "DELETE FROM system_history.access_history WHERE query_start < subtract_hours(NOW(), {retention_hours})" + +[[tables]] +name = "lineage_unresolved" +target = "databend::log::lineage" +create = "CREATE TABLE IF NOT EXISTS system_history.lineage_unresolved (query_id STRING NULL, event_time TIMESTAMP NULL, query_kind STRING NULL, lineage_kind STRING NULL, column_lineage_hash STRING NOT NULL, source_lineage_key STRING NULL, source_address_kind STRING NULL, source_catalog_type STRING NULL, source_object_type STRING NULL, source_catalog STRING NULL, source_database STRING NULL, source_name STRING NULL, source_id UInt64 NULL, target_lineage_key STRING NULL, target_address_kind STRING NULL, target_catalog_type STRING NULL, target_object_type STRING NULL, target_catalog STRING NULL, target_database STRING NULL, target_name STRING NULL, target_id UInt64 NULL, source_to_target_columns MAP(STRING, ARRAY(STRING)) NOT NULL, target_to_source_columns MAP(STRING, ARRAY(STRING)) NOT NULL) CLUSTER BY LINEAR(event_time, target_lineage_key, source_lineage_key)" +transform = "settings (timezone='Etc/UTC') MERGE INTO system_history.lineage_unresolved AS target USING (SELECT m['query_id']::STRING AS query_id, to_timestamp(m['event_time']::Int64) AS event_time, m['query_kind']::STRING AS query_kind, m['lineage_kind']::STRING AS lineage_kind, m['column_lineage_hash']::STRING AS column_lineage_hash, m['source']['lineage_key']::STRING AS source_lineage_key, m['source']['address_kind']::STRING AS source_address_kind, m['source']['catalog_type']::STRING AS source_catalog_type, m['source']['object_type']::STRING AS source_object_type, m['source']['catalog']::STRING AS source_catalog, m['source']['database']::STRING AS source_database, m['source']['name']::STRING AS source_name, m['source']['id']::UInt64 AS source_id, m['target']['lineage_key']::STRING AS target_lineage_key, m['target']['address_kind']::STRING AS target_address_kind, m['target']['catalog_type']::STRING AS target_catalog_type, m['target']['object_type']::STRING AS target_object_type, m['target']['catalog']::STRING AS target_catalog, m['target']['database']::STRING AS target_database, m['target']['name']::STRING AS target_name, m['target']['id']::UInt64 AS target_id, m['source_to_target_columns']::MAP(STRING, ARRAY(STRING)) AS source_to_target_columns, m['target_to_source_columns']::MAP(STRING, ARRAY(STRING)) AS target_to_source_columns FROM (SELECT parse_json(message) AS m FROM system_history.log_history WHERE target='databend::log::lineage' AND batch_number >= {batch_begin} AND batch_number < {batch_end}) WHERE coalesce(m['operation']::STRING, 'UPSERT_EDGE') = 'UPSERT_EDGE' QUALIFY row_number() OVER (PARTITION BY m['source']['lineage_key']::STRING, m['target']['lineage_key']::STRING, m['lineage_kind']::STRING, m['column_lineage_hash']::STRING ORDER BY m['event_time']::Int64 DESC, m['query_id']::STRING DESC) = 1) AS source ON target.source_lineage_key = source.source_lineage_key AND target.target_lineage_key = source.target_lineage_key AND target.lineage_kind = source.lineage_kind AND target.column_lineage_hash = source.column_lineage_hash WHEN MATCHED THEN UPDATE * WHEN NOT MATCHED THEN INSERT *;" +additional_transforms = ["DELETE FROM system_history.lineage_unresolved WHERE (source_address_kind = 'ID' AND source_id IN (SELECT DISTINCT m['object_id']::UInt64 FROM (SELECT parse_json(message) AS m FROM system_history.log_history WHERE target='databend::log::lineage' AND batch_number >= {batch_begin} AND batch_number < {batch_end}) WHERE m['operation']::STRING = 'DELETE_OBJECT')) OR (target_address_kind = 'ID' AND target_id IN (SELECT DISTINCT m['object_id']::UInt64 FROM (SELECT parse_json(message) AS m FROM system_history.log_history WHERE target='databend::log::lineage' AND batch_number >= {batch_begin} AND batch_number < {batch_end}) WHERE m['operation']::STRING = 'DELETE_OBJECT'))"] +delete = "DELETE FROM system_history.lineage_unresolved WHERE lineage_kind = 'DML' AND event_time < subtract_hours(NOW(), {retention_hours})" 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/config/src/config.rs b/src/query/config/src/config.rs index 792745c586fc1..8d85cc27d2342 100644 --- a/src/query/config/src/config.rs +++ b/src/query/config/src/config.rs @@ -2724,9 +2724,9 @@ pub struct HistoryLogTableConfig { /// Specifies the history table name pub table_name: String, - /// The retention period (in hours) for history logs. - /// Data older than this period will be deleted during retention tasks. - pub retention: usize, + /// The optional retention period (in hours) for history logs. + /// Unspecified tables use the default retention, except lineage, which is retained permanently. + pub retention: Option, /// Whether this history table is invisible for querying. /// Default is false. diff --git a/src/query/service/src/history_tables/global_history_log.rs b/src/query/service/src/history_tables/global_history_log.rs index 0e8b10da02f59..f6ad979e33ced 100644 --- a/src/query/service/src/history_tables/global_history_log.rs +++ b/src/query/service/src/history_tables/global_history_log.rs @@ -297,8 +297,8 @@ impl GlobalHistoryLog { .get_u64_from_meta(&format!("{}/{}/batch_number", meta_key, table.name)) .await? .unwrap_or(0); - let sql = if table.name == "log_history" { - table.assemble_log_history_transform(&self.stage_name, batch_number_begin) + let sqls = if table.name == "log_history" { + table.assemble_log_history_transforms(&self.stage_name, batch_number_begin) } else { batch_number_end = self .meta_handle @@ -308,9 +308,13 @@ impl GlobalHistoryLog { if batch_number_begin >= batch_number_end { return Ok(()); } - table.assemble_normal_transform(batch_number_begin, batch_number_end) + table.assemble_normal_transforms(batch_number_begin, batch_number_end) }; - self.execute_sql(&sql).await?; + // Advance the batch checkpoint only after every phase succeeds. Tables that configure + // additional phases must keep them replay-safe so a partial failure can retry the batch. + for sql in sqls { + self.execute_sql(&sql).await?; + } if table.name == "log_history" { self.meta_handle .set_u64_to_meta( @@ -348,6 +352,9 @@ impl GlobalHistoryLog { } pub async fn clean(&self, table: &HistoryTable, meta_key: &str) -> Result { + let Some(sql) = table.delete.as_ref() else { + return Ok(false); + }; let got_permit = self .meta_handle .check_should_perform_clean( @@ -357,7 +364,6 @@ impl GlobalHistoryLog { .await?; if got_permit { let start = Instant::now(); - let sql = &table.delete; self.execute_sql(sql).await?; let context = self.create_context().await?; let delete_elapsed = start.elapsed().as_secs(); diff --git a/src/query/service/src/interpreters/common/finish_hook.rs b/src/query/service/src/interpreters/common/finish_hook.rs index 33907a04a9839..224c9558a4b8c 100644 --- a/src/query/service/src/interpreters/common/finish_hook.rs +++ b/src/query/service/src/interpreters/common/finish_hook.rs @@ -19,6 +19,7 @@ use databend_common_exception::Result; use databend_common_pipeline::core::ExecutionInfo; use crate::interpreters::common::log_query_finished; +use crate::interpreters::common::log_query_lineage; use crate::interpreters::hook::vacuum_hook::hook_clear_m_cte_temp_table; use crate::interpreters::hook::vacuum_hook::hook_disk_temp_dir; use crate::interpreters::hook::vacuum_hook::hook_vacuum_temp_files; @@ -102,6 +103,9 @@ impl QueryFinishHooks { }; if self.log_finished { log_query_finished(&ctx, info.res.clone().err()); + if info.res.is_ok() { + log_query_lineage(&ctx); + } } info.res.clone().and(hooks_res) } diff --git a/src/query/service/src/interpreters/common/lineage_log.rs b/src/query/service/src/interpreters/common/lineage_log.rs new file mode 100644 index 0000000000000..efc48d0912521 --- /dev/null +++ b/src/query/service/src/interpreters/common/lineage_log.rs @@ -0,0 +1,467 @@ +// 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::sync::Arc; +use std::time::SystemTime; + +use databend_common_catalog::table_context::TableContextQueryIdentity; +use databend_common_catalog::table_context::TableContextQueryInfo; +use databend_common_config::GlobalConfig; +use databend_common_meta_app::schema::CatalogType; +use databend_common_sql::QueryLineage; +use databend_common_sql::QueryLineageColumn; +use databend_common_sql::QueryLineageKind; +use databend_common_sql::QueryLineageRelation; +use databend_common_sql::QueryLineageRelationKind; +use log::info; +use log::warn; +use serde::Serialize; +use sha2::Digest; +use sha2::Sha256; + +use crate::sessions::QueryContext; +use crate::sessions::convert_query_log_timestamp; + +const LINEAGE_LOG_TARGET: &str = "databend::log::lineage"; + +#[derive(Serialize)] +struct LineageLogEntry { + operation: &'static str, + query_id: String, + event_time: i64, + query_kind: String, + lineage_kind: &'static str, + source: LineageEndpoint, + target: LineageEndpoint, + column_lineage_hash: String, + source_to_target_columns: BTreeMap>, + target_to_source_columns: BTreeMap>, +} + +#[derive(Serialize)] +struct LineageObjectDeleteLogEntry { + operation: &'static str, + object_id: u64, + query_id: String, + event_time: i64, + query_kind: String, +} + +#[derive(Clone, Serialize)] +struct LineageEndpoint { + object_type: &'static str, + address_kind: &'static str, + catalog_type: Option<&'static str>, + catalog: String, + database: String, + name: String, + id: Option, + lineage_key: String, +} + +pub fn log_query_lineage(ctx: &Arc) { + let Some(lineage) = ctx.get_query_lineage() else { + return; + }; + + let query_id = ctx.get_id(); + let query_kind = ctx.get_query_kind().to_string(); + let event_time = convert_query_log_timestamp(SystemTime::now()); + + for entry in build_log_entries(lineage, query_id.clone(), query_kind.clone(), event_time) { + match serde_json::to_string(&entry) { + Ok(json) => info!(target: LINEAGE_LOG_TARGET, "{}", json), + Err(err) => warn!("Failed to serialize query lineage log: {:?}", err), + } + } +} + +/// Record an id-addressed object tombstone for asynchronous lineage cleanup. +/// The history transform is replay-safe, so this log must not affect the DDL result. +pub fn log_lineage_object_deletion(ctx: &Arc, object_id: u64) { + if !lineage_enabled() { + return; + } + + let entry = LineageObjectDeleteLogEntry { + operation: "DELETE_OBJECT", + object_id, + query_id: ctx.get_id(), + event_time: convert_query_log_timestamp(SystemTime::now()), + query_kind: ctx.get_query_kind().to_string(), + }; + match serde_json::to_string(&entry) { + Ok(json) => info!(target: LINEAGE_LOG_TARGET, "{}", json), + Err(err) => warn!("Failed to serialize lineage deletion log: {:?}", err), + } +} + +fn lineage_enabled() -> bool { + GlobalConfig::instance() + .log + .history + .is_table_enabled("lineage_unresolved") +} + +fn build_log_entries( + lineage: QueryLineage, + query_id: String, + query_kind: String, + event_time: i64, +) -> Vec { + let lineage_kind = lineage_kind_name(lineage.kind); + let mut entries = Vec::new(); + + for lineage_target in lineage.targets { + let target = endpoint_from_relation(lineage.kind, false, &lineage_target.relation); + for lineage_source in lineage_target.sources { + let source = endpoint_from_relation(lineage.kind, true, &lineage_source.relation); + if source.lineage_key == target.lineage_key { + continue; + } + + // Stages are object-level lineage endpoints. Their file fields are not stable schema + // columns, so COPY edges intentionally carry empty column maps. + let column_pairs = if source.object_type == "STAGE" || target.object_type == "STAGE" { + BTreeSet::new() + } else { + lineage_source + .columns + .into_iter() + .map(|edge| { + ( + column_address(&source, edge.source), + column_address(&target, edge.target), + ) + }) + .collect::>() + }; + let (source_to_target_columns, target_to_source_columns) = column_maps(&column_pairs); + + entries.push(LineageLogEntry { + operation: "UPSERT_EDGE", + query_id: query_id.clone(), + event_time, + query_kind: query_kind.clone(), + lineage_kind, + source: source.clone(), + target: target.clone(), + column_lineage_hash: column_lineage_hash(&column_pairs), + source_to_target_columns, + target_to_source_columns, + }); + } + } + + entries +} + +fn lineage_kind_name(kind: QueryLineageKind) -> &'static str { + match kind { + QueryLineageKind::Ctas => "CTAS", + QueryLineageKind::Dml => "DML", + QueryLineageKind::CreateView => "CREATE_VIEW", + } +} + +fn object_type_name(kind: &QueryLineageRelationKind) -> &'static str { + match kind { + QueryLineageRelationKind::Table => "TABLE", + QueryLineageRelationKind::View => "VIEW", + QueryLineageRelationKind::Stage => "STAGE", + } +} + +fn catalog_type_name(catalog_type: Option) -> Option<&'static str> { + catalog_type.map(|catalog_type| match catalog_type { + CatalogType::Default => "DEFAULT", + CatalogType::Hive => "HIVE", + CatalogType::Iceberg => "ICEBERG", + CatalogType::Paimon => "PAIMON", + }) +} + +fn endpoint_from_relation( + kind: QueryLineageKind, + is_source: bool, + relation: &QueryLineageRelation, +) -> LineageEndpoint { + let object_type = object_type_name(&relation.kind); + // View definitions retain source names. Renaming a referenced object should therefore make + // the edge inactive until that name resolves again. Other endpoints prefer stable IDs. + let relation_id = if relation.catalog_type != Some(CatalogType::Default) + || kind == QueryLineageKind::CreateView && is_source + { + None + } else { + relation.id + }; + let (address_kind, address_value) = match relation_id { + Some(id) => ("ID", id.to_string()), + None if relation.kind == QueryLineageRelationKind::Stage => ("NAME", relation.name.clone()), + None => ( + "NAME", + format!( + "{}.{}.{}", + relation.catalog, relation.database, relation.name + ), + ), + }; + let lineage_key = format!("{object_type}::{address_kind}::{address_value}"); + + LineageEndpoint { + object_type, + address_kind, + catalog_type: catalog_type_name(relation.catalog_type), + catalog: relation.catalog.clone(), + database: relation.database.clone(), + name: relation.name.clone(), + id: relation_id, + lineage_key, + } +} + +fn column_address(endpoint: &LineageEndpoint, column: QueryLineageColumn) -> String { + // View output column ids are query-plan ordinals rather than stable schema ids. Keep view + // columns name-addressed even though the view object itself has a stable table id. + if endpoint.object_type == "VIEW" || endpoint.address_kind == "NAME" { + column.name + } else { + column.id.to_string() + } +} + +fn column_maps( + column_pairs: &BTreeSet<(String, String)>, +) -> (BTreeMap>, BTreeMap>) { + let mut source_to_target = BTreeMap::>::new(); + let mut target_to_source = BTreeMap::>::new(); + for (source, target) in column_pairs { + source_to_target + .entry(source.clone()) + .or_default() + .insert(target.clone()); + target_to_source + .entry(target.clone()) + .or_default() + .insert(source.clone()); + } + + let into_map = |map: BTreeMap>| { + map.into_iter() + .map(|(key, values)| (key, values.into_iter().collect())) + .collect() + }; + (into_map(source_to_target), into_map(target_to_source)) +} + +fn column_lineage_hash(column_pairs: &BTreeSet<(String, String)>) -> String { + let mut hasher = Sha256::new(); + for (source, target) in column_pairs { + for value in [source, target] { + hasher.update((value.len() as u64).to_be_bytes()); + hasher.update(value.as_bytes()); + } + } + format!("{:x}", hasher.finalize()) +} + +#[cfg(test)] +mod tests { + use databend_common_sql::LineageSource; + use databend_common_sql::LineageTarget; + use databend_common_sql::QueryLineageColumnEdge; + + use super::*; + + #[test] + fn test_build_log_entries_uses_bidirectional_column_maps() { + let captured = lineage( + QueryLineageKind::Dml, + relation("source", Some(10)), + relation("target", Some(20)), + vec![ + edge(1, "a", 3, "x"), + edge(2, "b", 3, "x"), + edge(1, "a", 3, "x"), + ], + ); + + let entries = build_log_entries(captured, "query-1".into(), "INSERT".into(), 1); + assert_eq!(entries.len(), 1); + assert_eq!(entries[0].operation, "UPSERT_EDGE"); + assert_eq!( + entries[0].source_to_target_columns, + BTreeMap::from([ + ("1".to_string(), vec!["3".to_string()]), + ("2".to_string(), vec!["3".to_string()]), + ]) + ); + assert_eq!( + entries[0].target_to_source_columns, + BTreeMap::from([("3".to_string(), vec!["1".to_string(), "2".to_string()],)]) + ); + + let reordered = lineage( + QueryLineageKind::Dml, + relation("source", Some(10)), + relation("target", Some(20)), + vec![edge(2, "b", 3, "x"), edge(1, "a", 3, "x")], + ); + let reordered = build_log_entries(reordered, "query-2".into(), "INSERT".into(), 2); + assert_eq!( + entries[0].column_lineage_hash, + reordered[0].column_lineage_hash + ); + } + + #[test] + fn test_create_view_uses_source_names_and_target_ids() { + let lineage = lineage( + QueryLineageKind::CreateView, + relation("source", Some(10)), + view_relation("view", Some(20)), + vec![edge(1, "a", 3, "x")], + ); + + let entries = build_log_entries(lineage, "query-1".into(), "CREATE_VIEW".into(), 1); + assert_eq!(entries[0].source.address_kind, "NAME"); + assert_eq!(entries[0].target.address_kind, "ID"); + assert_eq!( + entries[0].source_to_target_columns, + BTreeMap::from([("a".to_string(), vec!["x".to_string()],)]) + ); + assert_eq!( + entries[0].target_to_source_columns, + BTreeMap::from([("x".to_string(), vec!["a".to_string()],)]) + ); + } + + #[test] + fn test_build_log_entries_skips_self_loop() { + let same = relation("table", Some(10)); + let lineage = lineage(QueryLineageKind::Dml, same.clone(), same, vec![edge( + 1, "a", 1, "a", + )]); + assert!(build_log_entries(lineage, "query-1".into(), "INSERT".into(), 1).is_empty()); + } + + #[test] + fn test_copy_stage_uses_name_addressing_without_column_maps() { + let lineage = lineage( + QueryLineageKind::Dml, + relation_with_kind("named_stage", None, QueryLineageRelationKind::Stage), + relation("target", Some(20)), + vec![edge(1, "file_col", 3, "x")], + ); + + let entries = build_log_entries(lineage, "query-1".into(), "COPY".into(), 1); + assert_eq!(entries.len(), 1); + assert_eq!(entries[0].source.object_type, "STAGE"); + assert_eq!(entries[0].source.address_kind, "NAME"); + assert_eq!(entries[0].source.catalog_type, None); + assert_eq!(entries[0].source.lineage_key, "STAGE::NAME::named_stage"); + assert!(entries[0].source_to_target_columns.is_empty()); + assert!(entries[0].target_to_source_columns.is_empty()); + } + + #[test] + fn test_external_catalog_uses_captured_names() { + let source = QueryLineageRelation { + catalog: "iceberg_catalog".to_string(), + database: "db".to_string(), + name: "source".to_string(), + id: None, + catalog_type: Some(CatalogType::Iceberg), + kind: QueryLineageRelationKind::Table, + }; + let lineage = lineage( + QueryLineageKind::Dml, + source, + relation("target", Some(20)), + vec![edge(1, "a", 3, "x")], + ); + + let entries = build_log_entries(lineage, "query-1".into(), "INSERT".into(), 1); + assert_eq!(entries[0].source.catalog_type, Some("ICEBERG")); + assert_eq!(entries[0].source.address_kind, "NAME"); + assert_eq!( + entries[0].source_to_target_columns, + BTreeMap::from([("a".to_string(), vec!["3".to_string()])]) + ); + } + + fn lineage( + kind: QueryLineageKind, + source: QueryLineageRelation, + target: QueryLineageRelation, + columns: Vec, + ) -> QueryLineage { + QueryLineage { + kind, + targets: vec![LineageTarget { + relation: target, + sources: vec![LineageSource { + relation: source, + columns, + }], + }], + } + } + + fn relation(name: &str, id: Option) -> QueryLineageRelation { + relation_with_kind(name, id, QueryLineageRelationKind::Table) + } + + fn view_relation(name: &str, id: Option) -> QueryLineageRelation { + relation_with_kind(name, id, QueryLineageRelationKind::View) + } + + fn relation_with_kind( + name: &str, + id: Option, + kind: QueryLineageRelationKind, + ) -> QueryLineageRelation { + let catalog_type = + (kind != QueryLineageRelationKind::Stage).then_some(CatalogType::Default); + QueryLineageRelation { + catalog: "default".to_string(), + database: "default".to_string(), + name: name.to_string(), + id, + catalog_type, + kind, + } + } + + fn edge( + source_id: u32, + source_name: &str, + target_id: u32, + target_name: &str, + ) -> QueryLineageColumnEdge { + QueryLineageColumnEdge { + source: QueryLineageColumn { + id: source_id, + name: source_name.to_string(), + }, + target: QueryLineageColumn { + id: target_id, + name: target_name.to_string(), + }, + } + } +} diff --git a/src/query/service/src/interpreters/common/mod.rs b/src/query/service/src/interpreters/common/mod.rs index f342544b52e5c..e8f65b2f5b0bb 100644 --- a/src/query/service/src/interpreters/common/mod.rs +++ b/src/query/service/src/interpreters/common/mod.rs @@ -15,6 +15,7 @@ mod column; mod finish_hook; mod grant; +mod lineage_log; mod metrics; mod notification; mod query_log; @@ -29,6 +30,8 @@ pub mod table_option_validation; pub use column::*; pub use finish_hook::QueryFinishHooks; pub use grant::validate_grant_object_exists; +pub use lineage_log::log_lineage_object_deletion; +pub use lineage_log::log_query_lineage; pub use log::*; pub use notification::get_notification_client_config; pub use query_log::InterpreterQueryLog; diff --git a/src/query/service/src/interpreters/interpreter.rs b/src/query/service/src/interpreters/interpreter.rs index 94d4d2b28011a..a77a94d86e1f6 100644 --- a/src/query/service/src/interpreters/interpreter.rs +++ b/src/query/service/src/interpreters/interpreter.rs @@ -45,6 +45,7 @@ use md5::Md5; use crate::interpreters::common::QueryFinishHooks; use crate::interpreters::common::log_query_finished; +use crate::interpreters::common::log_query_lineage; use crate::interpreters::common::log_query_start; use crate::interpreters::interpreter_txn_commit::execute_commit_statement; use crate::pipelines::PipelineBuildResult; @@ -199,6 +200,7 @@ async fn build_pipeline_before_execute( if build_res.main_pipeline.is_empty() { if hooks.log_finished { log_query_finished(&ctx, None); + log_query_lineage(&ctx); } return Ok(BuiltPipeline::Finished(Box::pin(DataBlockStream::create( None, diff --git a/src/query/service/src/interpreters/interpreter_copy_into_location.rs b/src/query/service/src/interpreters/interpreter_copy_into_location.rs index 50ffd1f8ba37b..565f99dd9ba9e 100644 --- a/src/query/service/src/interpreters/interpreter_copy_into_location.rs +++ b/src/query/service/src/interpreters/interpreter_copy_into_location.rs @@ -128,6 +128,7 @@ impl Interpreter for CopyIntoLocationInterpreter { debug!("ctx.id" = self.ctx.get_id().as_str(); "copy_into_location_interpreter_execute_v2"); if check_deduplicate_label(self.ctx.clone()).await? { + self.ctx.attach_query_lineage(None); return Ok(PipelineBuildResult::create()); } diff --git a/src/query/service/src/interpreters/interpreter_copy_into_table.rs b/src/query/service/src/interpreters/interpreter_copy_into_table.rs index 5efd01cbfdecc..7ede10916e87f 100644 --- a/src/query/service/src/interpreters/interpreter_copy_into_table.rs +++ b/src/query/service/src/interpreters/interpreter_copy_into_table.rs @@ -964,6 +964,7 @@ impl Interpreter for CopyIntoTableInterpreter { debug!("ctx.id" = self.ctx.get_id().as_str(); "copy_into_table_interpreter_execute_v2"); if check_deduplicate_label(self.ctx.clone()).await? { + self.ctx.attach_query_lineage(None); return Ok(PipelineBuildResult::create()); } @@ -977,9 +978,17 @@ impl Interpreter for CopyIntoTableInterpreter { ) .await?; + self.ctx.update_query_lineage_target_id( + plan.catalog_info.catalog_name(), + &plan.database_name, + &plan.table_name, + to_table.get_table_info().ident.table_id, + ); + to_table.check_mutable()?; if self.plan.no_file_to_copy { + self.ctx.attach_query_lineage(None); info!("no file to copy"); return self.on_no_files_to_copy().await; } diff --git a/src/query/service/src/interpreters/interpreter_factory.rs b/src/query/service/src/interpreters/interpreter_factory.rs index 427784b0e7c3d..f89544a0c9cb4 100644 --- a/src/query/service/src/interpreters/interpreter_factory.rs +++ b/src/query/service/src/interpreters/interpreter_factory.rs @@ -16,11 +16,13 @@ use std::sync::Arc; use databend_common_ast::ast::ExplainKind; use databend_common_catalog::lock::LockTableOption; +use databend_common_config::GlobalConfig; use databend_common_exception::ErrorCode; use databend_common_exception::Result; use databend_common_sql::binder::ExplainConfig; use databend_common_sql::plans::Mutation; use log::error; +use log::warn; use super::interpreter_catalog_create::CreateCatalogInterpreter; use super::interpreter_catalog_show_create::ShowCreateCatalogInterpreter; @@ -156,6 +158,18 @@ impl InterpreterFactory { let mut access_logger = AccessLogger::create(ctx.clone()); access_logger.log(plan); access_logger.output(); + + if lineage_enabled() { + match plan.query_lineage() { + Ok(lineage) => ctx.attach_query_lineage(lineage), + Err(err) => { + warn!("Failed to extract query lineage: {:?}", err); + ctx.attach_query_lineage(None); + } + } + } else { + ctx.attach_query_lineage(None); + } Self::get_warehouses_interpreter(ctx, plan, Self::get_inner) } @@ -937,3 +951,10 @@ impl InterpreterFactory { } } } + +fn lineage_enabled() -> bool { + GlobalConfig::instance() + .log + .history + .is_table_enabled("lineage_unresolved") +} diff --git a/src/query/service/src/interpreters/interpreter_insert.rs b/src/query/service/src/interpreters/interpreter_insert.rs index df8cea1bcf509..6def33f88acf2 100644 --- a/src/query/service/src/interpreters/interpreter_insert.rs +++ b/src/query/service/src/interpreters/interpreter_insert.rs @@ -291,6 +291,7 @@ impl Interpreter for InsertInterpreter { #[async_backtrace::framed] async fn execute2(&self) -> Result { if check_deduplicate_label(self.ctx.clone()).await? { + self.ctx.attach_query_lineage(None); return Ok(PipelineBuildResult::create()); } let table = if let Some(table_info) = &self.plan.table_info { @@ -310,6 +311,13 @@ impl Interpreter for InsertInterpreter { .await? }; + self.ctx.update_query_lineage_target_id( + &self.plan.catalog, + &self.plan.database, + &self.plan.table, + table.get_table_info().ident.table_id, + ); + let mut table_constraints = Vec::new(); // check mutability table.check_mutable()?; diff --git a/src/query/service/src/interpreters/interpreter_insert_multi_table.rs b/src/query/service/src/interpreters/interpreter_insert_multi_table.rs index e240db7f7933b..53f4796027ede 100644 --- a/src/query/service/src/interpreters/interpreter_insert_multi_table.rs +++ b/src/query/service/src/interpreters/interpreter_insert_multi_table.rs @@ -448,7 +448,14 @@ impl InsertMultiTableInterpreter { casted_schema, source_scalar_exprs, } = into; - let table = self.ctx.get_table(catalog, database, table).await?; + let table_name = table; + let table = self.ctx.get_table(catalog, database, table_name).await?; + self.ctx.update_query_lineage_target_id( + catalog, + database, + table_name, + table.get_table_info().ident.table_id, + ); branches.push( table, condition, diff --git a/src/query/service/src/interpreters/interpreter_mutation.rs b/src/query/service/src/interpreters/interpreter_mutation.rs index 1ac94e0c36abf..373281ae55f0b 100644 --- a/src/query/service/src/interpreters/interpreter_mutation.rs +++ b/src/query/service/src/interpreters/interpreter_mutation.rs @@ -94,6 +94,7 @@ impl Interpreter for MutationInterpreter { #[async_backtrace::framed] async fn execute2(&self) -> Result { if check_deduplicate_label(self.ctx.clone()).await? { + self.ctx.attach_query_lineage(None); return Ok(PipelineBuildResult::create()); } diff --git a/src/query/service/src/interpreters/interpreter_replace.rs b/src/query/service/src/interpreters/interpreter_replace.rs index 0ad331ec7dd5f..2d9ea11b6d992 100644 --- a/src/query/service/src/interpreters/interpreter_replace.rs +++ b/src/query/service/src/interpreters/interpreter_replace.rs @@ -94,6 +94,7 @@ impl Interpreter for ReplaceInterpreter { #[async_backtrace::framed] async fn execute2(&self) -> Result { if check_deduplicate_label(self.ctx.clone()).await? { + self.ctx.attach_query_lineage(None); return Ok(PipelineBuildResult::create()); } @@ -151,6 +152,13 @@ impl ReplaceInterpreter { .get_table(&plan.catalog, &plan.database, &plan.table) .await?; + self.ctx.update_query_lineage_target_id( + &plan.catalog, + &plan.database, + &plan.table, + table.get_table_info().ident.table_id, + ); + // check mutability table.check_mutable()?; diff --git a/src/query/service/src/interpreters/interpreter_table_create.rs b/src/query/service/src/interpreters/interpreter_table_create.rs index d0bb2dbdaef84..d7dc7c3cd8504 100644 --- a/src/query/service/src/interpreters/interpreter_table_create.rs +++ b/src/query/service/src/interpreters/interpreter_table_create.rs @@ -207,6 +207,7 @@ impl CreateTableInterpreter { let table_meta = req.table_meta.clone(); let reply = catalog.create_table(req.clone()).await?; if !reply.new_table && self.plan.create_option != CreateOption::CreateOrReplace { + self.ctx.attach_query_lineage(None); return Ok(PipelineBuildResult::create()); } if let Some(prefix) = req.table_meta.options.get(OPT_KEY_TEMP_PREFIX).cloned() { @@ -214,6 +215,12 @@ impl CreateTableInterpreter { } let table_id = reply.table_id; + self.ctx.update_query_lineage_target_id( + &self.plan.catalog, + &self.plan.database, + &self.plan.table, + table_id, + ); let prev_table_id = reply.prev_table_id; let orphan_table_name = reply.orphan_table_name.clone(); let table_id_seq = reply @@ -251,6 +258,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/service/src/interpreters/interpreter_table_drop.rs b/src/query/service/src/interpreters/interpreter_table_drop.rs index 2bf55ee8314f6..03b12a3870a96 100644 --- a/src/query/service/src/interpreters/interpreter_table_drop.rs +++ b/src/query/service/src/interpreters/interpreter_table_drop.rs @@ -30,6 +30,7 @@ use databend_common_users::UserApiProvider; use databend_storages_common_table_meta::table::OPT_KEY_TEMP_PREFIX; use crate::interpreters::Interpreter; +use crate::interpreters::common::log_lineage_object_deletion; use crate::pipelines::PipelineBuildResult; use crate::sessions::QueryContext; use crate::sessions::TableContextTableAccess; @@ -89,6 +90,8 @@ impl Interpreter for DropTableInterpreter { }; let is_temp = tbl.is_temp(); let table_id = tbl.get_table_info().ident.table_id; + let delete_lineage = !is_temp + && FuseTable::try_from_table(tbl.as_ref()).is_ok_and(|table| table.is_transient()); let engine = tbl.get_table_info().engine(); if matches!(engine, VIEW_ENGINE | STREAM_ENGINE) { @@ -133,6 +136,9 @@ impl Interpreter for DropTableInterpreter { .unwrap_or_default(), }) .await?; + if delete_lineage { + log_lineage_object_deletion(&self.ctx, table_id); + } if !is_temp && !catalog.is_external() { // iceberg table do not need to generate ownership diff --git a/src/query/service/src/interpreters/interpreter_vacuum_drop_tables.rs b/src/query/service/src/interpreters/interpreter_vacuum_drop_tables.rs index 98c11266ff2f9..fcbb283e82113 100644 --- a/src/query/service/src/interpreters/interpreter_vacuum_drop_tables.rs +++ b/src/query/service/src/interpreters/interpreter_vacuum_drop_tables.rs @@ -39,6 +39,7 @@ use databend_enterprise_vacuum_handler::get_vacuum_handler; use log::info; use crate::interpreters::Interpreter; +use crate::interpreters::common::log_lineage_object_deletion; use crate::pipelines::PipelineBuildResult; use crate::sessions::QueryContext; use crate::sessions::TableContextLicense; @@ -106,6 +107,11 @@ impl VacuumDropTablesInterpreter { drop_ids: c.to_vec(), }; num_meta_keys_removed += catalog.gc_drop_tables(req).await?; + for drop_id in c { + if let DroppedId::Table { id, .. } = drop_id { + log_lineage_object_deletion(&self.ctx, id.table_id); + } + } } // then gc drop db ids diff --git a/src/query/service/src/interpreters/interpreter_view_create.rs b/src/query/service/src/interpreters/interpreter_view_create.rs index a3de2eb7f978d..74fbdec81ebb9 100644 --- a/src/query/service/src/interpreters/interpreter_view_create.rs +++ b/src/query/service/src/interpreters/interpreter_view_create.rs @@ -137,7 +137,17 @@ impl Interpreter for CreateViewInterpreter { table_properties: None, table_partition: None, }; - catalog.create_table(plan).await?; + let reply = catalog.create_table(plan).await?; + if !reply.new_table && !self.plan.create_option.is_overriding() { + self.ctx.attach_query_lineage(None); + } else { + self.ctx.update_query_lineage_target_id( + &self.plan.catalog, + &self.plan.database, + &self.plan.view_name, + reply.table_id, + ); + } Ok(PipelineBuildResult::create()) } diff --git a/src/query/service/src/interpreters/interpreter_view_drop.rs b/src/query/service/src/interpreters/interpreter_view_drop.rs index 213de4d4d057f..20dead75072d1 100644 --- a/src/query/service/src/interpreters/interpreter_view_drop.rs +++ b/src/query/service/src/interpreters/interpreter_view_drop.rs @@ -23,6 +23,7 @@ use databend_common_storages_stream::stream_table::STREAM_ENGINE; use databend_storages_common_table_meta::table::OPT_KEY_TEMP_PREFIX; use crate::interpreters::Interpreter; +use crate::interpreters::common::log_lineage_object_deletion; use crate::pipelines::PipelineBuildResult; use crate::sessions::QueryContext; use crate::sessions::TableContextTableAccess; @@ -103,6 +104,7 @@ impl Interpreter for DropViewInterpreter { .unwrap_or_default(), }) .await?; + log_lineage_object_deletion(&self.ctx, table.get_id()); }; Ok(PipelineBuildResult::create()) diff --git a/src/query/service/src/sessions/query_ctx.rs b/src/query/service/src/sessions/query_ctx.rs index e5dff77b98c44..e316b3a2ba7a5 100644 --- a/src/query/service/src/sessions/query_ctx.rs +++ b/src/query/service/src/sessions/query_ctx.rs @@ -122,6 +122,7 @@ use databend_common_pipeline::core::LockGuard; use databend_common_pipeline::core::PlanProfile; use databend_common_settings::Settings; use databend_common_sql::IndexType; +use databend_common_sql::QueryLineage; use databend_common_storage::DataOperator; use databend_common_storage::FileStatus; use databend_common_storage::StageFileInfo; @@ -157,6 +158,7 @@ use jiff::Zoned; use jiff::tz::TimeZone; use log::debug; use log::info; +use log::warn; use parking_lot::Mutex; use parking_lot::RwLock; use tokio::sync::Semaphore; @@ -548,6 +550,56 @@ impl QueryContext { (finish_time - query_start_time) / 1_000 } + pub fn attach_query_lineage(&self, lineage: Option) { + self.shared.attach_query_lineage(lineage); + } + + pub fn get_query_lineage(&self) -> Option { + self.shared.get_query_lineage() + } + + /// Reconcile captured lineage after execution resolves the actual target. + /// + /// A missing captured id belongs to paths such as CTAS whose target is created during + /// execution. If a captured id changed, however, its bind-time column ids belong to the old + /// object and cannot safely be attached to the replacement table, so that target is skipped. + /// This does not change which table the query writes to. + pub fn update_query_lineage_target_id( + &self, + catalog: &str, + database: &str, + table: &str, + table_id: u64, + ) { + let Some(mut lineage) = self.get_query_lineage() else { + return; + }; + lineage.targets.retain_mut(|target| { + if target.relation.catalog != catalog + || target.relation.database != database + || target.relation.name != table + { + return true; + } + + if target + .relation + .id + .is_some_and(|captured_id| captured_id != table_id) + { + warn!( + "Skipping lineage target after table identity changed between bind and execution: {}.{}.{}", + catalog, database, table + ); + return false; + } + + target.relation.id = Some(table_id); + true + }); + self.attach_query_lineage((!lineage.targets.is_empty()).then_some(lineage)); + } + pub fn get_created_time(&self) -> SystemTime { self.shared.created_time } diff --git a/src/query/service/src/sessions/query_ctx_shared.rs b/src/query/service/src/sessions/query_ctx_shared.rs index 9b9ca5f8d3abb..e1de153809f91 100644 --- a/src/query/service/src/sessions/query_ctx_shared.rs +++ b/src/query/service/src/sessions/query_ctx_shared.rs @@ -57,6 +57,7 @@ use databend_common_meta_app::principal::UserInfo; use databend_common_meta_app::tenant::Tenant; use databend_common_pipeline::core::PlanProfile; use databend_common_settings::Settings; +use databend_common_sql::QueryLineage; use databend_common_storage::DataOperator; use databend_common_storage::StorageMetrics; use databend_common_storages_stream::stream_table::StreamTable; @@ -119,6 +120,7 @@ pub struct QueryContextShared { running_query_kind: Arc>>, running_query_text_hash: Arc>>, running_query_parameterized_hash: Arc>>, + query_lineage: Arc>>, aborting: Arc, pub(super) abort_notify: Arc, pub(super) tables_refs: Arc>>>, @@ -217,6 +219,7 @@ impl QueryContextShared { running_query_kind: Arc::new(RwLock::new(None)), running_query_text_hash: Arc::new(RwLock::new(None)), running_query_parameterized_hash: Arc::new(RwLock::new(None)), + query_lineage: Arc::new(RwLock::new(None)), aborting: Arc::new(AtomicBool::new(false)), abort_notify: Arc::new(WatchNotify::new()), tables_refs: Arc::new(Mutex::new(HashMap::new())), @@ -654,6 +657,14 @@ impl QueryContextShared { .unwrap_or(QueryKind::Unknown) } + pub fn attach_query_lineage(&self, lineage: Option) { + *self.query_lineage.write() = lineage; + } + + pub fn get_query_lineage(&self) -> Option { + self.query_lineage.read().clone() + } + pub fn get_connection_id(&self) -> String { self.session.get_id() } diff --git a/src/query/service/src/table_functions/get_lineage/edge_reader.rs b/src/query/service/src/table_functions/get_lineage/edge_reader.rs new file mode 100644 index 0000000000000..f4135442deb5f --- /dev/null +++ b/src/query/service/src/table_functions/get_lineage/edge_reader.rs @@ -0,0 +1,533 @@ +// 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::sync::Arc; + +use databend_common_catalog::catalog::CATALOG_DEFAULT; +use databend_common_catalog::plan::Filters; +use databend_common_catalog::plan::Projection; +use databend_common_catalog::plan::PushDownInfo; +use databend_common_catalog::table::Table; +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; +use databend_common_expression::ColumnRef; +use databend_common_expression::Constant; +use databend_common_expression::DataBlock; +use databend_common_expression::Expr; +use databend_common_expression::Scalar; +use databend_common_expression::ScalarRef; +use databend_common_expression::TableSchema; +use databend_common_expression::type_check::check_function; +use databend_common_expression::types::DataType; +use databend_common_expression::types::NumberScalar; +use databend_common_functions::BUILTIN_FUNCTIONS; +use databend_common_pipeline::core::always_callback; +use databend_common_sql::executor::table_read_plan::ToReadDataSourcePlan; +use futures::TryStreamExt; + +use crate::interpreters::QueryFinishHooks; +use crate::physical_plans::PhysicalPlan; +use crate::physical_plans::PhysicalPlanMeta; +use crate::physical_plans::TableScan; +use crate::pipelines::executor::ExecutorSettings; +use crate::pipelines::executor::PipelinePullingExecutor; +use crate::schedulers::build_local_pipeline; +use crate::sessions::QueryContext; +use crate::sessions::TableContext; +use crate::sessions::TableContextQueryState; +use crate::stream::PullingExecutorStream; + +const HISTORY_DATABASE: &str = "system_history"; +const LINEAGE_TABLE: &str = "lineage_unresolved"; +const FRONTIER_BATCH_SIZE: usize = 256; +const FIRST_LINEAGE_SCAN_ID: usize = 1_000; + +const EDGE_COLUMNS: &[&str] = &[ + "query_id", + "event_time", + "query_kind", + "lineage_kind", + "column_lineage_hash", + "source_lineage_key", + "source_address_kind", + "source_catalog_type", + "source_object_type", + "source_catalog", + "source_database", + "source_name", + "source_id", + "target_lineage_key", + "target_address_kind", + "target_catalog_type", + "target_object_type", + "target_catalog", + "target_database", + "target_name", + "target_id", + "source_to_target_columns", + "target_to_source_columns", +]; + +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub(super) enum AddressKind { + Id, + Name, +} + +impl AddressKind { + fn parse(value: Option) -> Result { + match value.as_deref() { + Some("ID") => Ok(Self::Id), + Some("NAME") => Ok(Self::Name), + other => Err(ErrorCode::Internal(format!( + "invalid lineage address kind: {other:?}" + ))), + } + } +} + +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub(super) enum LineageObjectType { + Table, + View, + Stage, +} + +impl LineageObjectType { + fn parse(value: Option) -> Result { + match value.as_deref() { + Some("TABLE") => Ok(Self::Table), + Some("VIEW") => Ok(Self::View), + Some("STAGE") => Ok(Self::Stage), + other => Err(ErrorCode::Internal(format!( + "invalid lineage object type: {other:?}" + ))), + } + } + + pub(super) fn as_str(self) -> &'static str { + match self { + Self::Table => "TABLE", + Self::View => "VIEW", + Self::Stage => "STAGE", + } + } +} + +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +pub(super) struct CapturedObject { + pub lineage_key: String, + pub address_kind: AddressKind, + pub catalog_type: String, + pub object_type: LineageObjectType, + pub catalog: String, + pub database: String, + pub name: String, + pub id: Option, +} + +impl CapturedObject { + pub(super) fn is_default_catalog(&self) -> bool { + self.catalog_type.eq_ignore_ascii_case("DEFAULT") + } +} + +#[derive(Clone, Debug)] +pub(super) struct RawLineageEdge { + pub query_id: Option, + pub event_time: Option, + pub query_kind: Option, + pub lineage_kind: Option, + pub column_lineage_hash: String, + pub source: CapturedObject, + pub target: CapturedObject, + pub source_to_target_columns: BTreeMap>, + pub target_to_source_columns: BTreeMap>, +} + +impl RawLineageEdge { + pub(super) fn newer_than(&self, other: &Self) -> bool { + ( + self.event_time, + self.query_id.as_deref().unwrap_or_default(), + self.lineage_kind.as_deref().unwrap_or_default(), + self.column_lineage_hash.as_str(), + ) > ( + other.event_time, + other.query_id.as_deref().unwrap_or_default(), + other.lineage_kind.as_deref().unwrap_or_default(), + other.column_lineage_hash.as_str(), + ) + } +} + +pub(super) struct LineageEdgeReader { + ctx: Arc, + table: Arc, + next_scan_id: usize, +} + +impl LineageEdgeReader { + pub(super) async fn try_create(ctx: Arc) -> Result { + let table = ctx + .get_table(CATALOG_DEFAULT, HISTORY_DATABASE, LINEAGE_TABLE) + .await?; + Ok(Self { + ctx, + table, + next_scan_id: FIRST_LINEAGE_SCAN_ID, + }) + } + + /// Read every edge whose directional lineage key is in the current frontier. + /// + /// The table object is pinned for the lifetime of the reader, so every level reads the same + /// Fuse snapshot even if the history transform commits concurrently. Scan ids are allocated + /// monotonically per batch instead of reserving a fixed-size range for each distance level. + pub(super) async fn read_frontier( + &mut self, + match_column: &str, + frontier: &BTreeSet, + ) -> Result> { + if frontier.is_empty() { + return Ok(vec![]); + } + + let mut edges = Vec::new(); + let keys = frontier.iter().cloned().collect::>(); + for chunk in keys.chunks(FRONTIER_BATCH_SIZE) { + let scan_id = self.next_scan_id; + self.next_scan_id = self + .next_scan_id + .checked_add(1) + .ok_or_else(|| ErrorCode::Internal("lineage scan id overflow"))?; + edges.extend(self.read_batch(match_column, chunk, scan_id).await?); + } + Ok(edges) + } + + async fn read_batch( + &self, + match_column: &str, + keys: &[String], + scan_id: usize, + ) -> Result> { + let parent_ctx = self + .ctx + .as_any() + .downcast_ref::() + .ok_or_else(|| ErrorCode::Internal("GET_LINEAGE requires QueryContext"))?; + let child_ctx = QueryContext::create_from(parent_ctx); + if child_ctx.check_aborting().is_err() { + return Err(ErrorCode::AbortedQuery("GET_LINEAGE query was aborted")); + } + let schema = self.table.schema(); + let projection = Projection::from_column_names(&schema, EDGE_COLUMNS)?; + let filters = build_key_filter(&schema, match_column, keys)?; + let push_downs = PushDownInfo { + projection: Some(projection), + filters: Some(filters), + ..Default::default() + }; + + let mut source = self + .table + .read_plan(child_ctx.clone(), Some(push_downs), None, false, false) + .await?; + source.scan_id = scan_id; + let output_schema = source.output_schema.clone(); + let name_mapping = source + .output_schema + .fields() + .iter() + .map(|field| (field.name().to_string(), field.name().to_string())) + .collect(); + let physical_plan = PhysicalPlan::new(TableScan { + meta: PhysicalPlanMeta::new("LineageEdgeScan"), + scan_id, + name_mapping, + source: Box::new(source), + internal_column: None, + table_index: None, + stat_info: None, + }); + + let mut build_res = build_local_pipeline(&child_ctx, &physical_plan).await?; + build_res.main_pipeline.set_on_finished(always_callback( + QueryFinishHooks::nested().into_callback(child_ctx.clone()), + )); + let settings = ExecutorSettings::try_create(child_ctx.clone())?; + let pulling_executor = PipelinePullingExecutor::from_pipelines(build_res, settings)?; + let child_executor = pulling_executor.get_inner(); + let abort_notify = child_ctx.get_abort_notify(); + let match_column = match_column.to_string(); + let frontier = keys.iter().cloned().collect::>(); + + // Child contexts share the outer executor slot, so set_executor() would hide the outer + // pipeline from query cancellation. Keep it registered and explicitly forward the shared + // abort notification to this child executor instead. + let mut stream = PullingExecutorStream::create(pulling_executor)?; + let abort = abort_notify.notified(); + tokio::pin!(abort); + let mut edges = Vec::new(); + loop { + tokio::select! { + block = stream.try_next() => match block? { + Some(block) => edges.extend(decode_edges( + &block, + &output_schema, + &match_column, + &frontier, + )?), + None => return Ok(edges), + }, + () = &mut abort => { + let cause = ErrorCode::AbortedQuery( + "GET_LINEAGE child scan was aborted", + ); + child_executor.finish(Some(cause.clone())); + return Err(cause); + } + } + } + } +} + +fn build_key_filter(schema: &TableSchema, column: &str, keys: &[String]) -> Result { + let field = schema.field_with_name(column)?; + let data_type = DataType::from(field.data_type()); + let column_expr = || { + Expr::ColumnRef(ColumnRef { + span: None, + id: column.to_string(), + data_type: data_type.clone(), + display_name: column.to_string(), + }) + }; + + let mut predicates = keys.iter().map(|key| { + check_function( + None, + "eq", + &[], + &[ + column_expr(), + Expr::Constant(Constant { + span: None, + scalar: Scalar::String(key.clone()), + data_type: DataType::String, + }), + ], + &BUILTIN_FUNCTIONS, + ) + }); + let mut filter = predicates + .next() + .ok_or_else(|| ErrorCode::Internal("lineage frontier must not be empty"))??; + for predicate in predicates { + filter = check_function( + None, + "or_filters", + &[], + &[filter, predicate?], + &BUILTIN_FUNCTIONS, + )?; + } + let inverted_filter = check_function(None, "not", &[], &[filter.clone()], &BUILTIN_FUNCTIONS)?; + Ok(Filters { + filter: filter.as_remote_expr(), + inverted_filter: inverted_filter.as_remote_expr(), + }) +} + +fn decode_edges( + block: &DataBlock, + schema: &TableSchema, + match_column: &str, + frontier: &BTreeSet, +) -> Result> { + let offset = |name: &str| schema.index_of(name); + let mut edges = Vec::with_capacity(block.num_rows()); + for row in 0..block.num_rows() { + let match_key = required_string(block, offset(match_column)?, row)?; + if !frontier.contains(&match_key) { + continue; + } + + let source = CapturedObject { + lineage_key: required_string(block, offset("source_lineage_key")?, row)?, + address_kind: AddressKind::parse(optional_string( + block, + offset("source_address_kind")?, + row, + )?)?, + catalog_type: optional_string(block, offset("source_catalog_type")?, row)? + .unwrap_or_default(), + object_type: LineageObjectType::parse(optional_string( + block, + offset("source_object_type")?, + row, + )?)?, + catalog: optional_string(block, offset("source_catalog")?, row)?.unwrap_or_default(), + database: optional_string(block, offset("source_database")?, row)?.unwrap_or_default(), + name: optional_string(block, offset("source_name")?, row)?.unwrap_or_default(), + id: optional_u64(block, offset("source_id")?, row)?, + }; + let target = CapturedObject { + lineage_key: required_string(block, offset("target_lineage_key")?, row)?, + address_kind: AddressKind::parse(optional_string( + block, + offset("target_address_kind")?, + row, + )?)?, + catalog_type: optional_string(block, offset("target_catalog_type")?, row)? + .unwrap_or_default(), + object_type: LineageObjectType::parse(optional_string( + block, + offset("target_object_type")?, + row, + )?)?, + catalog: optional_string(block, offset("target_catalog")?, row)?.unwrap_or_default(), + database: optional_string(block, offset("target_database")?, row)?.unwrap_or_default(), + name: optional_string(block, offset("target_name")?, row)?.unwrap_or_default(), + id: optional_u64(block, offset("target_id")?, row)?, + }; + edges.push(RawLineageEdge { + query_id: optional_string(block, offset("query_id")?, row)?, + event_time: optional_timestamp(block, offset("event_time")?, row)?, + query_kind: optional_string(block, offset("query_kind")?, row)?, + lineage_kind: optional_string(block, offset("lineage_kind")?, row)?, + column_lineage_hash: required_string(block, offset("column_lineage_hash")?, row)?, + source, + target, + source_to_target_columns: string_array_map( + block, + offset("source_to_target_columns")?, + row, + )?, + target_to_source_columns: string_array_map( + block, + offset("target_to_source_columns")?, + row, + )?, + }); + } + Ok(edges) +} + +fn scalar_at(block: &DataBlock, offset: usize, row: usize) -> Result> { + block + .get_by_offset(offset) + .index(row) + .ok_or_else(|| ErrorCode::Internal("lineage block row is out of bounds")) +} + +fn optional_string(block: &DataBlock, offset: usize, row: usize) -> Result> { + match scalar_at(block, offset, row)? { + ScalarRef::Null => Ok(None), + ScalarRef::String(value) => Ok(Some(value.to_string())), + value => Err(ErrorCode::Internal(format!( + "expected lineage string, got {value:?}" + ))), + } +} + +fn required_string(block: &DataBlock, offset: usize, row: usize) -> Result { + optional_string(block, offset, row)? + .ok_or_else(|| ErrorCode::Internal("required lineage string is NULL")) +} + +fn optional_u64(block: &DataBlock, offset: usize, row: usize) -> Result> { + match scalar_at(block, offset, row)? { + ScalarRef::Null => Ok(None), + ScalarRef::Number(NumberScalar::UInt64(value)) => Ok(Some(value)), + value => Err(ErrorCode::Internal(format!( + "expected lineage UInt64, got {value:?}" + ))), + } +} + +fn optional_timestamp(block: &DataBlock, offset: usize, row: usize) -> Result> { + match scalar_at(block, offset, row)? { + ScalarRef::Null => Ok(None), + ScalarRef::Timestamp(value) => Ok(Some(value)), + value => Err(ErrorCode::Internal(format!( + "expected lineage timestamp, got {value:?}" + ))), + } +} + +fn string_array_map( + block: &DataBlock, + offset: usize, + row: usize, +) -> Result>> { + let scalar = scalar_at(block, offset, row)?; + let entries = match scalar { + ScalarRef::EmptyMap | ScalarRef::Null => return Ok(BTreeMap::new()), + ScalarRef::Map(entries) => entries, + value => { + return Err(ErrorCode::Internal(format!( + "expected lineage column map, got {value:?}" + ))); + } + }; + + let mut result = BTreeMap::new(); + for entry in entries.iter() { + let ScalarRef::Tuple(fields) = entry else { + return Err(ErrorCode::Internal(format!( + "expected lineage column map entry, got {entry:?}" + ))); + }; + if fields.len() != 2 { + return Err(ErrorCode::Internal(format!( + "expected lineage column map entry with 2 fields, got {}", + fields.len() + ))); + } + + let key = match &fields[0] { + ScalarRef::String(value) => value.to_string(), + ScalarRef::Null => continue, + value => { + return Err(ErrorCode::Internal(format!( + "expected lineage column map key, got {value:?}" + ))); + } + }; + let values = match &fields[1] { + ScalarRef::EmptyArray | ScalarRef::Null => Vec::new(), + ScalarRef::Array(values) => values + .iter() + .filter_map(|value| match value { + ScalarRef::String(value) => Some(Ok(value.to_string())), + ScalarRef::Null => None, + value => Some(Err(ErrorCode::Internal(format!( + "expected lineage column reference, got {value:?}" + )))), + }) + .collect::>>()?, + value => { + return Err(ErrorCode::Internal(format!( + "expected lineage column reference array, got {value:?}" + ))); + } + }; + result.insert(key, values); + } + Ok(result) +} diff --git a/src/query/service/src/table_functions/get_lineage/mod.rs b/src/query/service/src/table_functions/get_lineage/mod.rs new file mode 100644 index 0000000000000..e53113f8a0d52 --- /dev/null +++ b/src/query/service/src/table_functions/get_lineage/mod.rs @@ -0,0 +1,386 @@ +// 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. + +mod edge_reader; +mod resolver; +mod traversal; + +use std::any::Any; +use std::sync::Arc; + +use databend_common_catalog::plan::DataSourcePlan; +use databend_common_catalog::plan::PartStatistics; +use databend_common_catalog::plan::Partitions; +use databend_common_catalog::plan::PushDownInfo; +use databend_common_catalog::table::Table; +use databend_common_catalog::table_args::TableArgs; +use databend_common_catalog::table_args::i64_value; +use databend_common_catalog::table_args::string_value; +use databend_common_catalog::table_function::TableFunction; +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; +use databend_common_expression::DataBlock; +use databend_common_expression::DataSchema; +use databend_common_expression::DataSchemaRef; +use databend_common_expression::FromData; +use databend_common_expression::TableDataType; +use databend_common_expression::TableField; +use databend_common_expression::TableSchema; +use databend_common_expression::TableSchemaRefExt; +use databend_common_expression::types::Int32Type; +use databend_common_expression::types::StringType; +use databend_common_meta_app::schema::TableIdent; +use databend_common_meta_app::schema::TableInfo; +use databend_common_meta_app::schema::TableMeta; +use databend_common_pipeline::core::OutputPort; +use databend_common_pipeline::core::Pipeline; +use databend_common_pipeline::core::processor::ProcessorPtr; +use databend_common_pipeline::sources::AsyncSource; +use databend_common_pipeline::sources::AsyncSourcer; +use databend_meta_client::types::MetaId; + +use crate::sessions::TableContext; + +const GET_LINEAGE_FUNC: &str = "get_lineage"; +const GET_LINEAGE_ENGINE: &str = "GET_LINEAGE"; +const DEFAULT_DISTANCE: u8 = 5; +const MAX_DISTANCE: u8 = 5; + +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +enum ObjectDomain { + Table, + View, + Stage, + Column, +} + +impl ObjectDomain { + fn parse(value: &str) -> Result { + match value.trim().to_ascii_uppercase().as_str() { + "TABLE" => Ok(Self::Table), + "VIEW" => Ok(Self::View), + "STAGE" => Ok(Self::Stage), + "COLUMN" => Ok(Self::Column), + other => Err(ErrorCode::BadArguments(format!( + "unsupported object_domain '{other}', expected TABLE, VIEW, STAGE, or COLUMN" + ))), + } + } +} + +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +enum QueryDirection { + /// Walk from a target object toward the source objects that feed it. + Upstream, + /// Walk from a source object toward the target objects that consume it. + Downstream, +} + +impl QueryDirection { + fn parse(value: &str) -> Result { + match value.trim().to_ascii_uppercase().as_str() { + "UPSTREAM" => Ok(Self::Upstream), + "DOWNSTREAM" => Ok(Self::Downstream), + other => Err(ErrorCode::BadArguments(format!( + "unsupported direction '{other}', expected UPSTREAM or DOWNSTREAM" + ))), + } + } + + fn match_column(self) -> &'static str { + match self { + Self::Upstream => "target_lineage_key", + Self::Downstream => "source_lineage_key", + } + } +} + +#[derive(Clone, Debug)] +struct GetLineageArgs { + object_name: String, + object_domain: ObjectDomain, + direction: QueryDirection, + distance: u8, +} + +impl GetLineageArgs { + fn parse(table_args: &TableArgs) -> Result { + let args = table_args.expect_all_positioned(GET_LINEAGE_FUNC, None)?; + if !(args.len() == 3 || args.len() == 4) { + return Err(ErrorCode::BadArguments(format!( + "{GET_LINEAGE_FUNC} requires 3 or 4 positioned arguments: object_name, object_domain, direction[, distance]" + ))); + } + + let distance = if args.len() == 4 { + let distance = i64_value(&args[3])?; + if !(1..=MAX_DISTANCE as i64).contains(&distance) { + return Err(ErrorCode::BadArguments(format!( + "distance must be an integer in the range [1, {MAX_DISTANCE}]" + ))); + } + distance as u8 + } else { + DEFAULT_DISTANCE + }; + + Ok(Self { + object_name: string_value(&args[0])?, + object_domain: ObjectDomain::parse(&string_value(&args[1])?)?, + direction: QueryDirection::parse(&string_value(&args[2])?)?, + distance, + }) + } +} + +#[derive(Clone, Debug)] +struct LineageResultRow { + distance: i32, + source_object_domain: Option, + source_object_name: Option, + source_column_name: Option, + target_object_domain: Option, + target_object_name: Option, + target_column_name: Option, + target_status: String, + process: Option, +} + +pub struct GetLineageTable { + table_info: TableInfo, + table_args: TableArgs, + args: GetLineageArgs, +} + +impl GetLineageTable { + pub fn create( + database_name: &str, + table_func_name: &str, + table_id: MetaId, + table_args: TableArgs, + ) -> Result> { + let args = GetLineageArgs::parse(&table_args)?; + let table_info = TableInfo { + ident: TableIdent::new(table_id, 0), + desc: format!("'{}'.'{}'", database_name, table_func_name), + name: table_func_name.to_string(), + meta: TableMeta { + schema: Self::schema(), + engine: GET_LINEAGE_ENGINE.to_string(), + ..Default::default() + }, + ..Default::default() + }; + Ok(Arc::new(Self { + table_info, + table_args, + args, + })) + } + + fn schema() -> Arc { + let nullable_string = || TableDataType::Nullable(Box::new(TableDataType::String)); + TableSchemaRefExt::create(vec![ + TableField::new( + "distance", + TableDataType::Number(databend_common_expression::types::NumberDataType::Int32), + ), + TableField::new("source_object_domain", nullable_string()), + TableField::new("source_object_name", nullable_string()), + TableField::new("source_column_name", nullable_string()), + TableField::new("target_object_domain", nullable_string()), + TableField::new("target_object_name", nullable_string()), + TableField::new("target_column_name", nullable_string()), + TableField::new("target_status", TableDataType::String), + TableField::new("process", nullable_string()), + ]) + } +} + +#[async_trait::async_trait] +impl Table for GetLineageTable { + fn as_any(&self) -> &dyn Any { + self + } + + fn get_table_info(&self) -> &TableInfo { + &self.table_info + } + + async fn read_partitions( + &self, + _ctx: Arc, + _push_downs: Option, + _dry_run: bool, + ) -> Result<(PartStatistics, Partitions)> { + Ok((PartStatistics::default(), Partitions::default())) + } + + fn table_args(&self) -> Option { + Some(self.table_args.clone()) + } + + fn read_data( + &self, + ctx: Arc, + _plan: &DataSourcePlan, + pipeline: &mut Pipeline, + _put_cache: bool, + ) -> Result<()> { + let args = self.args.clone(); + let schema = self.table_info.meta.schema.clone(); + pipeline.add_source( + |output| GetLineageSource::create(ctx.clone(), output, args.clone(), schema.clone()), + 1, + )?; + Ok(()) + } + + // The result depends on a history-table snapshot acquired inside this opaque source. + fn result_can_be_cached(&self) -> bool { + false + } +} + +impl TableFunction for GetLineageTable { + fn function_name(&self) -> &str { + GET_LINEAGE_FUNC + } + + fn as_table<'a>(self: Arc) -> Arc + where Self: 'a { + self + } +} + +struct GetLineageSource { + ctx: Arc, + args: GetLineageArgs, + schema: DataSchemaRef, + finished: bool, +} + +impl GetLineageSource { + fn create( + ctx: Arc, + output: Arc, + args: GetLineageArgs, + schema: Arc, + ) -> Result { + AsyncSourcer::create(ctx.get_scan_progress(), output, Self { + ctx, + args, + schema: Arc::new(DataSchema::from(schema.as_ref())), + finished: false, + }) + } + + fn build_block(&self, rows: Vec) -> DataBlock { + if rows.is_empty() { + return DataBlock::empty_with_schema(&self.schema); + } + + let mut distances = Vec::with_capacity(rows.len()); + let mut source_domains = Vec::with_capacity(rows.len()); + let mut source_names = Vec::with_capacity(rows.len()); + let mut source_columns = Vec::with_capacity(rows.len()); + let mut target_domains = Vec::with_capacity(rows.len()); + let mut target_names = Vec::with_capacity(rows.len()); + let mut target_columns = Vec::with_capacity(rows.len()); + let mut target_statuses = Vec::with_capacity(rows.len()); + let mut processes = Vec::with_capacity(rows.len()); + + for row in rows { + distances.push(row.distance); + source_domains.push(row.source_object_domain); + source_names.push(row.source_object_name); + source_columns.push(row.source_column_name); + target_domains.push(row.target_object_domain); + target_names.push(row.target_object_name); + target_columns.push(row.target_column_name); + target_statuses.push(row.target_status); + processes.push(row.process); + } + + DataBlock::new_from_columns(vec![ + Int32Type::from_data(distances), + StringType::from_opt_data(source_domains), + StringType::from_opt_data(source_names), + StringType::from_opt_data(source_columns), + StringType::from_opt_data(target_domains), + StringType::from_opt_data(target_names), + StringType::from_opt_data(target_columns), + StringType::from_data(target_statuses), + StringType::from_opt_data(processes), + ]) + } +} + +#[async_trait::async_trait] +impl AsyncSource for GetLineageSource { + const NAME: &'static str = "GetLineageSource"; + + async fn generate(&mut self) -> Result> { + if self.finished { + return Ok(None); + } + let rows = traversal::traverse(self.ctx.clone(), self.args.clone()).await?; + self.finished = true; + Ok(Some(self.build_block(rows))) + } +} + +#[cfg(test)] +mod tests { + use databend_common_expression::Scalar; + use databend_common_expression::types::NumberScalar; + + use super::*; + + #[test] + fn test_get_lineage_args() -> Result<()> { + let default_distance = TableArgs::new_positioned(vec![ + Scalar::String("db.table".to_string()), + Scalar::String("TABLE".to_string()), + Scalar::String("UPSTREAM".to_string()), + ]); + let parsed = GetLineageArgs::parse(&default_distance)?; + assert_eq!(parsed.object_domain, ObjectDomain::Table); + assert_eq!(parsed.direction, QueryDirection::Upstream); + assert_eq!(parsed.distance, DEFAULT_DISTANCE); + + let explicit_distance = TableArgs::new_positioned(vec![ + Scalar::String("db.table.col".to_string()), + Scalar::String("COLUMN".to_string()), + Scalar::String("DOWNSTREAM".to_string()), + Scalar::Number(NumberScalar::Int64(2)), + ]); + let parsed = GetLineageArgs::parse(&explicit_distance)?; + assert_eq!(parsed.object_domain, ObjectDomain::Column); + assert_eq!(parsed.direction, QueryDirection::Downstream); + assert_eq!(parsed.distance, 2); + Ok(()) + } + + #[test] + fn test_get_lineage_rejects_invalid_distance() { + let args = TableArgs::new_positioned(vec![ + Scalar::String("db.table".to_string()), + Scalar::String("TABLE".to_string()), + Scalar::String("UPSTREAM".to_string()), + Scalar::Number(NumberScalar::Int64(6)), + ]); + assert!(GetLineageArgs::parse(&args).is_err()); + } +} diff --git a/src/query/service/src/table_functions/get_lineage/resolver.rs b/src/query/service/src/table_functions/get_lineage/resolver.rs new file mode 100644 index 0000000000000..7bb280ffe730a --- /dev/null +++ b/src/query/service/src/table_functions/get_lineage/resolver.rs @@ -0,0 +1,481 @@ +// 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::BTreeSet; +use std::collections::HashMap; +use std::sync::Arc; + +use databend_common_ast::parser::parse_table_ref; +use databend_common_catalog::catalog::CATALOG_DEFAULT; +use databend_common_catalog::table::Table; +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; +use databend_common_expression::TableSchemaRef; +use databend_common_expression::infer_table_schema; +use databend_common_meta_api::kv_pb_api::KVPbApi; +use databend_common_meta_app::schema::TableIdToName; +use databend_common_sql::Planner; +use databend_common_sql::planner::NameResolutionContext; +use databend_common_sql::planner::normalize_identifier; +use databend_common_storages_basic::view_table::QUERY; +use databend_common_storages_basic::view_table::VIEW_ENGINE; +use databend_common_storages_stream::stream_table::STREAM_ENGINE; +use databend_common_users::GrantObjectVisibilityChecker; +use databend_common_users::Object; +use databend_common_users::UserApiProvider; +use log::warn; + +use super::ObjectDomain; +use super::edge_reader::AddressKind; +use super::edge_reader::CapturedObject; +use super::edge_reader::LineageObjectType; +use crate::meta_service_error; +use crate::sessions::TableContext; + +#[derive(Clone, Debug)] +pub(super) struct ResolvedObject { + pub object_type: LineageObjectType, + pub catalog_type: String, + pub catalog: String, + pub database: String, + pub name: String, + pub id: Option, + pub schema: Option, + pub masked_column_ids: BTreeSet, + pub object_key: String, + pub lookup_keys: BTreeSet, + pub expandable: bool, +} + +impl ResolvedObject { + pub(super) fn qualified_name(&self) -> String { + match self.object_type { + LineageObjectType::Stage => self.name.clone(), + _ => format!("{}.{}.{}", self.catalog, self.database, self.name), + } + } + + pub(super) fn column_by_name(&self, name: &str) -> Option<(String, String)> { + let field = self.schema.as_ref()?.field_with_name(name).ok()?; + Some((field.column_id.to_string(), field.name().to_string())) + } + + pub(super) fn column_by_id(&self, id: &str) -> Option<(String, String)> { + let id = id.parse::().ok()?; + let field = self + .schema + .as_ref()? + .fields() + .iter() + .find(|field| field.column_id == id)?; + Some((field.column_id.to_string(), field.name().to_string())) + } + + pub(super) fn is_column_masked(&self, id: &str) -> bool { + id.parse() + .ok() + .is_some_and(|id| self.masked_column_ids.contains(&id)) + } +} + +pub(super) struct ObjectResolver { + ctx: Arc, + visibility: Arc, + cache: HashMap>, + stages: Option>, +} + +impl ObjectResolver { + pub(super) async fn try_create(ctx: Arc) -> Result { + let visibility = ctx.get_visibility_checker(false, Object::All).await?; + Ok(Self { + ctx, + visibility, + cache: HashMap::new(), + stages: None, + }) + } + + pub(super) async fn resolve_start( + &mut self, + domain: ObjectDomain, + value: &str, + ) -> Result> { + if domain == ObjectDomain::Stage { + return self.resolve_stage(value.trim()).await; + } + + let expected = match domain { + ObjectDomain::Table => LineageObjectType::Table, + ObjectDomain::View => LineageObjectType::View, + ObjectDomain::Column => { + let (catalog, database, table_name) = parse_object_name(&self.ctx, value)?; + if let Some(table) = self + .resolve_table_by_name( + &catalog, + &database, + &table_name, + LineageObjectType::Table, + None, + ) + .await? + { + return Ok(Some(table)); + } + return self + .resolve_table_by_name( + &catalog, + &database, + &table_name, + LineageObjectType::View, + None, + ) + .await; + } + ObjectDomain::Stage => unreachable!(), + }; + let (catalog, database, table_name) = parse_object_name(&self.ctx, value)?; + self.resolve_table_by_name(&catalog, &database, &table_name, expected, None) + .await + } + + pub(super) async fn resolve( + &mut self, + captured: &CapturedObject, + ) -> Result> { + if let Some(resolved) = self.cache.get(captured) { + return Ok(resolved.clone()); + } + + let resolved = match captured.object_type { + LineageObjectType::Stage => self.resolve_stage(&captured.name).await?, + _ if !captured.is_default_catalog() => { + // External-catalog objects are name-addressed in v1. Resolve their current display + // name when possible, but keep them terminal because their stable-ID semantics are + // catalog-specific. + self.resolve_table_by_name( + &captured.catalog, + &captured.database, + &captured.name, + captured.object_type, + Some(captured), + ) + .await? + .map(|mut object| { + object.expandable = false; + object + }) + } + _ => match captured.address_kind { + AddressKind::Id => self.resolve_table_by_id(captured).await?, + AddressKind::Name => { + self.resolve_table_by_name( + &captured.catalog, + &captured.database, + &captured.name, + captured.object_type, + Some(captured), + ) + .await? + } + }, + }; + self.cache.insert(captured.clone(), resolved.clone()); + Ok(resolved) + } + + async fn resolve_table_by_id( + &self, + captured: &CapturedObject, + ) -> Result> { + let Some(table_id) = captured.id else { + return Ok(None); + }; + let meta = UserApiProvider::instance().get_meta_store_client(); + let Some(name_entry) = meta + .get_pb(&TableIdToName { table_id }) + .await + .map_err(meta_service_error)? + else { + return Ok(None); + }; + let catalog = self.ctx.get_catalog(CATALOG_DEFAULT).await?; + let database = match catalog.get_db_name_by_id(name_entry.data.db_id).await { + Ok(database) => database, + Err(error) if error.code() == ErrorCode::UNKNOWN_DATABASE_ID => return Ok(None), + Err(error) => return Err(error), + }; + self.resolve_table_by_name( + CATALOG_DEFAULT, + &database, + &name_entry.data.table_name, + captured.object_type, + Some(captured), + ) + .await + .map(|object| object.filter(|object| object.id == Some(table_id))) + } + + async fn resolve_table_by_name( + &self, + catalog_name: &str, + database_name: &str, + table_name: &str, + expected: LineageObjectType, + captured: Option<&CapturedObject>, + ) -> Result> { + let table = match self + .ctx + .get_table(catalog_name, database_name, table_name) + .await + { + Ok(table) => table, + Err(error) + if matches!( + error.code(), + ErrorCode::UNKNOWN_TABLE | ErrorCode::UNKNOWN_DATABASE + ) => + { + return Ok(None); + } + Err(error) => return Err(error), + }; + if !matches_expected_type(table.as_ref(), expected) { + return Ok(None); + } + + let table_id = table.get_id(); + let masked_column_ids = table + .get_table_info() + .meta + .column_mask_policy_columns_ids + .keys() + .copied() + .collect(); + let catalog = self.ctx.get_catalog(catalog_name).await?; + let is_default = catalog_name.eq_ignore_ascii_case(CATALOG_DEFAULT); + if is_default { + let database = match catalog + .get_database(&self.ctx.get_tenant(), database_name) + .await + { + Ok(database) => database, + Err(error) if error.code() == ErrorCode::UNKNOWN_DATABASE => return Ok(None), + Err(error) => return Err(error), + }; + let db_id = database.get_db_info().database_id.db_id; + if !self.visibility.check_table_visibility( + catalog_name, + database_name, + table_name, + db_id, + table_id, + ) { + return Ok(None); + } + } + + let address_kind = captured + .map(|captured| captured.address_kind) + .unwrap_or(AddressKind::Id); + let object_type = expected; + let id = Some(table_id); + let lookup_keys = + object_lookup_keys(object_type, catalog_name, database_name, table_name, id); + let object_key = object_key( + object_type, + address_kind, + catalog_name, + database_name, + table_name, + id, + ); + let schema = if object_type == LineageObjectType::View { + let Some(query) = table.options().get(QUERY) else { + warn!( + "Skipping lineage view without stored query: {}.{}.{}", + catalog_name, database_name, table_name + ); + return Ok(None); + }; + let mut planner = Planner::new(self.ctx.clone()); + let (plan, _) = match planner.plan_sql(query).await { + Ok(result) => result, + Err(error) => { + warn!( + "Skipping lineage view whose query cannot be resolved: {}.{}.{}, error: {}", + catalog_name, database_name, table_name, error + ); + return Ok(None); + } + }; + Some(infer_table_schema(&plan.schema())?) + } else { + Some(table.schema()) + }; + Ok(Some(ResolvedObject { + object_type, + catalog_type: captured + .map(|captured| captured.catalog_type.clone()) + .unwrap_or_else(|| { + if is_default { + "DEFAULT".to_string() + } else { + "EXTERNAL".to_string() + } + }), + catalog: catalog_name.to_string(), + database: database_name.to_string(), + name: table_name.to_string(), + id, + schema, + masked_column_ids, + object_key, + lookup_keys, + expandable: is_default, + })) + } + + async fn resolve_stage(&mut self, name: &str) -> Result> { + if name.is_empty() { + return Ok(None); + } + if self.stages.is_none() { + let stages = UserApiProvider::instance() + .get_stages(&self.ctx.get_tenant()) + .await? + .into_iter() + .filter(|stage| self.visibility.check_stage_visibility(&stage.stage_name)) + .map(|stage| stage.stage_name) + .collect(); + self.stages = Some(stages); + } + if !self + .stages + .as_ref() + .is_some_and(|stages| stages.contains(name)) + { + return Ok(None); + } + + let key = format!("STAGE::NAME::{name}"); + Ok(Some(ResolvedObject { + object_type: LineageObjectType::Stage, + catalog_type: "STAGE".to_string(), + catalog: String::new(), + database: String::new(), + name: name.to_string(), + id: None, + schema: None, + masked_column_ids: BTreeSet::new(), + object_key: key.clone(), + lookup_keys: BTreeSet::from([key]), + expandable: true, + })) + } +} + +fn matches_expected_type(table: &dyn Table, expected: LineageObjectType) -> bool { + match expected { + LineageObjectType::View => { + !table.is_temp() && table.engine().eq_ignore_ascii_case(VIEW_ENGINE) + } + LineageObjectType::Table => is_lineage_table_endpoint(table.engine(), table.is_temp()), + LineageObjectType::Stage => false, + } +} + +fn is_lineage_table_endpoint(engine: &str, is_temporary: bool) -> bool { + !is_temporary + && ![VIEW_ENGINE, STREAM_ENGINE, "MEMORY", "DELTA"] + .iter() + .any(|unsupported| engine.eq_ignore_ascii_case(unsupported)) +} + +fn object_lookup_keys( + object_type: LineageObjectType, + catalog: &str, + database: &str, + name: &str, + id: Option, +) -> BTreeSet { + let mut keys = BTreeSet::new(); + if let Some(id) = id { + keys.insert(format!("{}::ID::{id}", object_type.as_str())); + } + keys.insert(format!( + "{}::NAME::{catalog}.{database}.{name}", + object_type.as_str() + )); + keys +} + +fn object_key( + object_type: LineageObjectType, + address_kind: AddressKind, + catalog: &str, + database: &str, + name: &str, + id: Option, +) -> String { + match (address_kind, id) { + (AddressKind::Id, Some(id)) => format!("{}::ID::{id}", object_type.as_str()), + _ => format!( + "{}::NAME::{catalog}.{database}.{name}", + object_type.as_str() + ), + } +} + +pub(super) fn parse_object_name( + ctx: &Arc, + value: &str, +) -> Result<(String, String, String)> { + let value = value.trim(); + if value.is_empty() { + return Err(ErrorCode::BadArguments("object_name must not be empty")); + } + let settings = ctx.get_settings(); + let resolution = NameResolutionContext::try_from(settings.as_ref())?; + let dialect = settings.get_sql_dialect().unwrap_or_default(); + let table_ref = parse_table_ref(value, dialect).map_err(|error| { + ErrorCode::BadArguments(format!("invalid object_name '{value}': {}", error.1)) + })?; + let catalog = table_ref + .catalog + .map(|ident| normalize_identifier(&ident, &resolution).name) + .unwrap_or_else(|| ctx.get_current_catalog()); + let database = table_ref + .database + .map(|ident| normalize_identifier(&ident, &resolution).name) + .unwrap_or_else(|| ctx.get_current_database()); + let table = normalize_identifier(&table_ref.table, &resolution).name; + Ok((catalog, database, table)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_lineage_table_endpoint_engines() { + for engine in ["FUSE", "HIVE", "ICEBERG", "PAIMON"] { + assert!(is_lineage_table_endpoint(engine, false), "engine={engine}"); + } + for engine in [VIEW_ENGINE, STREAM_ENGINE, "MEMORY", "DELTA"] { + assert!(!is_lineage_table_endpoint(engine, false), "engine={engine}"); + } + assert!(!is_lineage_table_endpoint("FUSE", true)); + } +} diff --git a/src/query/service/src/table_functions/get_lineage/traversal.rs b/src/query/service/src/table_functions/get_lineage/traversal.rs new file mode 100644 index 0000000000000..069fb89382227 --- /dev/null +++ b/src/query/service/src/table_functions/get_lineage/traversal.rs @@ -0,0 +1,593 @@ +// 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::BTreeSet; +use std::collections::HashMap; +use std::sync::Arc; + +use databend_common_ast::parser::parse_table_ref; +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; +use databend_common_expression::utils::date_helper::DateConverter; +use databend_common_sql::planner::NameResolutionContext; +use databend_common_sql::planner::normalize_identifier; +use jiff::fmt::strtime; +use jiff::tz::TimeZone; + +use super::GetLineageArgs; +use super::LineageResultRow; +use super::ObjectDomain; +use super::QueryDirection; +use super::edge_reader::AddressKind; +use super::edge_reader::CapturedObject; +use super::edge_reader::LineageEdgeReader; +use super::edge_reader::LineageObjectType; +use super::edge_reader::RawLineageEdge; +use super::resolver::ObjectResolver; +use super::resolver::ResolvedObject; +use crate::sessions::TableContext; + +#[derive(Clone)] +struct TraversedObjectEdge { + distance: u8, + edge: RawLineageEdge, + source: ResolvedObject, + target: ResolvedObject, +} + +#[derive(Clone)] +struct ColumnFrontier { + object: ResolvedObject, + column_id: String, + column_name: String, +} + +#[derive(Clone)] +struct TraversedColumnEdge { + distance: u8, + edge: RawLineageEdge, + source: ResolvedObject, + source_column: String, + target: ResolvedObject, + target_column: String, + target_masked: bool, +} + +pub(super) async fn traverse( + ctx: Arc, + args: GetLineageArgs, +) -> Result> { + let mut reader = LineageEdgeReader::try_create(ctx.clone()).await?; + let mut resolver = ObjectResolver::try_create(ctx.clone()).await?; + match args.object_domain { + ObjectDomain::Column => traverse_columns(ctx, &mut reader, &mut resolver, args).await, + _ => traverse_objects(&mut reader, &mut resolver, args).await, + } +} + +async fn traverse_objects( + reader: &mut LineageEdgeReader, + resolver: &mut ObjectResolver, + args: GetLineageArgs, +) -> Result> { + let Some(start) = resolver + .resolve_start(args.object_domain, &args.object_name) + .await? + else { + return Ok(vec![]); + }; + let mut frontier = start.lookup_keys.clone(); + let mut visited = frontier.clone(); + let mut results: HashMap<(String, String), TraversedObjectEdge> = HashMap::new(); + + for distance in 1..=args.distance { + let raw_edges = reader + .read_frontier(args.direction.match_column(), &frontier) + .await?; + let mut level_edges: HashMap<(String, String), TraversedObjectEdge> = HashMap::new(); + for edge in raw_edges { + let Some(source) = resolver.resolve(&edge.source).await? else { + continue; + }; + let Some(target) = resolver.resolve(&edge.target).await? else { + continue; + }; + if same_object(&source, &target) { + continue; + } + let key = (source.object_key.clone(), target.object_key.clone()); + let candidate = TraversedObjectEdge { + distance, + edge, + source, + target, + }; + match level_edges.get_mut(&key) { + Some(current) if candidate.edge.newer_than(¤t.edge) => *current = candidate, + None => { + level_edges.insert(key, candidate); + } + _ => {} + } + } + + let mut next_frontier = BTreeSet::new(); + for (key, candidate) in level_edges { + match results.get_mut(&key) { + Some(current) + if candidate.distance < current.distance + || (candidate.distance == current.distance + && candidate.edge.newer_than(¤t.edge)) => + { + *current = candidate.clone(); + } + None => { + results.insert(key, candidate.clone()); + } + _ => {} + } + + let next = match args.direction { + QueryDirection::Upstream => &candidate.source, + QueryDirection::Downstream => &candidate.target, + }; + if next.expandable && next.lookup_keys.iter().all(|key| !visited.contains(key)) { + next_frontier.extend(next.lookup_keys.iter().cloned()); + } + } + if next_frontier.is_empty() { + break; + } + visited.extend(next_frontier.iter().cloned()); + frontier = next_frontier; + } + + let mut rows = results + .into_values() + .map(|result| LineageResultRow { + distance: i32::from(result.distance), + source_object_domain: Some(result.source.object_type.as_str().to_string()), + source_object_name: Some(result.source.qualified_name()), + source_column_name: None, + target_object_domain: Some(result.target.object_type.as_str().to_string()), + target_object_name: Some(result.target.qualified_name()), + target_column_name: None, + target_status: "ACTIVE".to_string(), + process: Some(process_json(&result.edge)), + }) + .collect::>(); + sort_rows(&mut rows); + Ok(rows) +} + +async fn traverse_columns( + ctx: Arc, + reader: &mut LineageEdgeReader, + resolver: &mut ObjectResolver, + args: GetLineageArgs, +) -> Result> { + let (table_name, column_name) = split_column_name(&args.object_name)?; + let Some(start) = resolver + .resolve_start(ObjectDomain::Column, &table_name) + .await? + else { + return Ok(vec![]); + }; + let column_name = normalize_column_name(&ctx, &column_name)?; + let Some((column_id, column_name)) = start.column_by_name(&column_name) else { + return Ok(vec![]); + }; + + let mut frontier = vec![ColumnFrontier { + object: start, + column_id, + column_name, + }]; + let mut visited = BTreeSet::new(); + visited.insert(column_frontier_key(&frontier[0])); + let mut results: HashMap<(String, String, String, String), TraversedColumnEdge> = + HashMap::new(); + + for distance in 1..=args.distance { + let lookup_keys = frontier + .iter() + .flat_map(|column| column.object.lookup_keys.iter().cloned()) + .collect::>(); + let raw_edges = reader + .read_frontier(args.direction.match_column(), &lookup_keys) + .await?; + let mut level_edges: HashMap<(String, String, String, String), TraversedColumnEdge> = + HashMap::new(); + let mut next_columns = Vec::new(); + + for edge in raw_edges { + let Some(source) = resolver.resolve(&edge.source).await? else { + continue; + }; + let Some(target) = resolver.resolve(&edge.target).await? else { + continue; + }; + if same_object(&source, &target) { + continue; + } + + let (current_captured, current_object, next_captured, next_object, column_map) = + match args.direction { + QueryDirection::Upstream => ( + &edge.target, + &target, + &edge.source, + &source, + &edge.target_to_source_columns, + ), + QueryDirection::Downstream => ( + &edge.source, + &source, + &edge.target, + &target, + &edge.source_to_target_columns, + ), + }; + + for current_column in frontier.iter().filter(|column| { + column + .object + .lookup_keys + .contains(¤t_captured.lineage_key) + }) { + let current_ref = match column_address_kind(current_captured) { + AddressKind::Id => ¤t_column.column_id, + AddressKind::Name => ¤t_column.column_name, + }; + let Some(mapped_refs) = column_map.get(current_ref) else { + continue; + }; + for mapped_ref in mapped_refs { + let Some((next_id, next_name)) = + resolve_column_ref(next_captured, next_object, mapped_ref) + else { + continue; + }; + let (source_column, target_column) = match args.direction { + QueryDirection::Upstream => { + (next_name.clone(), current_column.column_name.clone()) + } + QueryDirection::Downstream => { + (current_column.column_name.clone(), next_name.clone()) + } + }; + let target_masked = match args.direction { + QueryDirection::Upstream => { + target.is_column_masked(¤t_column.column_id) + } + QueryDirection::Downstream => target.is_column_masked(&next_id), + }; + let key = ( + source.object_key.clone(), + source_column.clone(), + target.object_key.clone(), + target_column.clone(), + ); + let candidate = TraversedColumnEdge { + distance, + edge: edge.clone(), + source: source.clone(), + source_column, + target: target.clone(), + target_column, + target_masked, + }; + match level_edges.get_mut(&key) { + Some(current) if candidate.edge.newer_than(¤t.edge) => { + *current = candidate + } + None => { + level_edges.insert(key, candidate); + } + _ => {} + } + if next_object.expandable { + next_columns.push(ColumnFrontier { + object: next_object.clone(), + column_id: next_id, + column_name: next_name, + }); + } + } + } + + // The current endpoint must resolve to the same active object represented by the + // frontier. This also prevents a stale name-addressed edge from crossing a replace. + let _ = current_object; + } + + for (key, candidate) in level_edges { + match results.get_mut(&key) { + Some(current) + if candidate.distance < current.distance + || (candidate.distance == current.distance + && candidate.edge.newer_than(¤t.edge)) => + { + *current = candidate; + } + None => { + results.insert(key, candidate); + } + _ => {} + } + } + + let mut deduplicated = HashMap::new(); + for column in next_columns { + let key = column_frontier_key(&column); + if !visited.contains(&key) { + deduplicated.entry(key).or_insert(column); + } + } + if deduplicated.is_empty() { + break; + } + visited.extend(deduplicated.keys().cloned()); + frontier = deduplicated.into_values().collect(); + } + + let mut rows = results + .into_values() + .map(|result| LineageResultRow { + distance: i32::from(result.distance), + source_object_domain: Some(result.source.object_type.as_str().to_string()), + source_object_name: Some(result.source.qualified_name()), + source_column_name: Some(result.source_column), + target_object_domain: Some(result.target.object_type.as_str().to_string()), + target_object_name: Some(result.target.qualified_name()), + target_column_name: Some(result.target_column), + target_status: target_status(result.target_masked).to_string(), + process: Some(process_json(&result.edge)), + }) + .collect::>(); + sort_rows(&mut rows); + Ok(rows) +} + +fn column_address_kind(object: &CapturedObject) -> AddressKind { + if object.object_type == LineageObjectType::View { + AddressKind::Name + } else { + object.address_kind + } +} + +fn resolve_column_ref( + captured: &CapturedObject, + object: &ResolvedObject, + column_ref: &str, +) -> Option<(String, String)> { + if !captured.is_default_catalog() { + return Some((column_ref.to_string(), column_ref.to_string())); + } + match column_address_kind(captured) { + AddressKind::Id => object.column_by_id(column_ref), + AddressKind::Name => object.column_by_name(column_ref), + } +} + +fn same_object(source: &ResolvedObject, target: &ResolvedObject) -> bool { + match (source.id, target.id) { + (Some(source_id), Some(target_id)) + if source.catalog_type.eq_ignore_ascii_case("DEFAULT") + && target.catalog_type.eq_ignore_ascii_case("DEFAULT") => + { + source_id == target_id + } + _ => { + source.object_type == target.object_type + && source.catalog == target.catalog + && source.database == target.database + && source.name == target.name + } + } +} + +fn column_frontier_key(column: &ColumnFrontier) -> (String, String) { + let column_key = if column.object.object_type == LineageObjectType::View { + column.column_name.clone() + } else { + column.column_id.clone() + }; + ( + column + .object + .lookup_keys + .iter() + .next() + .cloned() + .unwrap_or_else(|| column.object.object_key.clone()), + column_key, + ) +} + +fn process_json(edge: &RawLineageEdge) -> String { + // Keep embedded JSON timestamps stable across session timezones and self-describing. + let event_time = edge.event_time.map(|timestamp| { + strtime::format( + "%Y-%m-%dT%H:%M:%S%.6fZ", + ×tamp.to_timestamp(&TimeZone::UTC), + ) + .expect("valid lineage process timestamp format") + .to_string() + }); + serde_json::json!({ + "query_id": edge.query_id, + "query_kind": edge.query_kind, + "lineage_kind": edge.lineage_kind, + "event_time": event_time, + }) + .to_string() +} + +fn target_status(masked: bool) -> &'static str { + if masked { "MASKED" } else { "ACTIVE" } +} + +fn sort_rows(rows: &mut [LineageResultRow]) { + rows.sort_by(|left, right| { + ( + left.distance, + left.source_object_name.as_deref().unwrap_or_default(), + left.source_column_name.as_deref().unwrap_or_default(), + left.target_object_name.as_deref().unwrap_or_default(), + left.target_column_name.as_deref().unwrap_or_default(), + ) + .cmp(&( + right.distance, + right.source_object_name.as_deref().unwrap_or_default(), + right.source_column_name.as_deref().unwrap_or_default(), + right.target_object_name.as_deref().unwrap_or_default(), + right.target_column_name.as_deref().unwrap_or_default(), + )) + }); +} + +fn split_column_name(input: &str) -> Result<(String, String)> { + let input = input.trim(); + let Some(index) = last_unquoted_dot(input) else { + return Err(ErrorCode::BadArguments( + "column object_name must be qualified by table name", + )); + }; + let table = input[..index].trim(); + let column = input[index + 1..].trim(); + if table.is_empty() || column.is_empty() { + return Err(ErrorCode::BadArguments( + "column object_name must be in table.column format", + )); + } + Ok((table.to_string(), column.to_string())) +} + +fn normalize_column_name(ctx: &Arc, value: &str) -> Result { + let settings = ctx.get_settings(); + let resolution = NameResolutionContext::try_from(settings.as_ref())?; + let dialect = settings.get_sql_dialect().unwrap_or_default(); + let column_ref = parse_table_ref(value, dialect).map_err(|error| { + ErrorCode::BadArguments(format!("invalid column name '{value}': {}", error.1)) + })?; + if column_ref.catalog.is_some() || column_ref.database.is_some() { + return Err(ErrorCode::BadArguments(format!( + "invalid column name '{value}'" + ))); + } + Ok(normalize_identifier(&column_ref.table, &resolution).name) +} + +fn last_unquoted_dot(input: &str) -> Option { + let mut in_quote = false; + let mut last_dot = None; + let mut iter = input.char_indices().peekable(); + while let Some((index, ch)) = iter.next() { + match ch { + '"' => { + if in_quote && matches!(iter.peek(), Some((_, '"'))) { + iter.next(); + } else { + in_quote = !in_quote; + } + } + '.' if !in_quote => last_dot = Some(index), + _ => {} + } + } + last_dot +} + +#[cfg(test)] +mod tests { + use super::*; + + fn test_edge( + query_id: Option<&str>, + event_time: Option, + query_kind: Option<&str>, + lineage_kind: Option<&str>, + ) -> RawLineageEdge { + let endpoint = CapturedObject { + lineage_key: "TABLE::ID::1".to_string(), + address_kind: AddressKind::Id, + catalog_type: "DEFAULT".to_string(), + object_type: LineageObjectType::Table, + catalog: "default".to_string(), + database: "default".to_string(), + name: "t".to_string(), + id: Some(1), + }; + RawLineageEdge { + query_id: query_id.map(str::to_string), + event_time, + query_kind: query_kind.map(str::to_string), + lineage_kind: lineage_kind.map(str::to_string), + column_lineage_hash: String::new(), + source: endpoint.clone(), + target: endpoint, + source_to_target_columns: Default::default(), + target_to_source_columns: Default::default(), + } + } + + #[test] + fn test_process_json() { + let edge = test_edge(Some("query-1"), Some(0), Some("Insert"), Some("DML")); + let process: serde_json::Value = + serde_json::from_str(&process_json(&edge)).expect("valid process JSON"); + assert_eq!( + process, + serde_json::json!({ + "query_id": "query-1", + "query_kind": "Insert", + "lineage_kind": "DML", + "event_time": "1970-01-01T00:00:00.000000Z", + }) + ); + + let null_process: serde_json::Value = + serde_json::from_str(&process_json(&test_edge(None, None, None, None))) + .expect("valid process JSON"); + assert_eq!( + null_process, + serde_json::json!({ + "query_id": null, + "query_kind": null, + "lineage_kind": null, + "event_time": null, + }) + ); + } + + #[test] + fn test_split_quoted_column_name() -> Result<()> { + assert_eq!( + split_column_name(r#"db."table.with.dot"."column.with.dot""#)?, + ( + r#"db."table.with.dot""#.to_string(), + r#""column.with.dot""#.to_string(), + ) + ); + Ok(()) + } + + #[test] + fn test_target_status() { + assert_eq!(target_status(false), "ACTIVE"); + assert_eq!(target_status(true), "MASKED"); + } +} diff --git a/src/query/service/src/table_functions/mod.rs b/src/query/service/src/table_functions/mod.rs index 0f490612c5453..dc228fd094e2b 100644 --- a/src/query/service/src/table_functions/mod.rs +++ b/src/query/service/src/table_functions/mod.rs @@ -16,6 +16,7 @@ mod async_crash_me; mod billing_usage_daily; mod copy_history; mod fuse_vacuum2; +mod get_lineage; #[cfg(feature = "storage-stage")] pub(crate) mod infer_schema; mod inspect_parquet; @@ -42,6 +43,7 @@ mod udf_table; pub use billing_usage_daily::BillingUsageDailyTable; pub use copy_history::CopyHistoryTable; +pub use get_lineage::GetLineageTable; pub use numbers::NumbersPartInfo; pub use numbers::NumbersTable; pub use numbers::generate_numbers_parts; diff --git a/src/query/service/src/table_functions/table_function_factory.rs b/src/query/service/src/table_functions/table_function_factory.rs index 1a2f3bb7eabcf..4fb12e0eb0e26 100644 --- a/src/query/service/src/table_functions/table_function_factory.rs +++ b/src/query/service/src/table_functions/table_function_factory.rs @@ -66,6 +66,7 @@ use crate::table_functions::TableFunction; use crate::table_functions::async_crash_me::AsyncCrashMeTable; use crate::table_functions::copy_history::CopyHistoryTable; use crate::table_functions::fuse_vacuum2::FuseVacuum2Table; +use crate::table_functions::get_lineage::GetLineageTable; #[cfg(feature = "storage-stage")] use crate::table_functions::infer_schema::InferSchemaTable; use crate::table_functions::inspect_parquet::InspectParquetTable; @@ -453,6 +454,10 @@ impl TableFunctionFactory { "stream_backlog".to_string(), (next_id(), Arc::new(StreamBacklogTable::create)), ); + creators.insert( + "get_lineage".to_string(), + (next_id(), Arc::new(GetLineageTable::create)), + ); TableFunctionFactory { creators: RwLock::new(creators), 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..8c7f67070bf0c 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,12 +27,16 @@ 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::lineage_enabled; use crate::binder::util::TableIdentifier; use crate::binder::util::legacy_table_ref_removed_error; use crate::optimizer::ir::SExpr; @@ -189,17 +193,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 +292,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 +314,15 @@ impl Binder { column.table_name = Some(self.normalize_identifier(table).name); } } + if lineage_enabled() { + 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 +364,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..e43de7d42a900 100644 --- a/src/query/sql/src/planner/binder/ddl/view.rs +++ b/src/query/sql/src/planner/binder/ddl/view.rs @@ -18,6 +18,7 @@ use databend_common_ast::ast::DescribeViewStmt; use databend_common_ast::ast::DropViewStmt; use databend_common_ast::ast::ShowLimit; use databend_common_ast::ast::ShowViewsStmt; +use databend_common_ast::ast::Statement; use databend_common_ast::ast::quote::QuotedIdent; use databend_common_ast::ast::quote::QuotedString; use databend_common_ast::visit::WalkMut; @@ -26,11 +27,13 @@ use databend_common_expression::DataField; use databend_common_expression::DataSchemaRefExt; use databend_common_expression::types::DataType; use log::debug; +use log::warn; use crate::BindContext; use crate::SelectBuilder; use crate::ViewRewriter; use crate::binder::Binder; +use crate::binder::lineage_enabled; use crate::planner::semantic::normalize_identifier; use crate::plans::CreateViewPlan; use crate::plans::DescribeViewPlan; @@ -65,6 +68,17 @@ impl Binder { }; query.walk_mut(&mut visitor)?; let subquery = format!("{}", query); + let query_plan = if lineage_enabled() { + match self.view_query_plan(&query).await { + Ok(plan) => Some(Box::new(plan)), + Err(error) => { + warn!("Failed to bind CREATE VIEW query for lineage: {:?}", error); + None + } + } + } else { + None + }; let plan = CreateViewPlan { create_option: create_option.clone().into(), @@ -74,6 +88,7 @@ impl Binder { view_name, column_names, subquery, + query_plan, }; Ok(Plan::CreateView(plan.into())) } @@ -225,4 +240,9 @@ impl Binder { schema, }))) } + async fn view_query_plan(&mut self, query: &databend_common_ast::ast::Query) -> Result { + let stmt = Statement::Query(Box::new(query.clone())); + let mut bind_context = BindContext::new(); + self.bind_statement(&mut bind_context, &stmt).await + } } 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/binder/mod.rs b/src/query/sql/src/planner/binder/mod.rs index 41f6dc35e6517..3cd8e5ad025b5 100644 --- a/src/query/sql/src/planner/binder/mod.rs +++ b/src/query/sql/src/planner/binder/mod.rs @@ -79,6 +79,7 @@ pub use column_binding::ColumnBinding; pub use column_binding::ColumnBindingBuilder; pub use constraint_expr::ConstraintExprBinder; pub use constraint_expr::validate_constraints_by_schema; +use databend_common_config::GlobalConfig; pub use databend_common_expression::DummyColumnType; pub use ddl::database::DEFAULT_STORAGE_CONNECTION; pub use ddl::database::DEFAULT_STORAGE_PATH; @@ -105,3 +106,10 @@ pub use stream_column_factory::STREAM_COLUMN_FACTORY; pub use window::WindowFunctionInfo; pub use window::WindowOrderByInfo; pub use window::bind_window_function_info; + +pub(crate) fn lineage_enabled() -> bool { + GlobalConfig::instance() + .log + .history + .is_table_enabled("lineage_unresolved") +} 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..470b8a2e11281 --- /dev/null +++ b/src/query/sql/tests/it/planner/lineage.rs @@ -0,0 +1,590 @@ +// 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_config::InnerConfig; +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 databend_common_sql_test_support::init_testing_globals_with_config; +use databend_common_tracing::HistoryTableConfig; + +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> { + let mut config = InnerConfig::default(); + config.log.history.on = true; + config.log.history.tables.push(HistoryTableConfig { + table_name: "lineage_unresolved".to_string(), + retention: None, + invisible: false, + }); + // Lite globals are thread-local in debug builds. Initialize capture on the test's current + // thread before LiteTableContext::create() attempts to install the default configuration. + init_testing_globals_with_config(config); + 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) } diff --git a/tests/logging/history_table/check_logs_table.sh b/tests/logging/history_table/check_logs_table.sh index 40ddb96cd8f05..c553a3abca236 100755 --- a/tests/logging/history_table/check_logs_table.sh +++ b/tests/logging/history_table/check_logs_table.sh @@ -2,6 +2,9 @@ set -e +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "$SCRIPT_DIR/lineage_sqllogic.sh" + # Source query IDs from setup script if [ -z "$QUERY_ID" ]; then echo "Error: Query IDs not set. Run setup_test_data.sh first." @@ -71,6 +74,9 @@ check_query_log "basic-10" "$CREATE_VIEW_QUERY_ID" "select query_text from syste check_query_log "basic-11" null "SELECT count(*) FROM system_history.login_history WHERE session_id = '$SELECT_SESSION_ID' " "1" +# Lineage assertions mirror the grouped setup suite and clean up its dedicated database. +run_lineage_suite check + check_query_log "resource-1" "$SELECT_QUERY_ID" "SELECT count(*) FROM system_history.query_history WHERE resource_usage IS NOT NULL AND resource_usage['read_count']::UInt64 > 0 AND resource_usage['read_bytes']::UInt64 > 0 AND resource_usage['read_duration_ms'] IS NOT NULL AND resource_usage['cpu_time_ms'] IS NOT NULL AND resource_usage['wait_time_ms'] IS NOT NULL AND" "1" check_query_log "resource-2" "$INSERT_QUERY_ID" "SELECT count(*) FROM system_history.query_history WHERE resource_usage IS NOT NULL AND resource_usage['write_count']::UInt64 > 0 AND resource_usage['write_bytes']::UInt64 > 0 AND resource_usage['write_duration_ms'] IS NOT NULL AND resource_usage['cpu_time_ms'] IS NOT NULL AND resource_usage['wait_time_ms'] IS NOT NULL AND" "1" diff --git a/tests/logging/history_table/history_table.toml b/tests/logging/history_table/history_table.toml index ed04547f2fcf8..90c5167269e8e 100644 --- a/tests/logging/history_table/history_table.toml +++ b/tests/logging/history_table/history_table.toml @@ -13,3 +13,6 @@ table_name = "login_history" [[log.history.tables]] table_name = "access_history" + +[[log.history.tables]] +table_name = "lineage_unresolved" diff --git a/tests/logging/history_table/history_table_external.toml b/tests/logging/history_table/history_table_external.toml index 48289f10aff47..9b57ef07245bd 100644 --- a/tests/logging/history_table/history_table_external.toml +++ b/tests/logging/history_table/history_table_external.toml @@ -23,3 +23,6 @@ table_name = "login_history" [[log.history.tables]] table_name = "access_history" + +[[log.history.tables]] +table_name = "lineage_unresolved" diff --git a/tests/logging/history_table/lineage_sqllogic.sh b/tests/logging/history_table/lineage_sqllogic.sh new file mode 100644 index 0000000000000..8884c7bb73a26 --- /dev/null +++ b/tests/logging/history_table/lineage_sqllogic.sh @@ -0,0 +1,18 @@ +#!/bin/bash + +# Shared sqllogictest runner for lineage history setup and assertions. +LINEAGE_HISTORY_SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +LINEAGE_HISTORY_REPO_ROOT="$(cd "$LINEAGE_HISTORY_SCRIPT_DIR/../../.." && pwd)" +BUILD_PROFILE="${BUILD_PROFILE:-debug}" + +run_lineage_suite() { + local suite_dir="$1" + ( + cd "$LINEAGE_HISTORY_REPO_ROOT" + "target/${BUILD_PROFILE}/databend-sqllogictests" \ + --suites tests/logging/history_table/sqllogic \ + --run_dir "$suite_dir" \ + --handlers http \ + --parallel 1 + ) +} diff --git a/tests/logging/history_table/setup_test_data.sh b/tests/logging/history_table/setup_test_data.sh index 360b5a7711d79..7fc238e79fecc 100644 --- a/tests/logging/history_table/setup_test_data.sh +++ b/tests/logging/history_table/setup_test_data.sh @@ -2,6 +2,9 @@ set -e +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "$SCRIPT_DIR/lineage_sqllogic.sh" + execute_query() { local sql="$1" local extra_headers="$2" @@ -45,6 +48,10 @@ select_session_id=$(echo $response | jq -r '.session_id') echo "Select Query ID: $select_query_id" echo "Select Session ID: $select_session_id" +# Lineage setup is grouped in a dedicated sqllogictest suite so scenarios remain reviewable and +# can grow without adding SQL assertions to this shell script. +run_lineage_suite setup + execute_query_silent "drop user if exists wrong_pass_user" execute_query_silent "create user wrong_pass_user identified by 'secure_password'" @@ -58,6 +65,34 @@ for _ in {1..3}; do sleep 3 done +# History ingestion is asynchronous. Poll the transformed edge table instead of assuming the fixed +# delay above is sufficient on every CI runner. +lineage_ready=false +for _ in {1..30}; do + lineage_response=$(execute_query "SELECT count_if(source_database IN ('lineage_history_objects', 'lineage_history_columns', 'lineage_history_views', 'lineage_history_statements') OR target_database IN ('lineage_history_objects', 'lineage_history_columns', 'lineage_history_views', 'lineage_history_statements')) AS captured_edges, count_if(target_database = 'lineage_history_lifecycle') AS lifecycle_edges, count_if(source_catalog = 'lineage_history_iceberg_catalog' AND source_database = 'lineage_db' AND target_database = 'lineage_history_iceberg') AS iceberg_edges, count_if(target_database = 'lineage_history_views' AND target_name IN ('src_view', 'view_dst')) AS view_edges FROM system_history.lineage_unresolved") + if [ "$(echo "$lineage_response" | jq -r '.state')" = "Failed" ]; then + # The history transform creates its destination table asynchronously. Treat a missing table + # like any other not-ready state and report the last response if the poll eventually times out. + sleep 1 + continue + fi + lineage_count=$(echo "$lineage_response" | jq -r '.data[0][0] // 0') + lifecycle_count=$(echo "$lineage_response" | jq -r '.data[0][1] // 0') + iceberg_count=$(echo "$lineage_response" | jq -r '.data[0][2] // 0') + view_count=$(echo "$lineage_response" | jq -r '.data[0][3] // 0') + if [ "$lineage_count" -ge 16 ] 2>/dev/null && [ "$lifecycle_count" -eq 3 ] 2>/dev/null && [ "$iceberg_count" -ge 1 ] 2>/dev/null && [ "$view_count" -ge 2 ] 2>/dev/null; then + lineage_ready=true + break + fi + sleep 1 +done + +if [ "$lineage_ready" = false ]; then + echo "Lineage history was not transformed within 30 seconds" + echo "$lineage_response" + exit 1 +fi + # Export query IDs for use in other scripts export QUERY_ID="$query_id" export CREATE_QUERY_ID="$create_query_id" diff --git a/tests/logging/history_table/sqllogic/check/01_object_patterns.test b/tests/logging/history_table/sqllogic/check/01_object_patterns.test new file mode 100644 index 0000000000000..04018da1d13e2 --- /dev/null +++ b/tests/logging/history_table/sqllogic/check/01_object_patterns.test @@ -0,0 +1,22 @@ +# Object patterns correspond to setup/01_object_patterns.test. + +query I +SELECT count(*) +FROM system_history.lineage_unresolved +WHERE source_database = 'lineage_history_objects' + AND target_database = 'lineage_history_objects' + AND source_name = 'src' + AND target_name = 'mid' + AND lineage_kind IN ('CTAS', 'DML') +---- +2 + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_history_objects.mid', 'TABLE', 'UPSTREAM', 1) +ORDER BY distance, source_object_name, target_object_name +---- +1 default.lineage_history_objects.src default.lineage_history_objects.mid + +statement ok +DROP DATABASE lineage_history_objects diff --git a/tests/logging/history_table/sqllogic/check/02_column_multihop.test b/tests/logging/history_table/sqllogic/check/02_column_multihop.test new file mode 100644 index 0000000000000..6975c11502c4f --- /dev/null +++ b/tests/logging/history_table/sqllogic/check/02_column_multihop.test @@ -0,0 +1,48 @@ +# Column patterns and multi-hop traversal correspond to setup/02_column_multihop.test. + +query ITTTT rowsort +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_history_columns.mid.x', 'COLUMN', 'UPSTREAM', 1) +---- +1 default.lineage_history_columns.src a default.lineage_history_columns.mid x +1 default.lineage_history_columns.src b default.lineage_history_columns.mid x + +query ITTTT rowsort +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_history_columns.dst.z', 'COLUMN', 'UPSTREAM', 2) +---- +1 default.lineage_history_columns.mid x default.lineage_history_columns.dst z +2 default.lineage_history_columns.src a default.lineage_history_columns.mid x +2 default.lineage_history_columns.src b default.lineage_history_columns.mid x + +# Masking applies to a concrete target column, not to the whole object or traversal path. +query ITTTTT rowsort +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name, target_status +FROM get_lineage('lineage_history_columns.dst.z', 'COLUMN', 'UPSTREAM', 2) +---- +1 default.lineage_history_columns.mid x default.lineage_history_columns.dst z MASKED +2 default.lineage_history_columns.src a default.lineage_history_columns.mid x ACTIVE +2 default.lineage_history_columns.src b default.lineage_history_columns.mid x ACTIVE + +query T +SELECT target_status +FROM get_lineage('lineage_history_columns.dst', 'TABLE', 'UPSTREAM', 1) +---- +ACTIVE + +query ITTTT rowsort +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_history_columns.src.a', 'COLUMN', 'DOWNSTREAM', 2) +---- +1 default.lineage_history_columns.src a default.lineage_history_columns.mid x +1 default.lineage_history_columns.src a default.lineage_history_columns.mid y +2 default.lineage_history_columns.mid x default.lineage_history_columns.dst z + +statement ok +ALTER TABLE lineage_history_columns.dst MODIFY COLUMN z UNSET MASKING POLICY + +statement ok +DROP MASKING POLICY lineage_history_columns_mask + +statement ok +DROP DATABASE lineage_history_columns diff --git a/tests/logging/history_table/sqllogic/check/03_view_boundary.test b/tests/logging/history_table/sqllogic/check/03_view_boundary.test new file mode 100644 index 0000000000000..a337b5e8e3eda --- /dev/null +++ b/tests/logging/history_table/sqllogic/check/03_view_boundary.test @@ -0,0 +1,12 @@ +# View traversal corresponds to setup/03_view_boundary.test. + +query ITTTT rowsort +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_history_views.view_dst.z', 'COLUMN', 'UPSTREAM', 2) +---- +1 default.lineage_history_views.src_view v default.lineage_history_views.view_dst z +2 default.lineage_history_views.src a default.lineage_history_views.src_view v +2 default.lineage_history_views.src b default.lineage_history_views.src_view v + +statement ok +DROP DATABASE lineage_history_views diff --git a/tests/logging/history_table/sqllogic/check/04_statement_coverage.test b/tests/logging/history_table/sqllogic/check/04_statement_coverage.test new file mode 100644 index 0000000000000..7a6635d6eaa97 --- /dev/null +++ b/tests/logging/history_table/sqllogic/check/04_statement_coverage.test @@ -0,0 +1,170 @@ +# Statement coverage corresponds to setup/04_statement_coverage.test. + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_history_statements.self_insert_dst', 'TABLE', 'UPSTREAM', 1) +---- +1 default.lineage_history_statements.src default.lineage_history_statements.self_insert_dst + +query ISS +SELECT distance, source_column_name, target_column_name +FROM get_lineage('lineage_history_statements.self_insert_dst.a', 'COLUMN', 'UPSTREAM', 1) +---- +1 a a + +query I +SELECT count(*) +FROM system_history.lineage_unresolved +WHERE target_database = 'lineage_history_statements' + AND source_lineage_key = target_lineage_key +---- +0 + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_history_statements.insert_dst', 'TABLE', 'UPSTREAM', 1) +---- +1 default.lineage_history_statements.src default.lineage_history_statements.insert_dst + +# Process metadata matches the selected history edge. Its timestamp is stable UTC RFC3339, +# independent of the session timezone used to query GET_LINEAGE. +query BBBB +SETTINGS (timezone = 'UTC') +WITH result AS ( + SELECT parse_json(process) AS process + FROM get_lineage('lineage_history_statements.insert_dst', 'TABLE', 'UPSTREAM', 1) +), history AS ( + SELECT query_id, query_kind, lineage_kind, event_time + FROM system_history.lineage_unresolved + WHERE source_database = 'lineage_history_statements' + AND source_name = 'src' + AND target_database = 'lineage_history_statements' + AND target_name = 'insert_dst' + QUALIFY row_number() OVER (ORDER BY event_time DESC, query_id DESC) = 1 +) +SELECT result.process['query_id']::STRING = history.query_id, + result.process['query_kind']::STRING = history.query_kind, + result.process['lineage_kind']::STRING = history.lineage_kind, + result.process['event_time']::STRING = + concat(replace(to_string(history.event_time), ' ', 'T'), 'Z') +FROM result CROSS JOIN history +---- +1 1 1 1 + +query B +SETTINGS (timezone = 'Asia/Shanghai') +WITH result AS ( + SELECT parse_json(process)['event_time']::STRING AS event_time + FROM get_lineage('lineage_history_statements.insert_dst', 'TABLE', 'UPSTREAM', 1) +), history AS ( + SELECT event_time + FROM system_history.lineage_unresolved + WHERE source_database = 'lineage_history_statements' + AND source_name = 'src' + AND target_database = 'lineage_history_statements' + AND target_name = 'insert_dst' + QUALIFY row_number() OVER (ORDER BY event_time DESC, query_id DESC) = 1 +) +SELECT result.event_time = concat( + replace(to_string(convert_timezone('UTC', history.event_time)), ' ', 'T'), + 'Z' + ) +FROM result CROSS JOIN history +---- +1 + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_history_statements.multi_a', 'TABLE', 'UPSTREAM', 1) +---- +1 default.lineage_history_statements.src default.lineage_history_statements.multi_a + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_history_statements.multi_b', 'TABLE', 'UPSTREAM', 1) +---- +1 default.lineage_history_statements.src default.lineage_history_statements.multi_b + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_history_statements.replace_dst', 'TABLE', 'UPSTREAM', 1) +---- +1 default.lineage_history_statements.src default.lineage_history_statements.replace_dst + +# Stream data lineage passes through to the backing table. +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_history_statements.stream_dst', 'TABLE', 'UPSTREAM', 1) +---- +1 default.lineage_history_statements.stream_src default.lineage_history_statements.stream_dst + +query ISS +SELECT distance, source_column_name, target_column_name +FROM get_lineage('lineage_history_statements.stream_dst.x', 'COLUMN', 'UPSTREAM', 1) +---- +1 a x + +# The Stream wrapper itself must not be persisted as a lineage endpoint. +query I +SELECT count(*) +FROM system_history.lineage_unresolved +WHERE source_database = 'lineage_history_statements' + AND source_name = 'src_stream' + AND target_database = 'lineage_history_statements' + AND target_name = 'stream_dst' +---- +0 + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_history_statements.copy_dst', 'TABLE', 'UPSTREAM', 2) +ORDER BY distance, source_object_name, target_object_name +---- +1 lineage_history_statements_stage default.lineage_history_statements.copy_dst +2 default.lineage_history_statements.src lineage_history_statements_stage + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_history_statements_stage', 'STAGE', 'DOWNSTREAM', 1) +ORDER BY target_object_name +---- +1 lineage_history_statements_stage default.lineage_history_statements.copy_dst +1 lineage_history_statements_stage default.lineage_history_statements.insert_select_stage_dst +1 lineage_history_statements_stage default.lineage_history_statements.insert_stage_dst + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_history_statements.insert_select_stage_dst', 'TABLE', 'UPSTREAM', 1) +---- +1 lineage_history_statements_stage default.lineage_history_statements.insert_select_stage_dst + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_history_statements.insert_stage_dst', 'TABLE', 'UPSTREAM', 1) +---- +1 lineage_history_statements_stage default.lineage_history_statements.insert_stage_dst + +# Stage file fields are not stable columns, so Stage-to-table edges stay object-only. +query I +SELECT count(*) +FROM get_lineage('lineage_history_statements.copy_dst.a', 'COLUMN', 'UPSTREAM', 1) +---- +0 + +query I +SELECT count(*) +FROM get_lineage('lineage_history_statements.insert_select_stage_dst.a', 'COLUMN', 'UPSTREAM', 1) +---- +0 + +query I +SELECT count(*) +FROM get_lineage('lineage_history_statements.insert_stage_dst.a', 'COLUMN', 'UPSTREAM', 1) +---- +0 + +statement ok +DROP DATABASE lineage_history_statements + +statement ok +DROP STAGE lineage_history_statements_stage diff --git a/tests/logging/history_table/sqllogic/check/05_object_lifecycle.test b/tests/logging/history_table/sqllogic/check/05_object_lifecycle.test new file mode 100644 index 0000000000000..6db1d9db96e93 --- /dev/null +++ b/tests/logging/history_table/sqllogic/check/05_object_lifecycle.test @@ -0,0 +1,93 @@ +# Object lifecycle corresponds to setup/05_object_lifecycle.test. + +# The dropped Fuse edge and the two DataMovement column-lifecycle edges remain in history. +query TI rowsort +SELECT target_name, count(*) +FROM system_history.lineage_unresolved +WHERE target_database = 'lineage_history_lifecycle' +GROUP BY target_name +---- +dropped_fuse 1 +renamed_column_dst 1 +replaced_column_dst 1 + +# A retained physical edge is inactive while its target object is dropped. +query I +SELECT count(*) +FROM get_lineage('lineage_history_lifecycle.dropped_fuse', 'TABLE', 'UPSTREAM', 1) +---- +0 + +# UNDROP restores the same table id, so the retained edge becomes active without a rewrite. +statement ok +UNDROP TABLE lineage_history_lifecycle.dropped_fuse + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_history_lifecycle.dropped_fuse', 'TABLE', 'UPSTREAM', 1) +---- +1 default.lineage_history_lifecycle.src default.lineage_history_lifecycle.dropped_fuse + +# Renaming preserves column ids, so the existing DataMovement edge resolves both current names. +statement ok +ALTER TABLE lineage_history_lifecycle.renamed_column_src RENAME COLUMN a TO renamed_a + +statement ok +ALTER TABLE lineage_history_lifecycle.renamed_column_dst RENAME COLUMN a TO renamed_a + +query ITTTT +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage( + 'lineage_history_lifecycle.renamed_column_dst.renamed_a', + 'COLUMN', + 'UPSTREAM', + 1 +) +---- +1 default.lineage_history_lifecycle.renamed_column_src renamed_a default.lineage_history_lifecycle.renamed_column_dst renamed_a + +# Dropping and adding the same name allocates a new column id. The object edge remains valid, but +# the old column mapping must not attach itself to the replacement column. +statement ok +ALTER TABLE lineage_history_lifecycle.replaced_column_src DROP COLUMN a + +statement ok +ALTER TABLE lineage_history_lifecycle.replaced_column_src ADD COLUMN a INT + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage( + 'lineage_history_lifecycle.replaced_column_dst', + 'TABLE', + 'UPSTREAM', + 1 +) +---- +1 default.lineage_history_lifecycle.replaced_column_src default.lineage_history_lifecycle.replaced_column_dst + +query I +SELECT count(*) +FROM get_lineage( + 'lineage_history_lifecycle.replaced_column_dst.a', + 'COLUMN', + 'UPSTREAM', + 1 +) +---- +0 + +# Leave cleanup tombstones for the next history transform cycle. +statement ok +DROP TABLE lineage_history_lifecycle.dropped_fuse + +statement ok +SET data_retention_time_in_days = 0 + +statement ok +VACUUM DROP TABLE FROM lineage_history_lifecycle + +statement ok +UNSET data_retention_time_in_days + +statement ok +DROP DATABASE lineage_history_lifecycle diff --git a/tests/logging/history_table/sqllogic/check/06_iceberg_catalog.test b/tests/logging/history_table/sqllogic/check/06_iceberg_catalog.test new file mode 100644 index 0000000000000..1f21b968e7fc9 --- /dev/null +++ b/tests/logging/history_table/sqllogic/check/06_iceberg_catalog.test @@ -0,0 +1,48 @@ +# Iceberg lineage corresponds to setup/06_iceberg_catalog.test. + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_history_iceberg.dst', 'TABLE', 'UPSTREAM', 1) +---- +1 lineage_history_iceberg_catalog.lineage_db.src default.lineage_history_iceberg.dst + +query ITTTT +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_history_iceberg.dst.a', 'COLUMN', 'UPSTREAM', 1) +---- +1 lineage_history_iceberg_catalog.lineage_db.src a default.lineage_history_iceberg.dst a + +# External catalog ids are not stable across catalog reloads. Both the object and its column map +# must therefore be name-addressed in the aggregated physical record. +query IIII +SELECT + count_if(source_catalog_type = 'ICEBERG'), + count_if(source_address_kind = 'NAME'), + count_if(source_id IS NULL), + count_if(source_to_target_columns['a'] IS NOT NULL) +FROM system_history.lineage_unresolved +WHERE source_catalog = 'lineage_history_iceberg_catalog' + AND source_database = 'lineage_db' + AND source_name = 'src' + AND target_database = 'lineage_history_iceberg' + AND target_name = 'dst' +---- +1 1 1 1 + +statement ok +DROP DATABASE lineage_history_iceberg + +statement ok +USE CATALOG lineage_history_iceberg_catalog + +statement ok +DROP TABLE lineage_db.src + +statement ok +DROP DATABASE lineage_db + +statement ok +USE CATALOG default + +statement ok +DROP CATALOG lineage_history_iceberg_catalog diff --git a/tests/logging/history_table/sqllogic/setup/01_object_patterns.test b/tests/logging/history_table/sqllogic/setup/01_object_patterns.test new file mode 100644 index 0000000000000..44922f769165f --- /dev/null +++ b/tests/logging/history_table/sqllogic/setup/01_object_patterns.test @@ -0,0 +1,21 @@ +# Object patterns: CTAS and DML retain distinct column mappings for the same object edge. + +statement ok +DROP DATABASE IF EXISTS lineage_history_objects + +statement ok +CREATE DATABASE lineage_history_objects + +statement ok +CREATE TABLE lineage_history_objects.src(a INT, b INT) + +statement ok +INSERT INTO lineage_history_objects.src VALUES (1, 2) + +statement ok +CREATE TABLE lineage_history_objects.mid AS +SELECT a AS x, b AS y FROM lineage_history_objects.src + +statement ok +INSERT INTO lineage_history_objects.mid(x, y) +SELECT b, a FROM lineage_history_objects.src diff --git a/tests/logging/history_table/sqllogic/setup/02_column_multihop.test b/tests/logging/history_table/sqllogic/setup/02_column_multihop.test new file mode 100644 index 0000000000000..f13b78dd3c489 --- /dev/null +++ b/tests/logging/history_table/sqllogic/setup/02_column_multihop.test @@ -0,0 +1,35 @@ +# Column multi-hop: retain multiple mappings and extend them through another CTAS edge. + +statement ok +DROP DATABASE IF EXISTS lineage_history_columns + +statement ok +DROP MASKING POLICY IF EXISTS lineage_history_columns_mask + +statement ok +CREATE DATABASE lineage_history_columns + +statement ok +CREATE TABLE lineage_history_columns.src(a INT, b INT) + +statement ok +INSERT INTO lineage_history_columns.src VALUES (1, 2) + +statement ok +CREATE TABLE lineage_history_columns.mid AS +SELECT a AS x, b AS y FROM lineage_history_columns.src + +statement ok +INSERT INTO lineage_history_columns.mid(x, y) +SELECT b, a FROM lineage_history_columns.src + +statement ok +CREATE TABLE lineage_history_columns.dst AS +SELECT x AS z FROM lineage_history_columns.mid + +# Target status is column-specific: only dst.z has a masking policy. +statement ok +CREATE MASKING POLICY lineage_history_columns_mask AS (val INT) RETURNS INT -> -1 + +statement ok +ALTER TABLE lineage_history_columns.dst MODIFY COLUMN z SET MASKING POLICY lineage_history_columns_mask diff --git a/tests/logging/history_table/sqllogic/setup/03_view_boundary.test b/tests/logging/history_table/sqllogic/setup/03_view_boundary.test new file mode 100644 index 0000000000000..a0535b683b0d2 --- /dev/null +++ b/tests/logging/history_table/sqllogic/setup/03_view_boundary.test @@ -0,0 +1,18 @@ +# View boundary: view columns are name-addressed while the view object keeps its stable id. + +statement ok +DROP DATABASE IF EXISTS lineage_history_views + +statement ok +CREATE DATABASE lineage_history_views + +statement ok +CREATE TABLE lineage_history_views.src(a INT, b INT) + +statement ok +CREATE VIEW lineage_history_views.src_view AS +SELECT a + b AS v FROM lineage_history_views.src + +statement ok +CREATE TABLE lineage_history_views.view_dst AS +SELECT v AS z FROM lineage_history_views.src_view diff --git a/tests/logging/history_table/sqllogic/setup/04_statement_coverage.test b/tests/logging/history_table/sqllogic/setup/04_statement_coverage.test new file mode 100644 index 0000000000000..c5d858ae9a1da --- /dev/null +++ b/tests/logging/history_table/sqllogic/setup/04_statement_coverage.test @@ -0,0 +1,107 @@ +# Supported statements: capture object lineage for INSERT, multi-table INSERT, REPLACE, and COPY. + +statement ok +DROP DATABASE IF EXISTS lineage_history_statements + +statement ok +DROP STAGE IF EXISTS lineage_history_statements_stage + +statement ok +CREATE DATABASE lineage_history_statements + +statement ok +CREATE TABLE lineage_history_statements.src(a INT, b INT, c INT) + +statement ok +INSERT INTO lineage_history_statements.src VALUES (1, 2, 3) + +# The self source must be removed without dropping the other source edge. +statement ok +CREATE TABLE lineage_history_statements.self_insert_dst(a INT, b INT) + +statement ok +INSERT INTO lineage_history_statements.self_insert_dst VALUES (10, 100) + +statement ok +INSERT INTO lineage_history_statements.self_insert_dst +SELECT d.a + s.a, s.b +FROM lineage_history_statements.self_insert_dst AS d +CROSS JOIN lineage_history_statements.src AS s + +statement ok +CREATE TABLE lineage_history_statements.insert_dst(x INT, y INT) + +statement ok +INSERT INTO lineage_history_statements.insert_dst +SELECT a, b FROM lineage_history_statements.src + +statement ok +CREATE TABLE lineage_history_statements.multi_a(x INT, y INT) + +statement ok +CREATE TABLE lineage_history_statements.multi_b(x INT, y INT) + +statement ok +INSERT ALL + INTO lineage_history_statements.multi_a(x, y) VALUES(a, b) + INTO lineage_history_statements.multi_b(x, y) VALUES(b, c) +SELECT a, b, c FROM lineage_history_statements.src + +statement ok +CREATE TABLE lineage_history_statements.replace_dst(x INT, y INT) + +statement ok +REPLACE INTO lineage_history_statements.replace_dst ON(x) +SELECT a, b FROM lineage_history_statements.src + +# A Stream is transparent for lineage: its data columns belong to the backing table. +statement ok +CREATE TABLE lineage_history_statements.stream_src(a INT, b INT) + +statement ok +CREATE STREAM lineage_history_statements.src_stream +ON TABLE lineage_history_statements.stream_src APPEND_ONLY = TRUE + +statement ok +INSERT INTO lineage_history_statements.stream_src VALUES (1, 2) + +statement ok +CREATE TABLE lineage_history_statements.stream_dst(x INT) + +statement ok +INSERT INTO lineage_history_statements.stream_dst +SELECT a + 1 FROM lineage_history_statements.src_stream + +statement ok +CREATE STAGE lineage_history_statements_stage FILE_FORMAT = (TYPE = PARQUET) + +statement ok +COPY INTO @lineage_history_statements_stage +FROM (SELECT a, b FROM lineage_history_statements.src) +FILE_FORMAT = (TYPE = PARQUET) + +statement ok +CREATE TABLE lineage_history_statements.copy_dst(a INT, b INT) + +statement ok +COPY INTO lineage_history_statements.copy_dst +FROM @lineage_history_statements_stage +FILE_FORMAT = (TYPE = PARQUET) +FORCE = TRUE + +# A Stage scan inside a regular SELECT must remain a Stage lineage endpoint. +statement ok +CREATE TABLE lineage_history_statements.insert_select_stage_dst(a INT, b INT) + +statement ok +INSERT INTO lineage_history_statements.insert_select_stage_dst +SELECT a, b FROM @lineage_history_statements_stage + +# INSERT FROM Stage uses the COPY plan path and purges input files, so keep it last. +statement ok +CREATE TABLE lineage_history_statements.insert_stage_dst(a INT, b INT) + +statement ok +INSERT INTO lineage_history_statements.insert_stage_dst +FROM @lineage_history_statements_stage +FILE_FORMAT = (TYPE = PARQUET) diff --git a/tests/logging/history_table/sqllogic/setup/05_object_lifecycle.test b/tests/logging/history_table/sqllogic/setup/05_object_lifecycle.test new file mode 100644 index 0000000000000..c191d8ab4ddc5 --- /dev/null +++ b/tests/logging/history_table/sqllogic/setup/05_object_lifecycle.test @@ -0,0 +1,78 @@ +# Object lifecycle: DROP invalidates edges, transient/View DROP removes them, and VACUUM cleans Fuse edges. + +statement ok +DROP DATABASE IF EXISTS lineage_history_lifecycle + +statement ok +CREATE DATABASE lineage_history_lifecycle + +statement ok +CREATE TABLE lineage_history_lifecycle.src(a INT) + +# DataMovement column lineage uses stable column ids. The check suite mutates these schemas after +# history ingestion to verify rename and replacement behavior without emitting another edge. +statement ok +CREATE TABLE lineage_history_lifecycle.renamed_column_src(a INT) + +statement ok +INSERT INTO lineage_history_lifecycle.renamed_column_src VALUES (1) + +statement ok +CREATE TABLE lineage_history_lifecycle.renamed_column_dst(a INT) + +statement ok +INSERT INTO lineage_history_lifecycle.renamed_column_dst +SELECT a FROM lineage_history_lifecycle.renamed_column_src + +statement ok +CREATE TABLE lineage_history_lifecycle.replaced_column_src(a INT, keep INT) + +statement ok +INSERT INTO lineage_history_lifecycle.replaced_column_src VALUES (1, 0) + +statement ok +CREATE TABLE lineage_history_lifecycle.replaced_column_dst(a INT) + +statement ok +INSERT INTO lineage_history_lifecycle.replaced_column_dst +SELECT a FROM lineage_history_lifecycle.replaced_column_src + +# A dropped Fuse edge is physically removed only after VACUUM. +statement ok +CREATE TABLE lineage_history_lifecycle.vacuumed_fuse AS +SELECT a FROM lineage_history_lifecycle.src + +statement ok +DROP TABLE lineage_history_lifecycle.vacuumed_fuse + +statement ok +SET data_retention_time_in_days = 0 + +statement ok +VACUUM DROP TABLE FROM lineage_history_lifecycle + +statement ok +UNSET data_retention_time_in_days + +# This Fuse table is dropped after VACUUM, so its physical edge remains available to UNDROP. +statement ok +CREATE TABLE lineage_history_lifecycle.dropped_fuse AS +SELECT a FROM lineage_history_lifecycle.src + +statement ok +DROP TABLE lineage_history_lifecycle.dropped_fuse + +# Transient tables and Views emit deletion tombstones immediately on DROP. +statement ok +CREATE TRANSIENT TABLE lineage_history_lifecycle.dropped_transient AS +SELECT a FROM lineage_history_lifecycle.src + +statement ok +DROP TABLE lineage_history_lifecycle.dropped_transient + +statement ok +CREATE VIEW lineage_history_lifecycle.dropped_view AS +SELECT a FROM lineage_history_lifecycle.src + +statement ok +DROP VIEW lineage_history_lifecycle.dropped_view diff --git a/tests/logging/history_table/sqllogic/setup/06_iceberg_catalog.test b/tests/logging/history_table/sqllogic/setup/06_iceberg_catalog.test new file mode 100644 index 0000000000000..b78c1ec42006e --- /dev/null +++ b/tests/logging/history_table/sqllogic/setup/06_iceberg_catalog.test @@ -0,0 +1,43 @@ +# Iceberg catalog tables have catalog-local runtime ids, so lineage must use object and column names. + +statement ok +DROP DATABASE IF EXISTS lineage_history_iceberg + +statement ok +DROP CATALOG IF EXISTS lineage_history_iceberg_catalog + +statement ok +CREATE CATALOG lineage_history_iceberg_catalog +TYPE = ICEBERG +CONNECTION = ( + TYPE = 'rest' + ADDRESS = 'http://127.0.0.1:8181' + WAREHOUSE = 's3://iceberg-tpch/' + "s3.region" = 'us-east-1' + "s3.endpoint" = 'http://127.0.0.1:9002' + "s3.access-key-id" = 'admin' + "s3.secret-access-key" = 'password' +) + +statement ok +USE CATALOG lineage_history_iceberg_catalog + +statement ok +DROP DATABASE IF EXISTS lineage_db + +statement ok +CREATE DATABASE lineage_db + +statement ok +CREATE TABLE lineage_db.src(a INT) + +statement ok +USE CATALOG default + +statement ok +CREATE DATABASE lineage_history_iceberg + +# The source is intentionally empty; successful planning and execution must still emit lineage. +statement ok +CREATE TABLE lineage_history_iceberg.dst AS +SELECT a FROM lineage_history_iceberg_catalog.lineage_db.src diff --git a/tests/logging/test-history-tables.sh b/tests/logging/test-history-tables.sh index c31837967f5e2..b67cc9a037f25 100644 --- a/tests/logging/test-history-tables.sh +++ b/tests/logging/test-history-tables.sh @@ -20,12 +20,46 @@ BUILD_PROFILE="${BUILD_PROFILE:-debug}" SCRIPT_PATH="$(cd "$(dirname "$0")" >/dev/null 2>&1 && pwd)" cd "$SCRIPT_PATH/../../" || exit +ICEBERG_COMPOSE_FILE="$PWD/tests/sqllogictests/scripts/docker-compose-iceberg-tpch.yml" +iceberg_services_started=false + +stop_iceberg_services() { + if [ "$iceberg_services_started" = true ]; then + echo "Stopping Iceberg REST test services" + docker compose -f "$ICEBERG_COMPOSE_FILE" down || true + fi +} + +start_iceberg_services() { + echo "Starting Iceberg REST test services" + iceberg_services_started=true + docker compose -f "$ICEBERG_COMPOSE_FILE" up -d rustfs mc rest + + for _ in {1..60}; do + if curl -fsS http://127.0.0.1:9002/health/ready >/dev/null \ + && curl -fsS http://127.0.0.1:8181/v1/config >/dev/null \ + && docker compose -f "$ICEBERG_COMPOSE_FILE" exec -T mc \ + /usr/bin/mc stat rustfs/iceberg-tpch >/dev/null 2>&1; then + echo "Iceberg REST test services are ready" + return 0 + fi + sleep 1 + done + + echo "Iceberg REST test services did not become ready" + docker compose -f "$ICEBERG_COMPOSE_FILE" logs rustfs mc rest || true + return 1 +} + +trap stop_iceberg_services EXIT + echo "Cleaning up previous runs" killall -9 databend-query || true killall -9 databend-meta || true killall -9 vector || true rm -rf ./.databend +start_iceberg_services echo "Starting Databend Query cluster with 2 nodes enable history tables" @@ -68,13 +102,13 @@ echo 'Start databend-query node-1' nohup env RUST_BACKTRACE=1 target/${BUILD_PROFILE}/databend-query -c scripts/ci/deploy/config/databend-query-node-1.toml --internal-enable-sandbox-tenant >./.databend/query-1.out 2>&1 & echo "Waiting on node-1..." -python3 scripts/ci/wait_tcp.py --timeout 30 --port 9091 +python3 scripts/ci/wait_tcp.py --timeout 50 --port 9091 echo 'Start databend-query node-2' env "RUST_BACKTRACE=1" nohup target/${BUILD_PROFILE}/databend-query -c scripts/ci/deploy/config/databend-query-node-2.toml --internal-enable-sandbox-tenant >./.databend/query-2.out 2>&1 & echo "Waiting on node-2..." -python3 scripts/ci/wait_tcp.py --timeout 30 --port 9092 +python3 scripts/ci/wait_tcp.py --timeout 50 --port 9092 echo "Started 2-node cluster with history tables enabled..."