diff --git a/scripts/ci/deploy/config/databend-query-node-1.toml b/scripts/ci/deploy/config/databend-query-node-1.toml index 6259f5a7288d3..e7b0845f542ac 100644 --- a/scripts/ci/deploy/config/databend-query-node-1.toml +++ b/scripts/ci/deploy/config/databend-query-node-1.toml @@ -80,6 +80,9 @@ definition = "CREATE FUNCTION ping(STRING) RETURNS STRING LANGUAGE python HANDLE aggregate_spilling_memory_ratio = 60 join_spilling_memory_ratio = 60 +[query.lineage] +capture_enabled = true + [log] [log.file] diff --git a/scripts/ci/deploy/config/databend-query-node-2.toml b/scripts/ci/deploy/config/databend-query-node-2.toml index 3336d4b2c776f..9e01c509af9e0 100644 --- a/scripts/ci/deploy/config/databend-query-node-2.toml +++ b/scripts/ci/deploy/config/databend-query-node-2.toml @@ -51,6 +51,9 @@ auth_type = "no_password" name = "ping" definition = "CREATE FUNCTION ping(STRING) RETURNS STRING LANGUAGE python HANDLER = 'ping' ADDRESS = 'http://0.0.0.0:8815'" +[query.lineage] +capture_enabled = true + [log] [log.file] diff --git a/scripts/ci/deploy/config/databend-query-node-3.toml b/scripts/ci/deploy/config/databend-query-node-3.toml index 8c4aa7db21c08..07a5e222f6a82 100644 --- a/scripts/ci/deploy/config/databend-query-node-3.toml +++ b/scripts/ci/deploy/config/databend-query-node-3.toml @@ -52,6 +52,9 @@ auth_type = "no_password" name = "ping" definition = "CREATE FUNCTION ping(STRING) RETURNS STRING LANGUAGE python HANDLER = 'ping' ADDRESS = 'http://0.0.0.0:8815'" +[query.lineage] +capture_enabled = true + [log] [log.file] diff --git a/scripts/ci/deploy/config/databend-query-node-hive.toml b/scripts/ci/deploy/config/databend-query-node-hive.toml index 50f1fe5ac5e38..00a72016278b7 100644 --- a/scripts/ci/deploy/config/databend-query-node-hive.toml +++ b/scripts/ci/deploy/config/databend-query-node-hive.toml @@ -44,6 +44,9 @@ auth_type = "no_password" name = "ping" definition = "CREATE FUNCTION ping(STRING) RETURNS STRING LANGUAGE python HANDLER = 'ping' ADDRESS = 'http://0.0.0.0:8815'" +[query.lineage] +capture_enabled = true + [log] [log.file] diff --git a/scripts/ci/deploy/config/databend-query-node-otlp-logs.toml b/scripts/ci/deploy/config/databend-query-node-otlp-logs.toml index c2623a6e3760d..9c710f848a17c 100644 --- a/scripts/ci/deploy/config/databend-query-node-otlp-logs.toml +++ b/scripts/ci/deploy/config/databend-query-node-otlp-logs.toml @@ -40,6 +40,9 @@ auth_type = "no_password" name = "ping" definition = "CREATE FUNCTION ping(STRING) RETURNS STRING LANGUAGE python HANDLER = 'ping' ADDRESS = 'http://0.0.0.0:8815'" +[query.lineage] +capture_enabled = true + [log] [log.file] diff --git a/scripts/ci/deploy/config/databend-query-node-system-managed.toml b/scripts/ci/deploy/config/databend-query-node-system-managed.toml index ef9d6f2405718..751ac4f676ae0 100644 --- a/scripts/ci/deploy/config/databend-query-node-system-managed.toml +++ b/scripts/ci/deploy/config/databend-query-node-system-managed.toml @@ -51,6 +51,9 @@ definition = "CREATE FUNCTION ping(STRING) RETURNS STRING LANGUAGE python HANDLE type = "system_managed" node_group +[query.lineage] +capture_enabled = true + [log] [log.file] diff --git a/src/common/license/src/license.rs b/src/common/license/src/license.rs index bdc4232fdec7d..6ad81400f3fcb 100644 --- a/src/common/license/src/license.rs +++ b/src/common/license/src/license.rs @@ -37,6 +37,8 @@ pub enum Feature { StorageEncryption, #[serde(alias = "stream", alias = "STREAM")] Stream, + #[serde(alias = "lineage", alias = "LINEAGE")] + Lineage, #[serde(alias = "table_ref", alias = "TABLE_REF")] TableRef, #[serde(alias = "attach_table", alias = "ATTACH_TABLE")] @@ -77,6 +79,7 @@ impl fmt::Display for Feature { Feature::ComputedColumn => write!(f, "computed_column"), Feature::StorageEncryption => write!(f, "storage_encryption"), Feature::Stream => write!(f, "stream"), + Feature::Lineage => write!(f, "lineage"), Feature::TableRef => write!(f, "table_ref"), Feature::AttacheTable => write!(f, "attach_table"), Feature::AmendTable => write!(f, "amend_table"), @@ -211,6 +214,10 @@ mod tests { Feature::Stream, serde_json::from_str::("\"Stream\"").unwrap() ); + assert_eq!( + Feature::Lineage, + serde_json::from_str::("\"LINEAGE\"").unwrap() + ); assert_eq!( Feature::TableRef, serde_json::from_str::("\"TableRef\"").unwrap() @@ -270,6 +277,7 @@ mod tests { Feature::ComputedColumn, Feature::StorageEncryption, Feature::Stream, + Feature::Lineage, Feature::TableRef, Feature::AttacheTable, Feature::AmendTable, @@ -281,7 +289,7 @@ mod tests { }; assert_eq!( - "LicenseInfo{ type: enterprise, org: databend, tenants: [databend_tenant,foo], features: [amend_table,attach_table,computed_column,data_mask,license_info,private_task,row_access_policy,storage_encryption,stream,system_history,table_ref,vacuum,workload_group] }", + "LicenseInfo{ type: enterprise, org: databend, tenants: [databend_tenant,foo], features: [amend_table,attach_table,computed_column,data_mask,license_info,lineage,private_task,row_access_policy,storage_encryption,stream,system_history,table_ref,vacuum,workload_group] }", license_info.to_string() ); } @@ -297,6 +305,7 @@ mod tests { Feature::ComputedColumn, Feature::StorageEncryption, Feature::Stream, + Feature::Lineage, Feature::TableRef, Feature::AttacheTable, Feature::AmendTable, diff --git a/src/meta/api/src/api_impl/auto_increment_api_test_suite.rs b/src/meta/api/src/api_impl/auto_increment_api_test_suite.rs index e99259e54aab9..0950941f72879 100644 --- a/src/meta/api/src/api_impl/auto_increment_api_test_suite.rs +++ b/src/meta/api/src/api_impl/auto_increment_api_test_suite.rs @@ -128,6 +128,7 @@ impl AutoIncrementApiTestSuite { table_name: tbl_name.to_string(), }, table_meta: drop_table_meta(created_on), + lineage_updates: vec![], as_dropped: true, table_properties: None, table_partition: None, diff --git a/src/meta/api/src/api_impl/garbage_collection_api.rs b/src/meta/api/src/api_impl/garbage_collection_api.rs index aa09647daeb0a..4be5601489b61 100644 --- a/src/meta/api/src/api_impl/garbage_collection_api.rs +++ b/src/meta/api/src/api_impl/garbage_collection_api.rs @@ -73,6 +73,7 @@ use log::info; use log::warn; use super::index_api::IndexApi; +use super::lineage_api::append_delete_lineage_for_table_id_txn_ops; use crate::kv_app_error::KVAppError; use crate::kv_pb_api::KVPbApi; use crate::kv_pb_crud_api::KVPbCrudApi; @@ -662,6 +663,13 @@ async fn gc_dropped_table_by_id( // 3) remove_index_for_dropped_table(kv_api, tenant, table_id_ident, &mut txn).await?; + append_delete_lineage_for_table_id_txn_ops( + kv_api, + tenant, + &mut txn, + table_id_ident.table_id, + ) + .await?; // Count removed keys (approximate for DeleteByPrefix operations) let mut num_meta_keys_removed = 0; diff --git a/src/meta/api/src/api_impl/lineage_api.rs b/src/meta/api/src/api_impl/lineage_api.rs new file mode 100644 index 0000000000000..3cf0ab5f47378 --- /dev/null +++ b/src/meta/api/src/api_impl/lineage_api.rs @@ -0,0 +1,608 @@ +// 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 databend_common_meta_app::schema::LineageColumn; +use databend_common_meta_app::schema::LineageDetail; +use databend_common_meta_app::schema::LineageDirection; +use databend_common_meta_app::schema::LineageIdent; +use databend_common_meta_app::schema::LineageIdentity; +use databend_common_meta_app::schema::LineageKey; +use databend_common_meta_app::schema::LineageObjectRef; +use databend_common_meta_app::schema::LineageObjectType; +use databend_common_meta_app::schema::LineageUpdate; +use databend_common_meta_app::schema::LineageUpdateMode; +use databend_common_meta_app::tenant::Tenant; +use databend_meta_client::kvapi::DirName; +use databend_meta_client::kvapi::ListOptions; +use databend_meta_client::kvapi::StructKey; +use databend_meta_client::types::ConditionResult::Eq; +use databend_meta_client::types::MetaError; +use databend_meta_client::types::TxnRequest; +use fastrace::func_name; + +use crate::kv_app_error::KVAppError; +use crate::kv_fetch_util::mget_pb_values; +use crate::kv_pb_api::KVPbApi; +use crate::txn_backoff::txn_backoff; +use crate::txn_condition_util::txn_cond_seq; +use crate::txn_core_util::send_txn; +use crate::txn_op_builder_util::txn_del; +use crate::txn_op_builder_util::txn_put_pb; + +const LINEAGE_MERGE_MAX_RETRIES: u32 = 3; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct MergeLineageReq { + pub updates: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ListLineageReq { + pub tenant: Tenant, + pub direction: LineageDirection, + pub object: LineageObjectRef, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ListLineageReply { + pub entries: Vec<(LineageKey, LineageDetail)>, +} + +pub(crate) fn append_replace_lineage_txn_ops<'a>( + txn: &mut TxnRequest, + updates: impl IntoIterator, +) { + for update in updates { + debug_assert_eq!(update.mode, LineageUpdateMode::Replace); + if update.mode != LineageUpdateMode::Replace { + continue; + } + + let idents = LineageEdgeIdents::from_update(update); + txn.if_then + .push(txn_put_pb(&idents.downstream, &update.detail)); + txn.if_then + .push(txn_put_pb(&idents.upstream, &update.detail)); + } +} + +pub(crate) async fn append_delete_lineage_for_table_id_txn_ops( + kv: &KV, + tenant: &Tenant, + txn: &mut TxnRequest, + table_id: u64, +) -> Result<(), MetaError> +where + KV: KVPbApi + ?Sized, +{ + let object = LineageObjectRef { + object_type: LineageObjectType::Table, + identity: LineageIdentity::Id { + id: table_id.to_string(), + }, + }; + + append_delete_lineage_for_object_txn_ops(kv, tenant, txn, LineageDirection::Upstream, &object) + .await?; + append_delete_lineage_for_object_txn_ops( + kv, + tenant, + txn, + LineageDirection::Downstream, + &object, + ) + .await?; + Ok(()) +} + +async fn append_delete_lineage_for_object_txn_ops( + kv: &KV, + tenant: &Tenant, + txn: &mut TxnRequest, + direction: LineageDirection, + object: &LineageObjectRef, +) -> Result<(), MetaError> +where + KV: KVPbApi + ?Sized, +{ + let dir = lineage_object_dir(tenant.clone(), direction, object.clone()); + let entries = kv.list_pb_vec(ListOptions::unlimited(&dir)).await?; + + for (ident, _) in entries { + let key = ident.name(); + let reverse_key = LineageKey { + direction: reverse_direction(&key.direction), + object: key.related_object.clone(), + related_object: key.object.clone(), + }; + let reverse_ident = LineageIdent::new_generic(tenant.clone(), reverse_key); + txn.if_then.push(txn_del(&ident)); + txn.if_then.push(txn_del(&reverse_ident)); + } + + Ok(()) +} + +fn lineage_object_dir( + tenant: Tenant, + direction: LineageDirection, + object: LineageObjectRef, +) -> DirName { + let dummy_related = LineageObjectRef { + object_type: LineageObjectType::Table, + identity: LineageIdentity::Id { + id: "0".to_string(), + }, + }; + let prefix = LineageIdent::new_generic(tenant.clone(), LineageKey { + direction, + object, + related_object: dummy_related, + }); + // A related object is encoded as type / identity kind / identity. + DirName::new_with_level(prefix, 3) +} + +#[derive(Clone)] +struct LineageEdgeIdents { + downstream: LineageIdent, + upstream: LineageIdent, +} + +impl LineageEdgeIdents { + fn from_update(update: &LineageUpdate) -> Self { + Self::new( + Tenant::new_literal(&update.tenant_name), + update.upstream.clone(), + update.downstream.clone(), + ) + } + + fn new(tenant: Tenant, upstream: LineageObjectRef, downstream: LineageObjectRef) -> Self { + Self { + downstream: LineageIdent::new_generic(tenant.clone(), LineageKey { + direction: LineageDirection::Downstream, + object: upstream.clone(), + related_object: downstream.clone(), + }), + upstream: LineageIdent::new_generic(tenant, LineageKey { + direction: LineageDirection::Upstream, + object: downstream, + related_object: upstream, + }), + } + } +} + +fn reverse_direction(direction: &LineageDirection) -> LineageDirection { + match direction { + LineageDirection::Upstream => LineageDirection::Downstream, + LineageDirection::Downstream => LineageDirection::Upstream, + } +} + +fn merge_lineage_detail( + existing: Option, + incoming: &LineageDetail, +) -> LineageDetail { + let Some(mut detail) = existing else { + return incoming.clone(); + }; + + detail.last_query_id = incoming.last_query_id.clone(); + detail.updated_on = incoming.updated_on; + for column in &incoming.column_lineage { + if !detail + .column_lineage + .iter() + .any(|existing| lineage_column_eq(existing, column)) + { + detail.column_lineage.push(column.clone()); + } + } + + detail +} + +fn lineage_column_eq(left: &LineageColumn, right: &LineageColumn) -> bool { + left.upstream == right.upstream && left.downstream == right.downstream +} + +#[async_trait::async_trait] +pub trait LineageApi: Send + Sync { + async fn merge_lineage(&self, req: MergeLineageReq) -> Result<(), KVAppError>; + + async fn list_lineage(&self, req: ListLineageReq) -> Result; +} + +#[async_trait::async_trait] +impl + ?Sized> LineageApi for KV { + async fn merge_lineage(&self, req: MergeLineageReq) -> Result<(), KVAppError> { + let updates = aggregate_merge_updates(&req.updates); + if updates.is_empty() { + return Ok(()); + } + + let mut trials = txn_backoff(Some(LINEAGE_MERGE_MAX_RETRIES), func_name!()); + loop { + trials.next().unwrap()?.await; + let txn = build_merge_lineage_txn(self, &updates).await?; + let (success, _) = send_txn(self, txn).await?; + if success { + return Ok(()); + } + } + } + + async fn list_lineage(&self, req: ListLineageReq) -> Result { + let dir = lineage_object_dir(req.tenant, req.direction, req.object); + let entries = self + .list_pb_vec(ListOptions::unlimited(&dir)) + .await? + .into_iter() + .map(|(ident, seqv)| (ident.name().clone(), seqv.data)) + .collect(); + + Ok(ListLineageReply { entries }) + } +} + +struct AggregatedLineageUpdate { + idents: LineageEdgeIdents, + detail: LineageDetail, +} + +fn aggregate_merge_updates(updates: &[LineageUpdate]) -> Vec { + let mut aggregated = BTreeMap::new(); + + for update in updates + .iter() + .filter(|update| update.mode == LineageUpdateMode::Merge) + { + let idents = LineageEdgeIdents::from_update(update); + let key = idents.downstream.to_string_key(); + aggregated + .entry(key) + .and_modify(|existing: &mut AggregatedLineageUpdate| { + existing.detail = + merge_lineage_detail(Some(existing.detail.clone()), &update.detail); + }) + .or_insert_with(|| AggregatedLineageUpdate { + idents, + detail: update.detail.clone(), + }); + } + + aggregated.into_values().collect() +} + +async fn build_merge_lineage_txn( + kv: &KV, + updates: &[AggregatedLineageUpdate], +) -> Result +where + KV: KVPbApi + ?Sized, +{ + let mut merge_keys = Vec::with_capacity(updates.len() * 2); + for update in updates { + merge_keys.push(update.idents.downstream.to_string_key()); + merge_keys.push(update.idents.upstream.to_string_key()); + } + + let merge_values: Vec<(u64, Option)> = mget_pb_values(kv, &merge_keys).await?; + let mut txn = TxnRequest::default(); + for (update, values) in updates.iter().zip(merge_values.chunks_exact(2)) { + let (downstream_seq, downstream_detail) = values[0].clone(); + let (upstream_seq, _) = values[1].clone(); + let merged_detail = merge_lineage_detail(downstream_detail, &update.detail); + + txn.condition + .push(txn_cond_seq(&update.idents.downstream, Eq, downstream_seq)); + txn.condition + .push(txn_cond_seq(&update.idents.upstream, Eq, upstream_seq)); + txn.if_then + .push(txn_put_pb(&update.idents.downstream, &merged_detail)); + txn.if_then + .push(txn_put_pb(&update.idents.upstream, &merged_detail)); + } + + Ok(txn) +} + +#[cfg(test)] +mod tests { + use chrono::TimeZone; + use chrono::Utc; + use databend_common_meta_app::schema::ColumnRef; + use databend_common_meta_app::schema::LineageColumn; + use databend_common_meta_app::schema::LineageIdentity; + use databend_common_meta_app::schema::LineageKind; + use databend_common_meta_app::schema::LineageObjectType; + use databend_common_meta_app::tenant::Tenant; + use databend_meta_client::kvapi::StructKey; + + use super::*; + use crate::testing::new_local_meta_store; + + #[test] + fn test_merge_lineage_detail_updates_last_observed_fields_and_preserves_kind() { + let existing = LineageDetail { + kind: LineageKind::Ctas, + last_query_id: Some("query1".to_string()), + updated_on: lineage_time(), + column_lineage: vec![LineageColumn { + upstream: ColumnRef::Id(1), + downstream: ColumnRef::Id(2), + }], + }; + let incoming = LineageDetail { + kind: LineageKind::DataMovement, + last_query_id: Some("query2".to_string()), + updated_on: Utc.with_ymd_and_hms(2026, 7, 24, 0, 0, 0).unwrap(), + column_lineage: vec![ + LineageColumn { + upstream: ColumnRef::Id(1), + downstream: ColumnRef::Id(2), + }, + LineageColumn { + upstream: ColumnRef::Id(3), + downstream: ColumnRef::Id(4), + }, + ], + }; + + let got = merge_lineage_detail(Some(existing), &incoming); + + assert_eq!(got.kind, LineageKind::Ctas); + assert_eq!(got.last_query_id.as_deref(), Some("query2")); + assert_eq!(got.updated_on, incoming.updated_on); + assert_eq!(got.column_lineage.len(), 2); + } + + #[tokio::test] + async fn test_merge_lineage_aggregates_duplicate_edges_and_updates_both_directions() + -> anyhow::Result<()> { + let meta = new_local_meta_store().await; + let tenant = Tenant::new_literal("tenant1"); + let existing = LineageDetail { + kind: LineageKind::Ctas, + last_query_id: Some("query1".to_string()), + updated_on: lineage_time(), + column_lineage: vec![LineageColumn { + upstream: ColumnRef::Id(1), + downstream: ColumnRef::Id(2), + }], + }; + replace_lineage( + &meta, + &lineage_update(&tenant, "10", "20", existing, LineageUpdateMode::Replace), + ) + .await?; + + let updated_on = Utc.with_ymd_and_hms(2026, 7, 24, 0, 0, 0).unwrap(); + meta.merge_lineage(MergeLineageReq { + updates: vec![ + lineage_update( + &tenant, + "10", + "20", + LineageDetail { + kind: LineageKind::DataMovement, + last_query_id: Some("query2".to_string()), + updated_on: lineage_time(), + column_lineage: vec![LineageColumn { + upstream: ColumnRef::Id(1), + downstream: ColumnRef::Id(2), + }], + }, + LineageUpdateMode::Merge, + ), + lineage_update( + &tenant, + "10", + "20", + LineageDetail { + kind: LineageKind::DataMovement, + last_query_id: Some("query3".to_string()), + updated_on, + column_lineage: vec![ + LineageColumn { + upstream: ColumnRef::Id(1), + downstream: ColumnRef::Id(2), + }, + LineageColumn { + upstream: ColumnRef::Id(3), + downstream: ColumnRef::Id(4), + }, + ], + }, + LineageUpdateMode::Merge, + ), + ], + }) + .await?; + + let downstream = + list_entries(&meta, tenant.clone(), LineageDirection::Downstream, "10").await?; + let upstream = list_entries(&meta, tenant, LineageDirection::Upstream, "20").await?; + assert_eq!(downstream.len(), 1); + assert_eq!(upstream.len(), 1); + assert_eq!(downstream[0].1, upstream[0].1); + + let detail = &downstream[0].1; + assert_eq!(detail.kind, LineageKind::Ctas); + assert_eq!(detail.last_query_id.as_deref(), Some("query3")); + assert_eq!(detail.updated_on, updated_on); + assert_eq!(detail.column_lineage, vec![ + LineageColumn { + upstream: ColumnRef::Id(1), + downstream: ColumnRef::Id(2), + }, + LineageColumn { + upstream: ColumnRef::Id(3), + downstream: ColumnRef::Id(4), + }, + ]); + Ok(()) + } + + #[tokio::test] + async fn test_lineage_merge_retry_is_bounded() { + let mut trials = txn_backoff(Some(LINEAGE_MERGE_MAX_RETRIES), "test"); + for _ in 0..LINEAGE_MERGE_MAX_RETRIES { + trials.next().unwrap().unwrap().await; + } + assert!(trials.next().unwrap().is_err()); + } + + #[tokio::test] + async fn test_delete_lineage_for_table_id_is_scoped_to_tenant_and_object() -> anyhow::Result<()> + { + let meta = new_local_meta_store().await; + let tenant_a = Tenant::new_literal("tenant_a"); + let tenant_b = Tenant::new_literal("tenant_b"); + let detail = lineage_detail(); + + for (tenant, upstream, downstream) in [ + (tenant_a.clone(), "10", "20"), + (tenant_a.clone(), "20", "30"), + (tenant_a.clone(), "10", "21"), + (tenant_a.clone(), "200", "30"), + (tenant_b.clone(), "10", "20"), + ] { + replace_lineage( + &meta, + &lineage_update( + &tenant, + upstream, + downstream, + detail.clone(), + LineageUpdateMode::Replace, + ), + ) + .await?; + } + + let dir = lineage_object_dir( + tenant_a.clone(), + LineageDirection::Upstream, + table_ref("20"), + ); + assert_eq!( + dir.to_string_key(), + "__fd_lineage/tenant_a/upstream/table/id/20" + ); + + let mut txn = TxnRequest::default(); + append_delete_lineage_for_table_id_txn_ops(&meta, &tenant_a, &mut txn, 20).await?; + let (success, _) = send_txn(&meta, txn).await?; + assert!(success); + + assert!( + list_entries(&meta, tenant_a.clone(), LineageDirection::Upstream, "20") + .await? + .is_empty() + ); + assert!( + list_entries(&meta, tenant_a.clone(), LineageDirection::Downstream, "20") + .await? + .is_empty() + ); + assert_eq!( + list_entries(&meta, tenant_a.clone(), LineageDirection::Downstream, "10") + .await? + .len(), + 1 + ); + assert_eq!( + list_entries(&meta, tenant_a, LineageDirection::Downstream, "200") + .await? + .len(), + 1 + ); + assert_eq!( + list_entries(&meta, tenant_b, LineageDirection::Upstream, "20") + .await? + .len(), + 1 + ); + Ok(()) + } + + async fn list_entries( + kv: &KV, + tenant: Tenant, + direction: LineageDirection, + object_id: &str, + ) -> Result, MetaError> + where + KV: LineageApi + ?Sized, + { + Ok(kv + .list_lineage(ListLineageReq { + tenant, + direction, + object: table_ref(object_id), + }) + .await? + .entries) + } + + fn table_ref(id: &str) -> LineageObjectRef { + LineageObjectRef { + object_type: LineageObjectType::Table, + identity: LineageIdentity::Id { id: id.to_string() }, + } + } + + fn lineage_update( + tenant: &Tenant, + upstream: &str, + downstream: &str, + detail: LineageDetail, + mode: LineageUpdateMode, + ) -> LineageUpdate { + LineageUpdate { + tenant_name: tenant.tenant_name().to_string(), + upstream: table_ref(upstream), + downstream: table_ref(downstream), + detail, + mode, + } + } + + async fn replace_lineage(kv: &KV, update: &LineageUpdate) -> Result<(), MetaError> + where KV: KVPbApi + ?Sized { + let mut txn = TxnRequest::default(); + append_replace_lineage_txn_ops(&mut txn, [update]); + let (success, _) = send_txn(kv, txn).await?; + assert!(success); + Ok(()) + } + + fn lineage_detail() -> LineageDetail { + LineageDetail { + kind: LineageKind::DataMovement, + last_query_id: Some("query".to_string()), + updated_on: lineage_time(), + column_lineage: vec![], + } + } + + fn lineage_time() -> chrono::DateTime { + Utc.with_ymd_and_hms(2026, 7, 23, 0, 0, 0).unwrap() + } +} diff --git a/src/meta/api/src/api_impl/mod.rs b/src/meta/api/src/api_impl/mod.rs index a8908d6ad652e..519d6a192a32b 100644 --- a/src/meta/api/src/api_impl/mod.rs +++ b/src/meta/api/src/api_impl/mod.rs @@ -20,6 +20,7 @@ pub mod database_util; pub mod dictionary_api; pub mod garbage_collection_api; pub mod index_api; +pub mod lineage_api; pub mod lock_api2; pub mod ref_api; pub mod security_api; diff --git a/src/meta/api/src/api_impl/schema_api.rs b/src/meta/api/src/api_impl/schema_api.rs index f2234e5c2e939..cfa5d48064307 100644 --- a/src/meta/api/src/api_impl/schema_api.rs +++ b/src/meta/api/src/api_impl/schema_api.rs @@ -86,6 +86,7 @@ use super::database_util::get_db_or_err; use super::dictionary_api::DictionaryApi; use super::garbage_collection_api::GarbageCollectionApi; use super::index_api::IndexApi; +use super::lineage_api::append_delete_lineage_for_table_id_txn_ops; use super::lock_api2::LockApi2; use super::security_api::SecurityApi; use super::table_api::TableApi; @@ -238,6 +239,9 @@ pub async fn construct_drop_table_txn_operations( } tb_meta.drop_on = Some(Utc::now()); + if tb_meta.engine == "VIEW" || tb_meta.options.contains_key("TRANSIENT") { + append_delete_lineage_for_table_id_txn_ops(kv_api, tenant, txn, table_id).await?; + } // Delete table-policy references when dropping table // diff --git a/src/meta/api/src/api_impl/table_api.rs b/src/meta/api/src/api_impl/table_api.rs index 02b43f42f0db0..fa2ee441fa935 100644 --- a/src/meta/api/src/api_impl/table_api.rs +++ b/src/meta/api/src/api_impl/table_api.rs @@ -60,6 +60,10 @@ use databend_common_meta_app::schema::GetTableCopiedFileReply; use databend_common_meta_app::schema::GetTableCopiedFileReq; use databend_common_meta_app::schema::GetTableReq; use databend_common_meta_app::schema::LeastVisibleTime; +use databend_common_meta_app::schema::LineageIdentity; +use databend_common_meta_app::schema::LineageObjectType; +use databend_common_meta_app::schema::LineageUpdate; +use databend_common_meta_app::schema::LineageUpdateMode; use databend_common_meta_app::schema::ListDatabaseReq; use databend_common_meta_app::schema::ListDroppedTableReq; use databend_common_meta_app::schema::ListDroppedTableResp; @@ -118,6 +122,7 @@ use super::database_api::DatabaseApi; use super::database_util::get_db_or_err; use super::garbage_collection_api::ORPHAN_POSTFIX; use super::garbage_collection_api::get_history_tables_for_gc; +use super::lineage_api::append_replace_lineage_txn_ops; use super::schema_api::build_upsert_table_deduplicated_label; use super::schema_api::construct_drop_table_txn_operations; use super::schema_api::get_db_by_id_or_err; @@ -201,6 +206,29 @@ fn validate_create_table_request(req: &CreateTableReq) -> Result<(), KVAppError> Ok(()) } +fn bind_lineage_updates_to_table_id( + updates: &[LineageUpdate], + table_id: u64, +) -> Vec { + // CREATE TABLE allocates the table id inside meta-service. Lineage writers + // can only build downstream-by-name updates before this point, so bind the + // downstream object to the allocated table id before appending lineage kv ops. + updates + .iter() + .cloned() + .map(|mut update| { + if update.downstream.object_type == LineageObjectType::Table + && matches!(update.downstream.identity, LineageIdentity::Name { .. }) + { + update.downstream.identity = LineageIdentity::Id { + id: table_id.to_string(), + }; + } + update + }) + .collect() +} + /// TableApi defines APIs for table lifecycle and metadata management. /// /// This trait handles: @@ -483,6 +511,19 @@ where .extend(vec![txn_put_pb(&storage_ident, &storage_value)]); } + if !req.lineage_updates.is_empty() { + // Structural lineage (CREATE VIEW / CTAS) should become + // visible atomically with the created table metadata. + let lineage_updates = + bind_lineage_updates_to_table_id(&req.lineage_updates, table_id); + append_replace_lineage_txn_ops( + &mut txn, + lineage_updates + .iter() + .filter(|update| update.mode == LineageUpdateMode::Replace), + ); + } + let (succ, responses) = send_txn(self, txn).await?; debug!( @@ -1100,7 +1141,7 @@ where } tb_meta.drop_on = None; - let txn_req = TxnRequest::new( + let mut txn_req = TxnRequest::new( vec![ // db has not to change, i.e., no new table is created. // Renaming db is OK and does not affect the seq of db_meta. @@ -1124,6 +1165,19 @@ where ], ); + if !req.lineage_updates.is_empty() { + // CTAS creates the table as dropped first and commits it + // later; commit lineage in the same txn that makes it visible. + let lineage_updates = + bind_lineage_updates_to_table_id(&req.lineage_updates, table_id); + append_replace_lineage_txn_ops( + &mut txn_req, + lineage_updates + .iter() + .filter(|update| update.mode == LineageUpdateMode::Replace), + ); + } + let txn_response = txn_sender.send_txn(self, txn_req).await?; let succ = match txn_response { IdempotentKVTxnResponse::Success(_) => true, diff --git a/src/meta/api/src/lib.rs b/src/meta/api/src/lib.rs index bc6f60a27da01..9a8dd76cb0aaf 100644 --- a/src/meta/api/src/lib.rs +++ b/src/meta/api/src/lib.rs @@ -41,6 +41,7 @@ pub use api_impl::dictionary_api; pub(crate) use api_impl::errors; pub use api_impl::garbage_collection_api; pub use api_impl::index_api; +pub use api_impl::lineage_api; pub use api_impl::lock_api2; pub use api_impl::ref_api; pub(crate) use api_impl::row_access_policy_api; @@ -82,6 +83,10 @@ pub use kv_app_error::from_nested; pub use kv_fetch_util::deserialize_struct_get_response; pub use kv_fetch_util::fetch_id; pub use kv_fetch_util::mget_pb_values; +pub use lineage_api::LineageApi; +pub use lineage_api::ListLineageReply; +pub use lineage_api::ListLineageReq; +pub use lineage_api::MergeLineageReq; pub use lock_api2::LockApi2; pub use ref_api::RefApi; pub use row_access_policy_api::RowAccessPolicyApi; diff --git a/src/meta/app/src/schema/lineage.rs b/src/meta/app/src/schema/lineage.rs new file mode 100644 index 0000000000000..b0d41b8d36ee6 --- /dev/null +++ b/src/meta/app/src/schema/lineage.rs @@ -0,0 +1,173 @@ +// 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 chrono::DateTime; +use chrono::Utc; +use databend_meta_client::kvapi; + +use crate::tenant_key::ident::TIdent; + +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, kvapi::KeyCodec)] +pub enum LineageDirection { + Upstream, + Downstream, +} + +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, kvapi::KeyCodec)] +pub enum LineageObjectType { + /// Table-like objects that share the table namespace, including tables, views, + /// materialized views, and name-addressed foreign tables (like iceberg table.). + Table, + Stage, +} + +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, kvapi::KeyCodec)] +pub enum LineageIdentity { + Id { id: String }, + Name { name: String }, +} + +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, kvapi::KeyCodec)] +pub struct LineageObjectRef { + pub object_type: LineageObjectType, + pub identity: LineageIdentity, +} + +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, kvapi::KeyCodec)] +pub struct LineageKey { + pub direction: LineageDirection, + pub object: LineageObjectRef, + pub related_object: LineageObjectRef, +} + +pub type LineageIdent = TIdent; + +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub enum LineageKind { + View, + MaterializedView, + Ctas, + DataMovement, + Unknown(i32), +} + +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub struct LineageDetail { + pub kind: LineageKind, + pub last_query_id: Option, + pub updated_on: DateTime, + pub column_lineage: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub struct LineageUpdate { + pub tenant_name: String, + pub upstream: LineageObjectRef, + pub downstream: LineageObjectRef, + pub detail: LineageDetail, + pub mode: LineageUpdateMode, +} + +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub enum LineageUpdateMode { + Replace, + Merge, +} + +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub struct LineageColumn { + pub upstream: ColumnRef, + pub downstream: ColumnRef, +} + +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub enum ColumnRef { + Id(u64), + Name(String), +} + +pub use kvapi_impl::LineageRsc; + +mod kvapi_impl { + use crate::schema::LineageDetail; + use crate::tenant_key::resource::TenantResource; + + pub struct LineageRsc; + impl TenantResource for LineageRsc { + const PREFIX: &'static str = "__fd_lineage"; + const HAS_TENANT: bool = true; + type ValueType = LineageDetail; + } +} + +#[cfg(test)] +mod tests { + use databend_meta_client::kvapi::testing::assert_round_trip; + + use super::LineageDirection; + use super::LineageIdent; + use super::LineageIdentity; + use super::LineageKey; + use super::LineageObjectRef; + use super::LineageObjectType; + use crate::tenant::Tenant; + + #[test] + fn test_lineage_downstream_name_to_id_key() { + let ident = LineageIdent::new_generic(Tenant::new_literal("tenant_a"), LineageKey { + direction: LineageDirection::Downstream, + object: LineageObjectRef { + object_type: LineageObjectType::Stage, + identity: LineageIdentity::Name { + name: "catalog/db/stage/path".to_string(), + }, + }, + related_object: LineageObjectRef { + object_type: LineageObjectType::Table, + identity: LineageIdentity::Id { + id: "42".to_string(), + }, + }, + }); + + assert_round_trip( + ident, + "__fd_lineage/tenant_a/downstream/stage/name/catalog%2fdb%2fstage%2fpath/table/id/42", + ); + } + + #[test] + fn test_lineage_downstream_id_to_id_key() { + let ident = LineageIdent::new_generic(Tenant::new_literal("tenant_a"), LineageKey { + direction: LineageDirection::Downstream, + object: LineageObjectRef { + object_type: LineageObjectType::Table, + identity: LineageIdentity::Id { + id: "11".to_string(), + }, + }, + related_object: LineageObjectRef { + object_type: LineageObjectType::Table, + identity: LineageIdentity::Id { + id: "42".to_string(), + }, + }, + }); + + assert_round_trip( + ident, + "__fd_lineage/tenant_a/downstream/table/id/11/table/id/42", + ); + } +} diff --git a/src/meta/app/src/schema/mod.rs b/src/meta/app/src/schema/mod.rs index 0213e827e9358..470da8845fe49 100644 --- a/src/meta/app/src/schema/mod.rs +++ b/src/meta/app/src/schema/mod.rs @@ -36,6 +36,7 @@ pub mod index_id_to_name_ident; pub mod index_name_ident; mod least_visible_time; pub mod least_visible_time_ident; +mod lineage; mod lock; pub mod marked_deleted_index_id; pub mod marked_deleted_index_ident; @@ -88,6 +89,18 @@ pub use index::*; pub use index_name_ident::IndexNameIdent; pub use index_name_ident::IndexNameIdentRaw; pub use least_visible_time::LeastVisibleTime; +pub use lineage::ColumnRef; +pub use lineage::LineageColumn; +pub use lineage::LineageDetail; +pub use lineage::LineageDirection; +pub use lineage::LineageIdent; +pub use lineage::LineageIdentity; +pub use lineage::LineageKey; +pub use lineage::LineageKind; +pub use lineage::LineageObjectRef; +pub use lineage::LineageObjectType; +pub use lineage::LineageUpdate; +pub use lineage::LineageUpdateMode; pub use lock::CreateLockRevReply; pub use lock::CreateLockRevReq; pub use lock::DeleteLockRevReq; diff --git a/src/meta/app/src/schema/table/mod.rs b/src/meta/app/src/schema/table/mod.rs index 67574b2044943..2b86bf401fec1 100644 --- a/src/meta/app/src/schema/table/mod.rs +++ b/src/meta/app/src/schema/table/mod.rs @@ -39,6 +39,7 @@ use maplit::hashmap; use super::CatalogInfo; use super::CreateOption; use super::DatabaseId; +use super::LineageUpdate; use super::MarkedDeletedIndexMeta; use crate::schema::constraint::Constraint; use crate::schema::database_name_ident::DatabaseNameIdent; @@ -518,6 +519,7 @@ pub struct CreateTableReq { pub catalog_name: Option, pub name_ident: TableNameIdent, pub table_meta: TableMeta, + pub lineage_updates: Vec, /// Set it to true if a dropped table needs to be created, /// @@ -636,6 +638,7 @@ pub struct CommitTableMetaReq { pub table_id: MetaId, pub prev_table_id: Option, pub orphan_table_name: Option, + pub lineage_updates: Vec, } impl CommitTableMetaReq { diff --git a/src/meta/binaries/metabench/main.rs b/src/meta/binaries/metabench/main.rs index 6a4390b59f5c2..c73f97646d92c 100644 --- a/src/meta/binaries/metabench/main.rs +++ b/src/meta/binaries/metabench/main.rs @@ -458,6 +458,7 @@ async fn benchmark_table(client: &MetaStore, prefix: u64, client_num: u64, i: u6 catalog_name: None, name_ident: tb_name_ident(), table_meta: Default::default(), + lineage_updates: vec![], as_dropped: false, table_properties: None, table_partition: None, @@ -512,6 +513,7 @@ async fn benchmark_table(client: &MetaStore, prefix: u64, client_num: u64, i: u6 catalog_name: None, name_ident: tb_name_ident(), table_meta: Default::default(), + lineage_updates: vec![], as_dropped: false, table_properties: None, table_partition: None, @@ -589,6 +591,7 @@ async fn benchmark_create_tables( table_name: format!("t-{}", i), }, table_meta, + lineage_updates: vec![], as_dropped: false, table_properties: None, table_partition: None, diff --git a/src/meta/process/src/filter_tenant.rs b/src/meta/process/src/filter_tenant.rs index 0cb8fcb400516..33009d7cb085f 100644 --- a/src/meta/process/src/filter_tenant.rs +++ b/src/meta/process/src/filter_tenant.rs @@ -1176,6 +1176,7 @@ const TENANT_SCOPED_PREFIXES: &[&str] = &[ "__fd_dictionaries", "__fd_file_formats", "__fd_index", + "__fd_lineage", "__fd_mask_policy_apply_table_id", "__fd_network_policies", "__fd_object_owners", diff --git a/src/meta/process/src/pb_value_decoder.rs b/src/meta/process/src/pb_value_decoder.rs index 256351c8543a9..adaccdfed450b 100644 --- a/src/meta/process/src/pb_value_decoder.rs +++ b/src/meta/process/src/pb_value_decoder.rs @@ -44,6 +44,7 @@ use databend_common_meta_app::schema::DictionaryMeta; use databend_common_meta_app::schema::EmptyProto; use databend_common_meta_app::schema::IndexMeta; use databend_common_meta_app::schema::LeastVisibleTime; +use databend_common_meta_app::schema::LineageDetail; use databend_common_meta_app::schema::LockMeta; use databend_common_meta_app::schema::MarkedDeletedIndexMeta; use databend_common_meta_app::schema::ObjectTagIdRefValue; @@ -110,6 +111,7 @@ pub fn decode_pb_value(key: &str, bytes: &[u8]) -> String { "__fd_table_lvt/" => LeastVisibleTime, "__fd_table_lock/" => LockMeta, "__fd_vacuum_watermark_ts/" => VacuumWatermark, + "__fd_lineage/" => LineageDetail, // schema - sequence "__fd_sequence/" => SequenceMeta, diff --git a/src/meta/proto-conv/src/impls/lineage.rs b/src/meta/proto-conv/src/impls/lineage.rs new file mode 100644 index 0000000000000..718c07ae310ca --- /dev/null +++ b/src/meta/proto-conv/src/impls/lineage.rs @@ -0,0 +1,145 @@ +// 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. + +//! Conversion between protobuf lineage structs and meta structs. + +use chrono::DateTime; +use chrono::Utc; +use databend_common_meta_app::schema as mt; +use databend_common_protos::pb; + +use crate::FromToProto; +use crate::Incompatible; +use crate::MIN_READER_VER; +use crate::VER; +use crate::missing; +use crate::reader_check_msg; + +impl FromToProto for mt::LineageDetail { + type PB = pb::LineageDetail; + + fn get_pb_ver(p: &Self::PB) -> u64 { + p.ver + } + + fn from_pb(p: Self::PB) -> Result + where Self: Sized { + reader_check_msg(p.ver, p.min_reader_ver)?; + + Ok(Self { + kind: lineage_kind_from_pb(p.kind), + last_query_id: p.last_query_id, + updated_on: DateTime::::from_pb(p.updated_on)?, + column_lineage: p + .column_lineage + .into_iter() + .map(mt::LineageColumn::from_pb) + .collect::, _>>()?, + }) + } + + fn to_pb(&self) -> Self::PB { + pb::LineageDetail { + ver: VER, + min_reader_ver: MIN_READER_VER, + last_query_id: self.last_query_id.clone(), + kind: lineage_kind_to_pb(&self.kind), + updated_on: self.updated_on.to_pb(), + column_lineage: self.column_lineage.iter().map(|c| c.to_pb()).collect(), + } + } +} + +fn lineage_kind_from_pb(kind: i32) -> mt::LineageKind { + match kind { + value if value == pb::lineage_detail::LineageKind::View as i32 => mt::LineageKind::View, + value if value == pb::lineage_detail::LineageKind::MaterializedView as i32 => { + mt::LineageKind::MaterializedView + } + value if value == pb::lineage_detail::LineageKind::Ctas as i32 => mt::LineageKind::Ctas, + value if value == pb::lineage_detail::LineageKind::DataMovement as i32 => { + mt::LineageKind::DataMovement + } + other => mt::LineageKind::Unknown(other), + } +} + +fn lineage_kind_to_pb(kind: &mt::LineageKind) -> i32 { + match kind { + mt::LineageKind::View => pb::lineage_detail::LineageKind::View as i32, + mt::LineageKind::MaterializedView => { + pb::lineage_detail::LineageKind::MaterializedView as i32 + } + mt::LineageKind::Ctas => pb::lineage_detail::LineageKind::Ctas as i32, + mt::LineageKind::DataMovement => pb::lineage_detail::LineageKind::DataMovement as i32, + mt::LineageKind::Unknown(value) => *value, + } +} + +impl FromToProto for mt::LineageColumn { + type PB = pb::LineageColumn; + + fn get_pb_ver(_p: &Self::PB) -> u64 { + VER + } + + fn from_pb(p: Self::PB) -> Result + where Self: Sized { + Ok(Self { + upstream: p + .upstream + .map(mt::ColumnRef::from_pb) + .transpose()? + .ok_or_else(missing("LineageColumn.upstream"))?, + downstream: p + .downstream + .map(mt::ColumnRef::from_pb) + .transpose()? + .ok_or_else(missing("LineageColumn.downstream"))?, + }) + } + + fn to_pb(&self) -> Self::PB { + pb::LineageColumn { + upstream: Some(self.upstream.to_pb()), + downstream: Some(self.downstream.to_pb()), + } + } +} + +impl FromToProto for mt::ColumnRef { + type PB = pb::ColumnRef; + + fn get_pb_ver(_p: &Self::PB) -> u64 { + VER + } + + fn from_pb(p: Self::PB) -> Result + where Self: Sized { + match p.identity.ok_or_else(missing("ColumnRef.identity"))? { + pb::column_ref::Identity::Id(id) => Ok(Self::Id(id)), + pb::column_ref::Identity::Name(name) => Ok(Self::Name(name)), + } + } + + fn to_pb(&self) -> Self::PB { + let identity = match self { + mt::ColumnRef::Id(id) => pb::column_ref::Identity::Id(*id), + mt::ColumnRef::Name(name) => pb::column_ref::Identity::Name(name.clone()), + }; + pb::ColumnRef { + identity: Some(identity), + } + } +} diff --git a/src/meta/proto-conv/src/impls/mod.rs b/src/meta/proto-conv/src/impls/mod.rs index 2c67beb963554..f36b5215e069b 100644 --- a/src/meta/proto-conv/src/impls/mod.rs +++ b/src/meta/proto-conv/src/impls/mod.rs @@ -24,6 +24,7 @@ mod file_format; mod id; mod index; mod least_visible_time; +mod lineage; mod lock; mod mask_policy_table_id; mod owner; diff --git a/src/meta/proto-conv/src/util.rs b/src/meta/proto-conv/src/util.rs index 9d9e03af6f879..acae03f310a62 100644 --- a/src/meta/proto-conv/src/util.rs +++ b/src/meta/proto-conv/src/util.rs @@ -210,6 +210,7 @@ const META_CHANGE_LOG: &[(u64, &str)] = &[ (178, "2026-06-22: Add: config.proto StorageConfig variants for StorageParams::{Azblob,Ftp,Http,Ipfs,Memory}"), (179, "2026-07-06: Add: task.proto/TaskMessage.DeleteTask.task_id"), (180, "2026-07-10: Add: catalog.proto/PaimonCatalogOption"), + (181, "2026-07-14: Add: lineage.proto/LineageDetail"), // Dear developer: // If you're gonna add a new metadata version, you'll have to add a test for it. // You could just copy an existing test file(e.g., `../tests/it/v024_table_meta.rs`) diff --git a/src/meta/proto-conv/tests/it/main.rs b/src/meta/proto-conv/tests/it/main.rs index 096dac6f6c7f3..f51195586c33f 100644 --- a/src/meta/proto-conv/tests/it/main.rs +++ b/src/meta/proto-conv/tests/it/main.rs @@ -172,3 +172,4 @@ mod v177_arrow_file_format_params; mod v178_storage_config; mod v179_task_delete_task_id; mod v180_paimon_catalog_option; +mod v181_lineage_detail; diff --git a/src/meta/proto-conv/tests/it/v181_lineage_detail.rs b/src/meta/proto-conv/tests/it/v181_lineage_detail.rs new file mode 100644 index 0000000000000..572141691b351 --- /dev/null +++ b/src/meta/proto-conv/tests/it/v181_lineage_detail.rs @@ -0,0 +1,44 @@ +// Copyright 2026 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 chrono::TimeZone; +use chrono::Utc; +use databend_common_meta_app::schema as mt; +use fastrace::func_name; + +use crate::common; + +#[test] +fn test_decode_v181_lineage_detail() -> anyhow::Result<()> { + let lineage_detail_v181 = vec![ + 8, 4, 18, 2, 113, 49, 26, 23, 50, 48, 50, 54, 45, 48, 55, 45, 50, 51, 32, 48, 48, 58, 48, + 48, 58, 48, 48, 32, 85, 84, 67, 34, 8, 10, 2, 8, 1, 18, 2, 8, 2, 160, 6, 181, 1, 168, 6, + 24, + ]; + + let want = mt::LineageDetail { + kind: mt::LineageKind::DataMovement, + last_query_id: Some("q1".to_string()), + updated_on: Utc.with_ymd_and_hms(2026, 7, 23, 0, 0, 0).unwrap(), + column_lineage: vec![mt::LineageColumn { + upstream: mt::ColumnRef::Id(1), + downstream: mt::ColumnRef::Id(2), + }], + }; + + common::test_pb_from_to(func_name!(), want.clone())?; + common::test_load_old(func_name!(), lineage_detail_v181.as_slice(), 181, want)?; + + Ok(()) +} diff --git a/src/meta/protos/proto/lineage.proto b/src/meta/protos/proto/lineage.proto new file mode 100644 index 0000000000000..a798979e6308c --- /dev/null +++ b/src/meta/protos/proto/lineage.proto @@ -0,0 +1,47 @@ +// 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. + +syntax = "proto3"; + +package databend_proto; + +message LineageDetail { + enum LineageKind { + UNSPECIFIED = 0; + VIEW = 1; + MATERIALIZED_VIEW = 2; + CTAS = 3; + DATA_MOVEMENT = 4; + } + + uint64 ver = 100; + uint64 min_reader_ver = 101; + + LineageKind kind = 1; + optional string last_query_id = 2; + string updated_on = 3; + repeated LineageColumn column_lineage = 4; +} + +message LineageColumn { + ColumnRef upstream = 1; + ColumnRef downstream = 2; +} + +message ColumnRef { + oneof identity { + uint64 id = 1; + string name = 2; + } +} diff --git a/src/meta/schema-api-test-suite/src/db_table_harness.rs b/src/meta/schema-api-test-suite/src/db_table_harness.rs index 87b59a23c6a12..d54e37ccb8a77 100644 --- a/src/meta/schema-api-test-suite/src/db_table_harness.rs +++ b/src/meta/schema-api-test-suite/src/db_table_harness.rs @@ -190,6 +190,7 @@ where MT: kvapi::KVApi + TableApi table_name: self.tbl_name(), }, table_meta: table_meta.clone(), + lineage_updates: vec![], as_dropped: false, table_properties: None, table_partition: None, @@ -216,6 +217,7 @@ where MT: kvapi::KVApi + TableApi table_name: self.tbl_name(), }, table_meta: table_meta.clone(), + lineage_updates: vec![], as_dropped: false, table_properties: None, table_partition: None, diff --git a/src/meta/schema-api-test-suite/src/schema_api_test_suite.rs b/src/meta/schema-api-test-suite/src/schema_api_test_suite.rs index b576f67e42c9a..e9805172de659 100644 --- a/src/meta/schema-api-test-suite/src/schema_api_test_suite.rs +++ b/src/meta/schema-api-test-suite/src/schema_api_test_suite.rs @@ -1693,6 +1693,7 @@ impl SchemaApiTestSuite { }, table_meta: table_meta(created_on), + lineage_updates: vec![], as_dropped: false, table_properties: None, table_partition: None, @@ -1810,6 +1811,7 @@ impl SchemaApiTestSuite { table_name: tbl_name.to_string(), }, table_meta: table_meta(created_on), + lineage_updates: vec![], as_dropped: false, table_properties: None, table_partition: None, @@ -2061,6 +2063,7 @@ impl SchemaApiTestSuite { table_name: tbl_name.to_string(), }, table_meta: tbl_meta, + lineage_updates: vec![], as_dropped: true, table_properties: None, table_partition: None, @@ -2091,6 +2094,7 @@ impl SchemaApiTestSuite { table_id: create_table_as_dropped_resp.table_id, prev_table_id: create_table_as_dropped_resp.prev_table_id, orphan_table_name: create_table_as_dropped_resp.orphan_table_name.clone(), + lineage_updates: vec![], }; mt.commit_table_meta(commit_table_req).await?; @@ -4360,6 +4364,7 @@ impl SchemaApiTestSuite { catalog_name: None, name_ident, table_meta: create_table_meta.clone(), + lineage_updates: vec![], as_dropped: false, table_properties: None, table_partition: None, @@ -5187,6 +5192,7 @@ impl SchemaApiTestSuite { }, table_meta: table_meta(created_on), + lineage_updates: vec![], as_dropped: false, table_properties: None, table_partition: None, @@ -5675,6 +5681,7 @@ impl SchemaApiTestSuite { table_name: tbl_name.to_string(), }, table_meta: drop_table_meta(created_on), + lineage_updates: vec![], as_dropped: true, table_properties: None, table_partition: None, @@ -5690,6 +5697,7 @@ impl SchemaApiTestSuite { table_id: create_table_as_dropped_resp.table_id, prev_table_id: Some(1111), orphan_table_name: create_table_as_dropped_resp.orphan_table_name.clone(), + lineage_updates: vec![], }; let resp = mt.commit_table_meta(commit_table_req).await; use databend_common_meta_app::app_error::AppError; @@ -5725,6 +5733,7 @@ impl SchemaApiTestSuite { table_id: create_table_as_dropped_resp.table_id, prev_table_id: create_table_as_dropped_resp.prev_table_id, orphan_table_name: create_table_as_dropped_resp.orphan_table_name.clone(), + lineage_updates: vec![], }; let resp = mt.commit_table_meta(commit_table_req).await; use databend_common_meta_app::app_error::AppError; @@ -5773,6 +5782,7 @@ impl SchemaApiTestSuite { table_id: create_table_as_dropped_resp.table_id, prev_table_id: create_table_as_dropped_resp.prev_table_id, orphan_table_name: create_table_as_dropped_resp.orphan_table_name.clone(), + lineage_updates: vec![], }; let resp = mt.commit_table_meta(commit_table_req).await; use databend_common_meta_app::app_error::AppError; @@ -5795,6 +5805,7 @@ impl SchemaApiTestSuite { table_name: tbl_name.to_string(), }, table_meta: table_meta(created_on), + lineage_updates: vec![], as_dropped: true, table_properties: None, table_partition: None, @@ -5815,6 +5826,7 @@ impl SchemaApiTestSuite { table_name: tbl_name.to_string(), }, table_meta: drop_table_meta(created_on), + lineage_updates: vec![], as_dropped: true, table_properties: None, table_partition: None, @@ -5828,6 +5840,7 @@ impl SchemaApiTestSuite { table_id: create_table_as_dropped_resp.table_id, prev_table_id: create_table_as_dropped_resp.prev_table_id, orphan_table_name: create_table_as_dropped_resp.orphan_table_name.clone(), + lineage_updates: vec![], }; mt.commit_table_meta(commit_table_req).await?; } @@ -5851,6 +5864,7 @@ impl SchemaApiTestSuite { table_name: tbl_name.to_string(), }, table_meta: drop_table_meta(created_on), + lineage_updates: vec![], as_dropped: true, table_properties: None, table_partition: None, @@ -5944,6 +5958,7 @@ impl SchemaApiTestSuite { table_name: tbl_name.to_string(), }, table_meta: drop_table_meta(created_on), + lineage_updates: vec![], as_dropped: true, table_properties: None, table_partition: None, @@ -5978,6 +5993,7 @@ impl SchemaApiTestSuite { table_id: resp.table_id, prev_table_id: resp.prev_table_id, orphan_table_name: resp.orphan_table_name.clone(), + lineage_updates: vec![], }; let resp = arc_mt.commit_table_meta(commit_table_req).await; 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..717d5cb728a79 100644 --- a/src/query/config/src/config.rs +++ b/src/query/config/src/config.rs @@ -1959,6 +1959,10 @@ pub struct QueryConfig { /// Max number of async table hook jobs running concurrently. #[clap(long, value_name = "VALUE", default_value = "2")] pub table_hook_async_max_concurrency: usize, + + /// Lineage collection config. + #[clap(skip)] + pub lineage: LineageConfig, } impl Default for QueryConfig { @@ -1974,6 +1978,27 @@ impl Default for QueryConfig { } } +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Args)] +#[serde(default)] +pub struct LineageConfig { + /// Whether to collect lineage for successful DDL/DML commits. + #[clap( + long, + value_name = "VALUE", + value_parser = clap::value_parser!(bool), + default_value = "false" + )] + pub capture_enabled: bool, +} + +impl Default for LineageConfig { + fn default() -> Self { + Self { + capture_enabled: false, + } + } +} + impl TryInto for QueryConfig { type Error = ErrorCode; diff --git a/src/query/ee/src/attach_table/handler.rs b/src/query/ee/src/attach_table/handler.rs index 31687b69ee188..52ba91b8a3923 100644 --- a/src/query/ee/src/attach_table/handler.rs +++ b/src/query/ee/src/attach_table/handler.rs @@ -121,6 +121,7 @@ impl AttachTableHandler for RealAttachTableHandler { table_name: plan.table.to_string(), }, table_meta, + lineage_updates: vec![], as_dropped: false, table_properties: None, table_partition: None, diff --git a/src/query/ee/src/stream/handler.rs b/src/query/ee/src/stream/handler.rs index 52f3ccae65c23..2d92f3204bf5c 100644 --- a/src/query/ee/src/stream/handler.rs +++ b/src/query/ee/src/stream/handler.rs @@ -158,6 +158,7 @@ impl StreamHandler for RealStreamHandler { comment: plan.comment.clone().unwrap_or("".to_string()), ..Default::default() }, + lineage_updates: vec![], as_dropped: false, table_properties: None, table_partition: None, diff --git a/src/query/ee/tests/it/license/license_mgr.rs b/src/query/ee/tests/it/license/license_mgr.rs index 1f967976049e9..816c0b6cb257e 100644 --- a/src/query/ee/tests/it/license/license_mgr.rs +++ b/src/query/ee/tests/it/license/license_mgr.rs @@ -102,6 +102,7 @@ async fn test_license_features() -> databend_common_exception::Result<()> { Feature::LicenseInfo, Feature::Vacuum, Feature::Stream, + Feature::Lineage, ]), ), Duration::from_hours(2), @@ -137,7 +138,13 @@ async fn test_license_features() -> databend_common_exception::Result<()> { assert!( license_mgr - .check_enterprise_enabled(token, Feature::Stream) + .check_enterprise_enabled(token.clone(), Feature::Stream) + .is_ok() + ); + + assert!( + license_mgr + .check_enterprise_enabled(token, Feature::Lineage) .is_ok() ); diff --git a/src/query/service/src/interpreters/common/lineage_writer.rs b/src/query/service/src/interpreters/common/lineage_writer.rs new file mode 100644 index 0000000000000..ec2a5bd2752c2 --- /dev/null +++ b/src/query/service/src/interpreters/common/lineage_writer.rs @@ -0,0 +1,605 @@ +// 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::HashMap; +use std::sync::Arc; + +use chrono::Utc; +use databend_common_base::runtime::GlobalIORuntime; +use databend_common_catalog::table::is_temp_table_by_table_info; +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; +use databend_common_meta_api::LineageApi; +use databend_common_meta_api::MergeLineageReq; +use databend_common_meta_app::schema::ColumnRef; +use databend_common_meta_app::schema::LineageColumn; +use databend_common_meta_app::schema::LineageDetail; +use databend_common_meta_app::schema::LineageIdentity; +use databend_common_meta_app::schema::LineageKind; +use databend_common_meta_app::schema::LineageObjectRef; +use databend_common_meta_app::schema::LineageObjectType; +use databend_common_meta_app::schema::LineageUpdate; +use databend_common_meta_app::schema::LineageUpdateMode; +use databend_common_meta_app::schema::TableInfo; +use databend_common_pipeline::core::ExecutionInfo; +use databend_common_pipeline::core::Pipeline; +use databend_common_pipeline::core::basic_callback; +use databend_common_sql::LineageUpstream; +use databend_common_sql::QueryLineage; +use databend_common_sql::QueryLineageColumnEdge; +use databend_common_sql::QueryLineageKind; +use databend_common_sql::QueryLineageRelation; +use databend_common_users::UserApiProvider; +use log::warn; + +use crate::sessions::QueryContext; +use crate::sessions::TableContext; + +pub async fn build_query_lineage_updates( + ctx: &Arc, + target_table_info: Option<&TableInfo>, +) -> Result> { + let query_lineage = ctx.get_query_lineage(); + build_query_lineage_updates_with_lineage( + ctx.as_ref(), + query_lineage.as_deref(), + target_table_info, + ) + .await +} + +pub async fn build_query_lineage_updates_with_lineage( + ctx: &dyn TableContext, + query_lineage: Option<&QueryLineage>, + target_table_info: Option<&TableInfo>, +) -> Result> { + let Some(query_lineage) = query_lineage else { + return Ok(vec![]); + }; + + build_lineage_updates(ctx, query_lineage, target_table_info).await +} + +pub fn attach_query_lineage_on_finished(pipeline: &mut Pipeline, updates: Vec) { + let updates = updates + .into_iter() + .filter(|update| update.mode == LineageUpdateMode::Merge) + .collect::>(); + if updates.is_empty() { + return; + } + + pipeline.set_on_finished(basic_callback(move |info: &ExecutionInfo| { + if info.res.is_ok() { + let _: Result<()> = GlobalIORuntime::instance().block_on(async move { + let meta = UserApiProvider::instance().get_meta_store_client(); + handle_lineage_persistence_result( + meta.merge_lineage(MergeLineageReq { updates }).await, + ) + }); + } + + Ok(()) + })); +} + +fn handle_lineage_persistence_result( + result: std::result::Result<(), E>, +) -> Result<()> { + if let Err(error) = result { + warn!("failed to persist query lineage: {error}"); + } + Ok(()) +} + +async fn build_lineage_updates( + ctx: &dyn TableContext, + query_lineage: &QueryLineage, + target_table_info: Option<&TableInfo>, +) -> Result> { + let kind = lineage_kind(query_lineage.kind); + let mode = lineage_update_mode(query_lineage.kind); + let mut updates = Vec::new(); + + for downstream in &query_lineage.downstreams { + let downstream_relation = if let Some(table_info) = target_table_info { + if !is_lineage_supported_table_info(table_info) { + warn!( + "skip query lineage target relation: kind={:?}, relation={}, reason=unsupported table engine '{}'", + query_lineage.kind, + full_table_name(&downstream.relation), + table_info.engine(), + ); + continue; + } + ResolvedRelation::from_table_info(&downstream.relation, table_info) + } else if query_lineage.kind == QueryLineageKind::CreateView { + ResolvedRelation::unresolved(downstream.relation.clone()) + } else { + match ResolvedRelation::resolve(ctx, &downstream.relation, true).await? { + ResolveRelationResult::Resolved(relation) => relation, + ResolveRelationResult::Skipped(reason) => { + warn!( + "skip query lineage target relation: kind={:?}, relation={}, reason={}", + query_lineage.kind, + full_table_name(&downstream.relation), + reason, + ); + continue; + } + } + }; + + for upstream in &downstream.upstreams { + let require_upstream_id_match = query_lineage.kind != QueryLineageKind::CreateView; + let upstream_relation = + ResolvedRelation::resolve(ctx, &upstream.relation, require_upstream_id_match) + .await?; + let upstream_relation = match upstream_relation { + ResolveRelationResult::Resolved(relation) => relation, + ResolveRelationResult::Skipped(reason) => { + warn!( + "skip query lineage upstream relation: kind={:?}, relation={}, reason={}", + query_lineage.kind, + full_table_name(&upstream.relation), + reason, + ); + continue; + } + }; + if upstream_relation.is_same_object(&downstream_relation) { + continue; + } + + let detail = LineageDetail { + kind: kind.clone(), + last_query_id: Some(ctx.get_id()), + updated_on: Utc::now(), + column_lineage: column_lineage( + query_lineage.kind, + &upstream_relation, + &downstream_relation, + upstream, + ), + }; + + updates.push(LineageUpdate { + tenant_name: ctx.get_tenant().tenant_name().to_string(), + upstream: upstream_relation.object_ref(query_lineage.kind, true), + downstream: downstream_relation.object_ref(query_lineage.kind, false), + detail, + mode: mode.clone(), + }); + } + } + + Ok(updates) +} + +#[derive(Clone)] +struct ResolvedRelation { + relation: QueryLineageRelation, + current_table_id: Option, + stable_table_id: bool, + column_ids_by_name: HashMap, + column_names_by_id: HashMap, +} + +enum ResolveRelationResult { + Resolved(ResolvedRelation), + Skipped(String), +} + +impl ResolvedRelation { + fn unresolved(relation: QueryLineageRelation) -> Self { + Self { + relation, + current_table_id: None, + stable_table_id: false, + column_ids_by_name: HashMap::new(), + column_names_by_id: HashMap::new(), + } + } + + fn from_table_info(relation: &QueryLineageRelation, table_info: &TableInfo) -> Self { + let mut column_ids_by_name = HashMap::new(); + let mut column_names_by_id = HashMap::new(); + for field in table_info.schema().fields() { + column_ids_by_name.insert(field.name().clone(), u64::from(field.column_id)); + column_names_by_id.insert(u64::from(field.column_id), field.name().clone()); + } + + Self { + relation: relation.clone(), + current_table_id: Some(table_info.ident.table_id), + stable_table_id: has_stable_lineage_table_id(table_info), + column_ids_by_name, + column_names_by_id, + } + } + + async fn resolve( + ctx: &dyn TableContext, + relation: &QueryLineageRelation, + require_id_match: bool, + ) -> Result { + let table = match ctx + .get_table(&relation.catalog, &relation.database, &relation.name) + .await + { + Ok(table) => table, + Err(error) if is_unknown_relation_error(&error) => { + return Ok(ResolveRelationResult::Skipped(error.message())); + } + Err(error) => return Err(error), + }; + if table.is_temp() { + return Ok(ResolveRelationResult::Skipped( + "temporary table is not captured".to_string(), + )); + } + let table_info = table.get_table_info(); + if !is_lineage_supported_table_info(table_info) { + return Ok(ResolveRelationResult::Skipped(format!( + "unsupported table engine '{}'", + table_info.engine() + ))); + } + let current_table_id = table_info.ident.table_id; + let stable_table_id = has_stable_lineage_table_id(table_info); + + if require_id_match + && stable_table_id + && relation.id.is_some_and(|id| id != current_table_id) + { + return Ok(ResolveRelationResult::Skipped(format!( + "table id changed, expected {}, got {}", + relation.id.unwrap(), + current_table_id + ))); + } + + let mut column_ids_by_name = HashMap::new(); + let mut column_names_by_id = HashMap::new(); + for field in table.schema().fields() { + column_ids_by_name.insert(field.name().clone(), u64::from(field.column_id)); + column_names_by_id.insert(u64::from(field.column_id), field.name().clone()); + } + + Ok(ResolveRelationResult::Resolved(Self { + relation: relation.clone(), + current_table_id: Some(current_table_id), + stable_table_id, + column_ids_by_name, + column_names_by_id, + })) + } + + fn object_ref(&self, kind: QueryLineageKind, is_upstream: bool) -> LineageObjectRef { + let identity = if self.uses_column_name_ref(kind, is_upstream) { + LineageIdentity::Name { + name: full_table_name(&self.relation), + } + } else { + LineageIdentity::Id { + id: self.current_table_id.unwrap().to_string(), + } + }; + + LineageObjectRef { + object_type: LineageObjectType::Table, + identity, + } + } + + fn uses_column_name_ref(&self, kind: QueryLineageKind, is_upstream: bool) -> bool { + (is_upstream && (kind == QueryLineageKind::CreateView || self.relation.id.is_none())) + || !self.stable_table_id + || self.current_table_id.is_none() + } + + fn is_same_object(&self, other: &Self) -> bool { + if self.stable_table_id && other.stable_table_id { + if let (Some(self_id), Some(other_id)) = (self.current_table_id, other.current_table_id) + { + return self_id == other_id; + } + } + + self.relation.catalog == other.relation.catalog + && self.relation.database == other.relation.database + && self.relation.name == other.relation.name + } +} + +fn is_unknown_relation_error(error: &ErrorCode) -> bool { + matches!( + error.code(), + ErrorCode::UNKNOWN_CATALOG + | ErrorCode::UNKNOWN_DATABASE + | ErrorCode::UNKNOWN_TABLE + | ErrorCode::UNKNOWN_VIEW + ) +} + +fn is_lineage_supported_table_info(table_info: &TableInfo) -> bool { + !is_temp_table_by_table_info(table_info) + && !matches!( + table_info.engine().to_ascii_uppercase().as_str(), + "MEMORY" | "DELTA" + ) +} + +fn has_stable_lineage_table_id(table_info: &TableInfo) -> bool { + table_info.ident.table_id != 0 + && !matches!( + table_info.engine().to_ascii_uppercase().as_str(), + "ICEBERG" | "HIVE" + ) +} + +fn lineage_kind(kind: QueryLineageKind) -> LineageKind { + match kind { + QueryLineageKind::CreateView => LineageKind::View, + QueryLineageKind::Ctas => LineageKind::Ctas, + QueryLineageKind::Dml => LineageKind::DataMovement, + } +} + +fn lineage_update_mode(kind: QueryLineageKind) -> LineageUpdateMode { + match kind { + QueryLineageKind::Dml => LineageUpdateMode::Merge, + QueryLineageKind::CreateView | QueryLineageKind::Ctas => LineageUpdateMode::Replace, + } +} + +fn column_lineage( + kind: QueryLineageKind, + upstream_relation: &ResolvedRelation, + downstream_relation: &ResolvedRelation, + upstream: &LineageUpstream, +) -> Vec { + upstream + .columns + .iter() + .filter_map( + |edge| match column_edge(kind, upstream_relation, downstream_relation, edge) { + Ok(column) => Some(column), + Err(reason) => { + warn!( + "skip query lineage column edge: kind={:?}, upstream_relation={}, downstream_relation={}, upstream_column={}({}), downstream_column={}({}), reason={}", + kind, + full_table_name(&upstream_relation.relation), + full_table_name(&downstream_relation.relation), + edge.upstream.name, + edge.upstream.id, + edge.downstream.name, + edge.downstream.id, + reason, + ); + None + } + }, + ) + .collect() +} + +fn column_edge( + kind: QueryLineageKind, + upstream_relation: &ResolvedRelation, + downstream_relation: &ResolvedRelation, + edge: &QueryLineageColumnEdge, +) -> std::result::Result { + let upstream = if upstream_relation.uses_column_name_ref(kind, true) { + upstream_relation + .column_ids_by_name + .contains_key(&edge.upstream.name) + .then(|| ColumnRef::Name(edge.upstream.name.clone())) + .ok_or_else(|| format!("upstream column '{}' is not found", edge.upstream.name))? + } else { + upstream_relation + .column_names_by_id + .contains_key(&u64::from(edge.upstream.id)) + .then_some(ColumnRef::Id(u64::from(edge.upstream.id))) + .ok_or_else(|| format!("upstream column id {} is not found", edge.upstream.id))? + }; + + let downstream = if downstream_relation.uses_column_name_ref(kind, false) { + (downstream_relation.column_ids_by_name.is_empty() + || downstream_relation + .column_ids_by_name + .contains_key(&edge.downstream.name)) + .then(|| ColumnRef::Name(edge.downstream.name.clone())) + .ok_or_else(|| format!("downstream column '{}' is not found", edge.downstream.name))? + } else { + downstream_relation + .column_ids_by_name + .get(&edge.downstream.name) + .copied() + .or_else(|| { + downstream_relation + .column_names_by_id + .contains_key(&u64::from(edge.downstream.id)) + .then_some(u64::from(edge.downstream.id)) + }) + .map(ColumnRef::Id) + .ok_or_else(|| { + format!( + "downstream column '{}'({}) is not found", + edge.downstream.name, edge.downstream.id + ) + })? + }; + + Ok(LineageColumn { + upstream, + downstream, + }) +} + +fn full_table_name(relation: &QueryLineageRelation) -> String { + match relation.catalog.as_str() { + "default" => format!("{}.{}", relation.database, relation.name), + _ => format!( + "{}.{}.{}", + relation.catalog, relation.database, relation.name + ), + } +} + +#[cfg(test)] +mod tests { + use databend_common_sql::QueryLineageColumn; + use databend_common_sql::QueryLineageRelationKind; + + use super::*; + + #[test] + fn test_column_lineage_skips_unexpressible_column_edge() { + let upstream = resolved_relation(Some(1), true, vec![("a", 11)]); + let downstream = resolved_relation(Some(2), true, vec![("b", 22)]); + let from = LineageUpstream { + relation: upstream.relation.clone(), + columns: vec![QueryLineageColumnEdge { + upstream: column("a", 11), + downstream: column("missing", 999), + }], + }; + + let got = column_lineage(QueryLineageKind::Ctas, &upstream, &downstream, &from); + + assert!(got.is_empty()); + } + + #[test] + fn test_column_lineage_keeps_object_edge_when_columns_are_empty() { + let upstream = resolved_relation(Some(1), true, vec![("a", 11)]); + let downstream = resolved_relation(Some(2), true, vec![("b", 22)]); + let from = LineageUpstream { + relation: upstream.relation.clone(), + columns: vec![], + }; + + let got = column_lineage(QueryLineageKind::Ctas, &upstream, &downstream, &from); + + assert!(got.is_empty()); + } + + #[test] + fn test_column_edge_uses_name_for_unresolved_downstream() { + let upstream = resolved_relation(Some(1), true, vec![("a", 11)]); + let downstream = resolved_relation(None, false, vec![]); + let edge = QueryLineageColumnEdge { + upstream: column("a", 11), + downstream: column("v", 101), + }; + + let got = column_edge(QueryLineageKind::CreateView, &upstream, &downstream, &edge) + .expect("unresolved view target should use column name"); + + assert_eq!(got.upstream, ColumnRef::Name("a".to_string())); + assert_eq!(got.downstream, ColumnRef::Name("v".to_string())); + } + + #[test] + fn test_lineage_supported_table_engines() { + let mut memory = TableInfo::default(); + memory.meta.engine = "MEMORY".to_string(); + assert!(!is_lineage_supported_table_info(&memory)); + + let mut delta = TableInfo::default(); + delta.meta.engine = "DELTA".to_string(); + assert!(!is_lineage_supported_table_info(&delta)); + + let mut fuse = TableInfo::default(); + fuse.meta.engine = "FUSE".to_string(); + assert!(is_lineage_supported_table_info(&fuse)); + } + + #[test] + fn test_unknown_relation_errors_are_skippable() { + assert!(is_unknown_relation_error(&ErrorCode::UnknownTable("t"))); + assert!(is_unknown_relation_error(&ErrorCode::UnknownDatabase("db"))); + assert!(is_unknown_relation_error(&ErrorCode::UnknownCatalog( + "catalog" + ))); + assert!(!is_unknown_relation_error(&ErrorCode::MetaServiceError( + "meta error" + ))); + } + + #[test] + fn test_lineage_persistence_failure_is_best_effort() { + assert!(handle_lineage_persistence_result::<&str>(Err("meta unavailable")).is_ok()); + } + + #[test] + fn test_same_lineage_object_prefers_stable_table_id() { + let relation = resolved_relation(Some(1), true, vec![]); + let same = resolved_relation(Some(1), true, vec![]); + let replaced = resolved_relation(Some(2), true, vec![]); + + assert!(relation.is_same_object(&same)); + assert!(!relation.is_same_object(&replaced)); + } + + #[test] + fn test_same_lineage_object_falls_back_to_name() { + let relation = resolved_relation(None, false, vec![]); + let same = resolved_relation(None, false, vec![]); + let mut other = resolved_relation(None, false, vec![]); + other.relation.name = "other".to_string(); + + assert!(relation.is_same_object(&same)); + assert!(!relation.is_same_object(&other)); + } + + fn resolved_relation( + current_table_id: Option, + stable_table_id: bool, + columns: Vec<(&str, u64)>, + ) -> ResolvedRelation { + let mut column_ids_by_name = HashMap::new(); + let mut column_names_by_id = HashMap::new(); + for (name, id) in columns { + column_ids_by_name.insert(name.to_string(), id); + column_names_by_id.insert(id, name.to_string()); + } + + ResolvedRelation { + relation: relation("t", current_table_id), + current_table_id, + stable_table_id, + column_ids_by_name, + column_names_by_id, + } + } + + fn relation(name: &str, id: Option) -> QueryLineageRelation { + QueryLineageRelation { + catalog: "default".to_string(), + database: "db".to_string(), + name: name.to_string(), + id, + kind: QueryLineageRelationKind::Table, + } + } + + fn column(name: &str, id: u64) -> QueryLineageColumn { + QueryLineageColumn { + name: name.to_string(), + id: id as u32, + } + } +} diff --git a/src/query/service/src/interpreters/common/mod.rs b/src/query/service/src/interpreters/common/mod.rs index f342544b52e5c..769b5a65a3561 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_writer; mod metrics; mod notification; mod query_log; @@ -29,6 +30,9 @@ pub mod table_option_validation; pub use column::*; pub use finish_hook::QueryFinishHooks; pub use grant::validate_grant_object_exists; +pub use lineage_writer::attach_query_lineage_on_finished; +pub use lineage_writer::build_query_lineage_updates; +pub use lineage_writer::build_query_lineage_updates_with_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_copy_into_table.rs b/src/query/service/src/interpreters/interpreter_copy_into_table.rs index 5efd01cbfdecc..44ab716a1e28a 100644 --- a/src/query/service/src/interpreters/interpreter_copy_into_table.rs +++ b/src/query/service/src/interpreters/interpreter_copy_into_table.rs @@ -67,6 +67,8 @@ use log::info; use crate::interpreters::HookOperator; use crate::interpreters::Interpreter; use crate::interpreters::SelectInterpreter; +use crate::interpreters::common::attach_query_lineage_on_finished; +use crate::interpreters::common::build_query_lineage_updates; use crate::interpreters::common::check_deduplicate_label; use crate::interpreters::common::dml_build_update_stream_req; use crate::physical_plans::CopyIntoTable; @@ -874,6 +876,7 @@ impl CopyIntoTableInterpreter { path_prefix, )?; + let lineage_updates = build_query_lineage_updates(&ctx, None).await?; to_table.commit_insertion( ctx.clone(), main_pipeline, @@ -884,6 +887,7 @@ impl CopyIntoTableInterpreter { deduplicated_label, table_meta_timestamps, )?; + attach_query_lineage_on_finished(main_pipeline, lineage_updates); } // Purge files. diff --git a/src/query/service/src/interpreters/interpreter_factory.rs b/src/query/service/src/interpreters/interpreter_factory.rs index 427784b0e7c3d..3c237b1c64fe2 100644 --- a/src/query/service/src/interpreters/interpreter_factory.rs +++ b/src/query/service/src/interpreters/interpreter_factory.rs @@ -16,6 +16,7 @@ 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; @@ -153,6 +154,7 @@ impl InterpreterFactory { error!("Access.denied(v2): {:?}", e); } })?; + initialize_query_lineage(&ctx, plan); let mut access_logger = AccessLogger::create(ctx.clone()); access_logger.log(plan); access_logger.output(); @@ -937,3 +939,24 @@ impl InterpreterFactory { } } } + +fn initialize_query_lineage(ctx: &QueryContext, plan: &Plan) { + ctx.init_query_lineage(|| { + if !GlobalConfig::instance() + .query + .common + .lineage + .capture_enabled + { + return None; + } + + match plan.query_lineage() { + Ok(lineage) => lineage, + Err(error) => { + log::warn!("failed to extract query lineage: {error:?}"); + None + } + } + }); +} diff --git a/src/query/service/src/interpreters/interpreter_insert.rs b/src/query/service/src/interpreters/interpreter_insert.rs index df8cea1bcf509..9066354cbb757 100644 --- a/src/query/service/src/interpreters/interpreter_insert.rs +++ b/src/query/service/src/interpreters/interpreter_insert.rs @@ -61,6 +61,8 @@ use crate::clusters::ClusterHelper; use crate::interpreters::HookOperator; use crate::interpreters::Interpreter; use crate::interpreters::InterpreterPtr; +use crate::interpreters::common::attach_query_lineage_on_finished; +use crate::interpreters::common::build_query_lineage_updates; use crate::interpreters::common::check_deduplicate_label; use crate::interpreters::common::dml_build_update_stream_req; use crate::physical_plans::ConstantTableScan; @@ -468,6 +470,8 @@ impl Interpreter for InsertInterpreter { build_query_pipeline_without_render_result_set(&self.ctx, &insert_select_plan) .await?; + let lineage_updates = + build_query_lineage_updates(&self.ctx, Some(table.get_table_info())).await?; table.commit_insertion( self.ctx.clone(), &mut build_res.main_pipeline, @@ -478,6 +482,7 @@ impl Interpreter for InsertInterpreter { unsafe { self.ctx.get_settings().get_deduplicate_label()? }, table_meta_timestamps, )?; + attach_query_lineage_on_finished(&mut build_res.main_pipeline, lineage_updates); // Execute the hook operator. if self.plan.branch.is_none() { 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..a626fcbe0f7fa 100644 --- a/src/query/service/src/interpreters/interpreter_insert_multi_table.rs +++ b/src/query/service/src/interpreters/interpreter_insert_multi_table.rs @@ -44,6 +44,8 @@ use databend_common_storages_fuse::FuseTable; use super::HookOperator; use crate::interpreters::Interpreter; use crate::interpreters::InterpreterPtr; +use crate::interpreters::common::attach_query_lineage_on_finished; +use crate::interpreters::common::build_query_lineage_updates; use crate::interpreters::common::dml_build_update_stream_req; use crate::physical_plans::CastSchema; use crate::physical_plans::ChunkAppendData; @@ -107,6 +109,9 @@ impl Interpreter for InsertMultiTableInterpreter { let physical_plan = self.build_physical_plan(false).await?; let mut build_res = build_query_pipeline_without_render_result_set(&self.ctx, &physical_plan).await?; + let lineage_updates = build_query_lineage_updates(&self.ctx, None).await?; + attach_query_lineage_on_finished(&mut build_res.main_pipeline, lineage_updates); + // Execute hook. if self .ctx diff --git a/src/query/service/src/interpreters/interpreter_mutation.rs b/src/query/service/src/interpreters/interpreter_mutation.rs index 1ac94e0c36abf..72e700c084621 100644 --- a/src/query/service/src/interpreters/interpreter_mutation.rs +++ b/src/query/service/src/interpreters/interpreter_mutation.rs @@ -178,6 +178,7 @@ impl MutationInterpreter { let mut builder = PhysicalPlanBuilder::new(mutation.metadata.clone(), self.ctx.clone(), dry_run); builder.set_mutation_build_info(mutation_build_info); + builder.set_query_lineage(self.ctx.get_query_lineage()); builder .build(&self.s_expr, *mutation.required_columns.clone()) .await diff --git a/src/query/service/src/interpreters/interpreter_replace.rs b/src/query/service/src/interpreters/interpreter_replace.rs index 0ad331ec7dd5f..7ebcbf6561f94 100644 --- a/src/query/service/src/interpreters/interpreter_replace.rs +++ b/src/query/service/src/interpreters/interpreter_replace.rs @@ -47,6 +47,7 @@ use crate::interpreters::HookOperator; use crate::interpreters::Interpreter; use crate::interpreters::InterpreterPtr; use crate::interpreters::SelectInterpreter; +use crate::interpreters::common::build_query_lineage_updates; use crate::interpreters::common::check_deduplicate_label; use crate::interpreters::common::dml_build_update_stream_req; #[cfg(feature = "storage-stage")] @@ -358,6 +359,7 @@ impl ReplaceInterpreter { }); } + let lineage_updates = build_query_lineage_updates(&self.ctx, None).await?; root = PhysicalPlan::new(CommitSink { input: root, snapshot: base_snapshot, @@ -370,6 +372,7 @@ impl ReplaceInterpreter { deduplicated_label: unsafe { self.ctx.get_settings().get_deduplicate_label()? }, table_meta_timestamps, recluster_info: None, + lineage_updates, meta: PhysicalPlanMeta::new("CommitSink"), }); diff --git a/src/query/service/src/interpreters/interpreter_table_create.rs b/src/query/service/src/interpreters/interpreter_table_create.rs index bfe7c753e8f3e..9fd3b27925406 100644 --- a/src/query/service/src/interpreters/interpreter_table_create.rs +++ b/src/query/service/src/interpreters/interpreter_table_create.rs @@ -69,6 +69,7 @@ use log::info; use crate::interpreters::InsertInterpreter; use crate::interpreters::Interpreter; +use crate::interpreters::common::build_query_lineage_updates; use crate::interpreters::common::table_option_validation::is_valid_analyze_count_min_sketch_error_rate; use crate::interpreters::common::table_option_validation::is_valid_analyze_frequency_columns; use crate::interpreters::common::table_option_validation::is_valid_analyze_histogram_algorithm; @@ -240,6 +241,7 @@ impl CreateTableInterpreter { TableIdent::new(table_id, table_id_seq), table_meta, ); + let lineage_updates = build_query_lineage_updates(&self.ctx, Some(&table_info)).await?; let insert_plan = Insert { catalog: self.plan.catalog.clone(), @@ -250,6 +252,7 @@ impl CreateTableInterpreter { overwrite: false, source: InsertInputSource::SelectPlan(select_plan), table_info: Some(table_info), + lineage_target_table_id: None, }; let mut pipeline = InsertInterpreter::try_create(self.ctx.clone(), insert_plan)? @@ -287,6 +290,7 @@ impl CreateTableInterpreter { "create_table_as_select {} success, commit table meta data by table id {}", qualified_table_name, table_id ); + let lineage_updates = lineage_updates.clone(); let fut = async move { let req = CommitTableMetaReq { name_ident: TableNameIdent { @@ -298,6 +302,7 @@ impl CreateTableInterpreter { table_id, prev_table_id, orphan_table_name, + lineage_updates, }; catalog.commit_table_meta(req).await }; @@ -535,6 +540,7 @@ impl CreateTableInterpreter { table_name: self.plan.table.to_string(), }, table_meta, + lineage_updates: vec![], as_dropped: false, table_properties: self.plan.table_properties.clone(), table_partition: self.plan.table_partition.as_ref().map(|table_partition| { diff --git a/src/query/service/src/interpreters/interpreter_table_recluster.rs b/src/query/service/src/interpreters/interpreter_table_recluster.rs index 53df92b6250d4..11a489cc833af 100644 --- a/src/query/service/src/interpreters/interpreter_table_recluster.rs +++ b/src/query/service/src/interpreters/interpreter_table_recluster.rs @@ -440,6 +440,7 @@ impl ReclusterTableInterpreter { deduplicated_label: None, table_meta_timestamps, recluster_info, + lineage_updates: vec![], meta: PhysicalPlanMeta::new("CommitSink"), }) } diff --git a/src/query/service/src/interpreters/interpreter_view_create.rs b/src/query/service/src/interpreters/interpreter_view_create.rs index 7a36c1fdcfbbb..f59929d2dfdf2 100644 --- a/src/query/service/src/interpreters/interpreter_view_create.rs +++ b/src/query/service/src/interpreters/interpreter_view_create.rs @@ -27,6 +27,7 @@ use databend_common_storages_basic::view_table::QUERY; use databend_common_storages_basic::view_table::VIEW_ENGINE; use crate::interpreters::Interpreter; +use crate::interpreters::common::build_query_lineage_updates; use crate::interpreters::util::check_view_circular_dependency; use crate::pipelines::PipelineBuildResult; use crate::sessions::QueryContext; @@ -115,6 +116,7 @@ impl Interpreter for CreateViewInterpreter { }; options.insert(QUERY.to_string(), subquery); + let lineage_updates = build_query_lineage_updates(&self.ctx, None).await?; let plan = CreateTableReq { create_option: self.plan.create_option, catalog_name: if self.plan.create_option.is_overriding() { @@ -132,6 +134,7 @@ impl Interpreter for CreateViewInterpreter { options, ..Default::default() }, + lineage_updates, as_dropped: false, table_properties: None, table_partition: None, diff --git a/src/query/service/src/physical_plans/physical_commit_sink.rs b/src/query/service/src/physical_plans/physical_commit_sink.rs index cac8d733ce4f1..c71a53bc8a2ab 100644 --- a/src/query/service/src/physical_plans/physical_commit_sink.rs +++ b/src/query/service/src/physical_plans/physical_commit_sink.rs @@ -18,6 +18,7 @@ use std::sync::Arc; use databend_common_catalog::plan::ReclusterInfoSideCar; use databend_common_exception::Result; use databend_common_expression::DataSchemaRef; +use databend_common_meta_app::schema::LineageUpdate; use databend_common_meta_app::schema::TableInfo; use databend_common_meta_app::schema::UpdateStreamMetaReq; use databend_common_pipeline::core::ExecutionInfo; @@ -35,6 +36,7 @@ use databend_storages_common_table_meta::meta::TableMetaTimestamps; use databend_storages_common_table_meta::meta::TableSnapshot; use databend_storages_common_table_meta::readers::snapshot_reader::TableSnapshotAccessor; +use crate::interpreters::common::attach_query_lineage_on_finished; use crate::physical_plans::physical_plan::IPhysicalPlan; use crate::physical_plans::physical_plan::PhysicalPlan; use crate::physical_plans::physical_plan::PhysicalPlanMeta; @@ -52,6 +54,8 @@ pub struct CommitSink { pub commit_type: CommitType, pub update_stream_meta: Vec, pub deduplicated_label: Option, + // Best-effort lineage updates are attached after the table commit sink is built. + pub lineage_updates: Vec, pub table_meta_timestamps: TableMetaTimestamps, // Used for recluster. @@ -95,6 +99,7 @@ impl IPhysicalPlan for CommitSink { commit_type: self.commit_type.clone(), update_stream_meta: self.update_stream_meta.clone(), deduplicated_label: self.deduplicated_label.clone(), + lineage_updates: self.lineage_updates.clone(), table_meta_timestamps: self.table_meta_timestamps, recluster_info: self.recluster_info.clone(), }) @@ -211,7 +216,12 @@ impl IPhysicalPlan for CommitSink { self.deduplicated_label.clone(), self.table_meta_timestamps, ) - }) + })?; + attach_query_lineage_on_finished( + &mut builder.main_pipeline, + self.lineage_updates.clone(), + ); + Ok(()) } } } diff --git a/src/query/service/src/physical_plans/physical_compact_source.rs b/src/query/service/src/physical_plans/physical_compact_source.rs index db28e9857f532..bccfbec3693e2 100644 --- a/src/query/service/src/physical_plans/physical_compact_source.rs +++ b/src/query/service/src/physical_plans/physical_compact_source.rs @@ -286,6 +286,7 @@ impl PhysicalPlanBuilder { deduplicated_label: None, recluster_info: None, table_meta_timestamps, + lineage_updates: vec![], meta: PhysicalPlanMeta::new("CommitSink"), }); diff --git a/src/query/service/src/physical_plans/physical_mutation.rs b/src/query/service/src/physical_plans/physical_mutation.rs index e37402097604b..36d427b0608b2 100644 --- a/src/query/service/src/physical_plans/physical_mutation.rs +++ b/src/query/service/src/physical_plans/physical_mutation.rs @@ -73,6 +73,7 @@ use tokio::sync::Semaphore; use super::ColumnMutation; use super::CommitType; +use crate::interpreters::common::build_query_lineage_updates_with_lineage; use crate::physical_plans::CommitSink; use crate::physical_plans::Exchange; use crate::physical_plans::MutationManipulate; @@ -320,6 +321,12 @@ impl PhysicalPlanBuilder { if *truncate_table { // Do truncate. + let lineage_updates = build_query_lineage_updates_with_lineage( + self.ctx.as_ref(), + self.query_lineage.as_deref(), + None, + ) + .await?; plan = PhysicalPlan::new(CommitSink { input: plan, snapshot: mutation_build_info.table_snapshot, @@ -333,6 +340,7 @@ impl PhysicalPlanBuilder { recluster_info: None, meta: PhysicalPlanMeta::new("CommitSink"), table_meta_timestamps: mutation_build_info.table_meta_timestamps, + lineage_updates, }); plan.adjust_plan_id(&mut 0); return Ok(plan); @@ -412,6 +420,12 @@ impl PhysicalPlanBuilder { }); } + let lineage_updates = build_query_lineage_updates_with_lineage( + self.ctx.as_ref(), + self.query_lineage.as_deref(), + None, + ) + .await?; plan = PhysicalPlan::new(CommitSink { input: plan, snapshot: mutation_build_info.table_snapshot, @@ -426,6 +440,7 @@ impl PhysicalPlanBuilder { meta: PhysicalPlanMeta::new("CommitSink"), recluster_info: None, table_meta_timestamps: mutation_build_info.table_meta_timestamps, + lineage_updates, }); plan.adjust_plan_id(&mut 0); @@ -653,6 +668,12 @@ impl PhysicalPlanBuilder { }; // build mutation_aggregate + let lineage_updates = build_query_lineage_updates_with_lineage( + self.ctx.as_ref(), + self.query_lineage.as_deref(), + None, + ) + .await?; let mut physical_plan = PhysicalPlan::new(CommitSink { input: plan, snapshot: mutation_build_info.table_snapshot, @@ -667,6 +688,7 @@ impl PhysicalPlanBuilder { recluster_info: None, meta: PhysicalPlanMeta::new("CommitSink"), table_meta_timestamps: mutation_build_info.table_meta_timestamps, + lineage_updates, }); physical_plan.adjust_plan_id(&mut 0); diff --git a/src/query/service/src/physical_plans/physical_plan_builder.rs b/src/query/service/src/physical_plans/physical_plan_builder.rs index 8616c62826182..2f4d04f281ecd 100644 --- a/src/query/service/src/physical_plans/physical_plan_builder.rs +++ b/src/query/service/src/physical_plans/physical_plan_builder.rs @@ -24,6 +24,7 @@ use databend_common_meta_app::schema::TableInfo; use databend_common_meta_app::schema::UpdateStreamMetaReq; use databend_common_sql::ColumnSet; use databend_common_sql::MetadataRef; +use databend_common_sql::QueryLineage; use databend_common_sql::optimizer::ir::RelExpr; use databend_common_sql::optimizer::ir::SExpr; use databend_common_sql::plans::RelOperator; @@ -44,6 +45,7 @@ pub struct PhysicalPlanBuilder { pub cte_required_columns: HashMap, pub is_cte_required_columns_collected: bool, pub build_depth: usize, + pub(crate) query_lineage: Option>, } impl PhysicalPlanBuilder { @@ -58,6 +60,7 @@ impl PhysicalPlanBuilder { cte_required_columns: HashMap::new(), is_cte_required_columns_collected: false, build_depth: 0, + query_lineage: None, } } @@ -190,6 +193,10 @@ impl PhysicalPlanBuilder { self.mutation_build_info = Some(mutation_build_info); } + pub(crate) fn set_query_lineage(&mut self, query_lineage: Option>) { + self.query_lineage = query_lineage; + } + pub fn set_metadata(&mut self, metadata: MetadataRef) { self.metadata = metadata; } diff --git a/src/query/service/src/sessions/query_ctx.rs b/src/query/service/src/sessions/query_ctx.rs index e5dff77b98c44..e71290dff00b3 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; @@ -933,6 +934,16 @@ impl QueryContext { self.shared.table_meta_timestamps.lock().clear(); } + pub fn init_query_lineage(&self, init: impl FnOnce() -> Option) { + self.shared + .query_lineage + .get_or_init(|| init().map(Arc::new)); + } + + pub fn get_query_lineage(&self) -> Option> { + self.shared.query_lineage.get().cloned().flatten() + } + pub fn get_materialized_cte_senders( &self, cte_name: &str, diff --git a/src/query/service/src/sessions/query_ctx_shared.rs b/src/query/service/src/sessions/query_ctx_shared.rs index 9b9ca5f8d3abb..e0b56a12bc3d3 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; @@ -145,6 +146,7 @@ pub struct QueryContextShared { pub(super) user_agent: Arc>, pub(super) query_profiles: Arc, PlanProfile>>>, + pub(super) query_lineage: OnceLock>>, pub(super) runtime_filter_state: RuntimeFilterState, @@ -241,6 +243,7 @@ impl QueryContextShared { window_partition_spill_progress: Arc::new(Progress::create()), query_cache_metrics: DataCacheMetrics::new(), query_profiles: Arc::new(RwLock::new(HashMap::new())), + query_lineage: Default::default(), runtime_filter_state: Default::default(), merge_into_join: Default::default(), query_queued_duration: Arc::new(RwLock::new(Duration::from_secs(0))), diff --git a/src/query/service/src/table_functions/get_lineage.rs b/src/query/service/src/table_functions/get_lineage.rs new file mode 100644 index 0000000000000..4b7037e3ef7d8 --- /dev/null +++ b/src/query/service/src/table_functions/get_lineage.rs @@ -0,0 +1,1699 @@ +// 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::any::Any; +use std::collections::HashSet; +use std::sync::Arc; + +use chrono::DateTime; +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::TableSchemaRef; +use databend_common_expression::TableSchemaRefExt; +use databend_common_expression::infer_table_schema; +use databend_common_expression::types::Int32Type; +use databend_common_expression::types::StringType; +use databend_common_license::license::Feature; +use databend_common_license::license_manager::LicenseManagerSwitch; +use databend_common_meta_api::LineageApi; +use databend_common_meta_api::ListLineageReq; +use databend_common_meta_api::kv_pb_api::KVPbApi; +use databend_common_meta_app::schema::ColumnRef; +use databend_common_meta_app::schema::LineageColumn; +use databend_common_meta_app::schema::LineageDetail; +use databend_common_meta_app::schema::LineageDirection; +use databend_common_meta_app::schema::LineageIdentity; +use databend_common_meta_app::schema::LineageKey; +use databend_common_meta_app::schema::LineageKind; +use databend_common_meta_app::schema::LineageObjectRef; +use databend_common_meta_app::schema::LineageObjectType; +use databend_common_meta_app::schema::TableIdToName; +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_store::MetaStore; +use databend_common_pipeline::core::OutputPort; +use databend_common_pipeline::core::Pipeline; +use databend_common_pipeline::core::ProcessorPtr; +use databend_common_pipeline::sources::AsyncSource; +use databend_common_pipeline::sources::AsyncSourcer; +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::UserApiProvider; +use futures::TryStreamExt; +use futures::stream; +use futures::stream::StreamExt; +use serde_json::json; + +use crate::meta_service_error; +use crate::sessions::TableContext; +use crate::sql::Planner; +use crate::table_functions::object_name::TableNameParser; + +const GET_LINEAGE_FUNC: &str = "get_lineage"; +const GET_LINEAGE_ENGINE: &str = "GET_LINEAGE"; +const DEFAULT_DISTANCE: u8 = 5; +const MAX_DISTANCE: u8 = 5; +const LINEAGE_META_LOOKUP_MAX_CONCURRENCY: usize = 8; + +pub struct GetLineageTable { + table_info: TableInfo, + args: GetLineageArgs, + table_args: TableArgs, +} + +impl GetLineageTable { + pub fn create( + database_name: &str, + table_func_name: &str, + table_id: u64, + 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(), + created_on: DateTime::from_timestamp(0, 0).unwrap(), + updated_on: DateTime::from_timestamp(0, 0).unwrap(), + ..Default::default() + }, + ..Default::default() + }; + + Ok(Arc::new(Self { + table_info, + args, + table_args, + })) + } + + fn schema() -> TableSchemaRef { + TableSchemaRefExt::create(vec![ + TableField::new( + "distance", + TableDataType::Number(databend_common_expression::types::NumberDataType::Int32), + ), + TableField::new("source_object_domain", TableDataType::String), + TableField::new("source_object_name", TableDataType::String), + TableField::new( + "source_column_name", + TableDataType::Nullable(Box::new(TableDataType::String)), + ), + TableField::new("target_object_domain", TableDataType::String), + TableField::new("target_object_name", TableDataType::String), + TableField::new( + "target_column_name", + TableDataType::Nullable(Box::new(TableDataType::String)), + ), + TableField::new("target_status", TableDataType::String), + TableField::new("process", TableDataType::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_backtrace::framed] + 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<()> { + LicenseManagerSwitch::instance() + .check_enterprise_enabled(ctx.get_license_key(), Feature::Lineage)?; + + pipeline.add_source( + |output| { + GetLineageSource::create( + ctx.clone(), + output, + self.args.clone(), + self.table_info.meta.schema.clone(), + ) + }, + 1, + )?; + Ok(()) + } +} + +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, + finished: bool, + args: GetLineageArgs, + schema: DataSchemaRef, +} + +impl GetLineageSource { + fn create( + ctx: Arc, + output: Arc, + args: GetLineageArgs, + schema: TableSchemaRef, + ) -> Result { + let schema = Arc::new(DataSchema::from(schema.as_ref())); + AsyncSourcer::create(ctx.get_scan_progress(), output, Self { + ctx, + finished: false, + args, + schema, + }) + } +} + +#[async_trait::async_trait] +impl AsyncSource for GetLineageSource { + const NAME: &'static str = GET_LINEAGE_FUNC; + + #[async_backtrace::framed] + async fn generate(&mut self) -> Result> { + if self.finished { + return Ok(None); + } + self.finished = true; + + let block = + collect_lineage(self.ctx.clone(), self.args.clone(), self.schema.clone()).await?; + Ok(Some(block)) + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +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, Copy, Debug, PartialEq, Eq)] +enum ObjectDomain { + Table, + View, + Column, + Stage, +} + +impl ObjectDomain { + fn parse(input: &str) -> Result { + match input.trim().to_ascii_uppercase().as_str() { + "TABLE" => Ok(Self::Table), + "VIEW" => Ok(Self::View), + "COLUMN" => Ok(Self::Column), + "STAGE" => Ok(Self::Stage), + other => Err(ErrorCode::BadArguments(format!( + "unsupported object_domain '{other}' for get_lineage" + ))), + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum QueryDirection { + Upstream, + Downstream, +} + +impl QueryDirection { + fn parse(input: &str) -> Result { + match input.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 lineage_direction(self) -> LineageDirection { + match self { + QueryDirection::Upstream => LineageDirection::Upstream, + QueryDirection::Downstream => LineageDirection::Downstream, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +struct FrontierNode { + resolved: ResolvedObject, + lookup_objects: Vec, + column_filter: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +struct ResolvedObject { + object: LineageObjectRef, + domain: String, + name: String, + addressing: ObjectAddressing, + columns: Option, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum ObjectAddressing { + TableId, + NameAddressed, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +struct LineageEdge { + kind: LineageKind, + source_object: LineageObjectRef, + target_object: LineageObjectRef, + source_column: Option, + target_column: Option, + last_query_id: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +struct ColumnMap { + columns: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +struct ColumnFilter { + columns: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +struct ResolvedColumn { + name: String, + id: Option, +} + +#[derive(Clone, Debug)] +struct LineageListRequest { + current: ResolvedObject, + object: LineageObjectRef, + column_filter: Option, +} + +#[derive(Debug)] +struct LineageListResult { + current: ResolvedObject, + column_filter: Option, + entries: Vec<(LineageKey, LineageDetail)>, +} + +#[derive(Debug)] +struct LineageResolveRequest { + current: ResolvedObject, + edge: LineageEdge, +} + +#[derive(Debug)] +struct ResolvedLineageEdge { + kind: LineageKind, + last_query_id: Option, + source_column: Option, + target_column: Option, + source: ResolvedObject, + target: ResolvedObject, +} + +#[derive(Clone)] +struct LineageSpec { + kind: LineageKind, +} + +#[derive(Default)] +struct LineageRows { + distances: Vec, + source_object_domains: Vec, + source_object_names: Vec, + source_column_names: Vec>, + target_object_domains: Vec, + target_object_names: Vec, + target_column_names: Vec>, + target_statuses: Vec, + processes: Vec, +} + +impl LineageRows { + fn push(&mut self, distance: u8, edge: ResolvedLineageEdge) { + self.distances.push(distance as i32); + self.source_object_domains.push(edge.source.domain); + self.source_object_names.push(edge.source.name); + self.source_column_names.push( + edge.source_column + .as_ref() + .map(|column| column.name.clone()), + ); + self.target_object_domains.push(edge.target.domain); + self.target_object_names.push(edge.target.name); + self.target_column_names.push( + edge.target_column + .as_ref() + .map(|column| column.name.clone()), + ); + self.target_statuses.push("ACTIVE".to_string()); + self.processes + .push(process_json(edge.last_query_id.as_deref())); + } + + fn into_block(self) -> DataBlock { + DataBlock::new_from_columns(vec![ + Int32Type::from_data(self.distances), + StringType::from_data(self.source_object_domains), + StringType::from_data(self.source_object_names), + StringType::from_opt_data(self.source_column_names), + StringType::from_data(self.target_object_domains), + StringType::from_data(self.target_object_names), + StringType::from_opt_data(self.target_column_names), + StringType::from_data(self.target_statuses), + StringType::from_data(self.processes), + ]) + } +} + +async fn collect_lineage( + ctx: Arc, + args: GetLineageArgs, + schema: DataSchemaRef, +) -> Result { + let rows = match args.direction { + QueryDirection::Upstream => collect_upstream_lineage(ctx, args).await?, + QueryDirection::Downstream => collect_downstream_lineage(ctx, args).await?, + }; + + if rows.distances.is_empty() { + Ok(DataBlock::empty_with_schema(&schema)) + } else { + Ok(rows.into_block()) + } +} + +async fn collect_upstream_lineage( + ctx: Arc, + args: GetLineageArgs, +) -> Result { + let table_name_parser = Arc::new(TableNameParser::new(&ctx)?); + let start = resolve_start_object(&ctx, table_name_parser.as_ref(), &args).await?; + let meta = UserApiProvider::instance().get_meta_store_client(); + let tenant = ctx.get_tenant(); + + let mut rows = LineageRows::default(); + let mut frontier = vec![start] + .into_iter() + .filter_map(upstream_start_frontier) + .collect::>(); + let mut emitted = HashSet::new(); + + for distance in 1..=args.distance { + let mut next_frontier = Vec::new(); + let list_requests = lineage_list_requests(&frontier); + let list_results = list_lineage_entries_concurrently( + meta.clone(), + tenant.clone(), + args.direction, + list_requests, + ) + .await?; + + let mut resolve_requests = Vec::new(); + for LineageListResult { + current, + column_filter, + entries, + } in list_results + { + for (key, detail) in entries { + resolve_requests.extend( + upstream_edges_from_entry(&key, &detail, column_filter.as_ref()) + .into_iter() + .map(|edge| LineageResolveRequest { + current: current.clone(), + edge, + }), + ); + } + } + + let resolved_edges = resolve_upstream_lineage_edges_concurrently( + ctx.clone(), + table_name_parser.clone(), + resolve_requests, + ) + .await?; + + for resolved_edge in resolved_edges { + if !emitted.insert(resolved_edge_dedup_key( + distance, + &resolved_edge.source, + &resolved_edge.target, + resolved_edge.source_column.as_ref(), + resolved_edge.target_column.as_ref(), + )) { + continue; + } + + let next = next_upstream_frontier_for_edge(&resolved_edge); + rows.push(distance, resolved_edge); + if let Some(next) = next { + next_frontier.push(next); + } + } + frontier = next_frontier; + if frontier.is_empty() { + break; + } + } + + Ok(rows) +} + +async fn collect_downstream_lineage( + ctx: Arc, + args: GetLineageArgs, +) -> Result { + let table_name_parser = Arc::new(TableNameParser::new(&ctx)?); + let start = resolve_start_object(&ctx, table_name_parser.as_ref(), &args).await?; + let meta = UserApiProvider::instance().get_meta_store_client(); + let tenant = ctx.get_tenant(); + + let mut rows = LineageRows::default(); + let mut frontier = vec![downstream_start_frontier(start)]; + let mut emitted = HashSet::new(); + + for distance in 1..=args.distance { + let mut next_frontier = Vec::new(); + let list_requests = lineage_list_requests(&frontier); + let list_results = list_lineage_entries_concurrently( + meta.clone(), + tenant.clone(), + QueryDirection::Downstream, + list_requests, + ) + .await?; + + let mut resolve_requests = Vec::new(); + for LineageListResult { + current, + column_filter, + entries, + } in list_results + { + for (key, detail) in entries { + resolve_requests.extend( + downstream_edges_from_entry(&key, &detail, column_filter.as_ref()) + .into_iter() + .map(|edge| LineageResolveRequest { + current: current.clone(), + edge, + }), + ); + } + } + + let resolved_edges = resolve_downstream_lineage_edges_concurrently( + ctx.clone(), + table_name_parser.clone(), + resolve_requests, + ) + .await?; + + for resolved_edge in resolved_edges { + if !emitted.insert(resolved_edge_dedup_key( + distance, + &resolved_edge.source, + &resolved_edge.target, + resolved_edge.source_column.as_ref(), + resolved_edge.target_column.as_ref(), + )) { + continue; + } + + let next = next_downstream_frontier_for_edge(&resolved_edge); + rows.push(distance, resolved_edge); + if let Some(next) = next { + next_frontier.push(next); + } + } + frontier = next_frontier; + if frontier.is_empty() { + break; + } + } + + Ok(rows) +} + +fn upstream_start_frontier(mut start: FrontierNode) -> Option { + (start.resolved.object.object_type == LineageObjectType::Table + && start.resolved.addressing == ObjectAddressing::TableId) + .then(|| { + start.lookup_objects = upstream_lookup_objects(&start.resolved); + start + }) +} + +fn downstream_start_frontier(mut start: FrontierNode) -> FrontierNode { + start.lookup_objects = downstream_lookup_objects(&start.resolved); + start +} + +fn lineage_list_requests(frontier: &[FrontierNode]) -> Vec { + frontier + .iter() + .flat_map(|node| { + node.lookup_objects + .iter() + .cloned() + .map(|object| LineageListRequest { + current: node.resolved.clone(), + object, + column_filter: node.column_filter.clone(), + }) + }) + .collect() +} + +async fn list_lineage_entries_concurrently( + meta: Arc, + tenant: databend_common_meta_app::tenant::Tenant, + direction: QueryDirection, + requests: Vec, +) -> Result> { + stream::iter(requests.into_iter().map(|request| { + let meta = meta.clone(); + let tenant = tenant.clone(); + async move { + let entries = meta + .as_ref() + .list_lineage(ListLineageReq { + tenant, + direction: direction.lineage_direction(), + object: request.object, + }) + .await + .map_err(meta_service_error)? + .entries; + Ok(LineageListResult { + current: request.current, + column_filter: request.column_filter, + entries, + }) + } + })) + .buffer_unordered(LINEAGE_META_LOOKUP_MAX_CONCURRENCY) + .try_collect() + .await +} + +async fn resolve_downstream_lineage_edges_concurrently( + ctx: Arc, + table_name_parser: Arc, + requests: Vec, +) -> Result> { + let resolved: Vec> = + stream::iter(requests.into_iter().map(|request| { + let ctx = ctx.clone(); + let table_name_parser = table_name_parser.clone(); + async move { resolve_downstream_lineage_edge(ctx, table_name_parser, request).await } + })) + .buffer_unordered(LINEAGE_META_LOOKUP_MAX_CONCURRENCY) + .try_collect() + .await?; + + Ok(resolved.into_iter().flatten().collect()) +} + +async fn resolve_upstream_lineage_edges_concurrently( + ctx: Arc, + table_name_parser: Arc, + requests: Vec, +) -> Result> { + let resolved: Vec> = + stream::iter(requests.into_iter().map(|request| { + let ctx = ctx.clone(); + let table_name_parser = table_name_parser.clone(); + async move { resolve_upstream_lineage_edge(ctx, table_name_parser, request).await } + })) + .buffer_unordered(LINEAGE_META_LOOKUP_MAX_CONCURRENCY) + .try_collect() + .await?; + + Ok(resolved.into_iter().flatten().collect()) +} + +async fn resolve_downstream_lineage_edge( + ctx: Arc, + table_name_parser: Arc, + request: LineageResolveRequest, +) -> Result> { + let LineageResolveRequest { current, edge } = request; + let Some(target) = + resolve_lineage_object(&ctx, table_name_parser.as_ref(), &edge.target_object).await? + else { + return Ok(None); + }; + Ok(resolved_downstream_edge_from_current(current, edge, target)) +} + +async fn resolve_upstream_lineage_edge( + ctx: Arc, + table_name_parser: Arc, + request: LineageResolveRequest, +) -> Result> { + let LineageResolveRequest { current, edge } = request; + let Some(source) = + resolve_lineage_object(&ctx, table_name_parser.as_ref(), &edge.source_object).await? + else { + return Ok(None); + }; + Ok(resolved_upstream_edge_from_current(current, edge, source)) +} + +fn resolved_downstream_edge_from_current( + current: ResolvedObject, + edge: LineageEdge, + target: ResolvedObject, +) -> Option { + let source_column = resolve_column_ref(¤t, edge.source_column.as_ref())?; + let target_column = resolve_column_ref(&target, edge.target_column.as_ref())?; + Some(ResolvedLineageEdge { + kind: edge.kind, + last_query_id: edge.last_query_id, + source_column, + target_column, + source: current, + target, + }) +} + +fn resolved_upstream_edge_from_current( + current: ResolvedObject, + edge: LineageEdge, + source: ResolvedObject, +) -> Option { + let source_column = resolve_column_ref(&source, edge.source_column.as_ref())?; + let target_column = resolve_column_ref(¤t, edge.target_column.as_ref())?; + Some(ResolvedLineageEdge { + kind: edge.kind, + last_query_id: edge.last_query_id, + source_column, + target_column, + source, + target: current, + }) +} + +fn column_map_from_schema(schema: &TableSchemaRef) -> ColumnMap { + ColumnMap { + columns: schema + .fields() + .iter() + .map(|field| ResolvedColumn { + name: field.name().to_string(), + id: Some(field.column_id() as u64), + }) + .collect(), + } +} + +fn resolve_column_filter(object: &ResolvedObject, column: ColumnRef) -> Result { + let Some(Some(resolved_column)) = resolve_column_ref(object, Some(&column)) else { + return Err(ErrorCode::BadArguments(format!( + "column '{}' does not exist in object '{}'", + column_ref_display(&column), + object.name + ))); + }; + Ok(ColumnFilter { + columns: vec![resolved_column], + }) +} + +fn resolve_column_ref( + object: &ResolvedObject, + column: Option<&ColumnRef>, +) -> Option> { + let Some(column) = column else { + return Some(None); + }; + let Some(column_map) = &object.columns else { + return Some(Some(resolved_column_from_ref(column))); + }; + column_map + .columns + .iter() + .find(|candidate| column_ref_matches_resolved(column, candidate)) + .cloned() + .map(Some) +} + +fn column_matches_filter(column: &ColumnRef, filter: &ColumnFilter) -> bool { + filter + .columns + .iter() + .any(|candidate| column_ref_matches_resolved(column, candidate)) +} + +fn column_ref_matches_resolved(column: &ColumnRef, resolved: &ResolvedColumn) -> bool { + match column { + ColumnRef::Id(id) => resolved.id == Some(*id), + ColumnRef::Name(name) => resolved.name.eq_ignore_ascii_case(name), + } +} + +fn resolved_column_from_ref(column: &ColumnRef) -> ResolvedColumn { + match column { + ColumnRef::Id(id) => ResolvedColumn { + name: id.to_string(), + id: Some(*id), + }, + ColumnRef::Name(name) => ResolvedColumn { + name: name.clone(), + id: None, + }, + } +} + +fn column_ref_display(column: &ColumnRef) -> String { + match column { + ColumnRef::Id(id) => id.to_string(), + ColumnRef::Name(name) => name.clone(), + } +} + +fn downstream_edges_from_entry( + key: &LineageKey, + detail: &LineageDetail, + column_filter: Option<&ColumnFilter>, +) -> Vec { + lineage_edges_from_entry( + QueryDirection::Downstream, + key.object.clone(), + key.related_object.clone(), + detail, + column_filter, + ) +} + +fn upstream_edges_from_entry( + key: &LineageKey, + detail: &LineageDetail, + column_filter: Option<&ColumnFilter>, +) -> Vec { + lineage_edges_from_entry( + QueryDirection::Upstream, + key.related_object.clone(), + key.object.clone(), + detail, + column_filter, + ) +} + +fn lineage_edges_from_entry( + direction: QueryDirection, + source_object: LineageObjectRef, + target_object: LineageObjectRef, + detail: &LineageDetail, + column_filter: Option<&ColumnFilter>, +) -> Vec { + let Some(spec) = LineageSpec::new(&detail.kind) else { + return Vec::new(); + }; + if let Some(column_name) = column_filter { + return detail + .column_lineage + .iter() + .filter_map(|column| spec.column_edge(direction, column, column_name)) + .map(|(source_column_name, target_column_name)| LineageEdge { + kind: detail.kind.clone(), + source_object: source_object.clone(), + target_object: target_object.clone(), + source_column: Some(source_column_name), + target_column: Some(target_column_name), + last_query_id: detail.last_query_id.clone(), + }) + .collect(); + } + + vec![LineageEdge { + kind: detail.kind.clone(), + source_object, + target_object, + source_column: None, + target_column: None, + last_query_id: detail.last_query_id.clone(), + }] +} + +impl LineageSpec { + fn new(kind: &LineageKind) -> Option { + match kind { + LineageKind::Unknown(_) => None, + _ => Some(Self { kind: kind.clone() }), + } + } + + fn column_edge( + &self, + direction: QueryDirection, + column: &LineageColumn, + column_filter: &ColumnFilter, + ) -> Option<(ColumnRef, ColumnRef)> { + match direction { + QueryDirection::Downstream => column_matches_filter(&column.upstream, column_filter) + .then(|| (column.upstream.clone(), column.downstream.clone())), + QueryDirection::Upstream => column_matches_filter(&column.downstream, column_filter) + .then(|| (column.upstream.clone(), column.downstream.clone())), + } + } + + fn upstream_next_lookup_objects(&self, resolved: &ResolvedObject) -> Vec { + match &self.kind { + LineageKind::View | LineageKind::MaterializedView => { + table_like_id_and_name_lookup_objects(resolved) + } + LineageKind::Ctas | LineageKind::DataMovement | LineageKind::Unknown(_) => { + canonical_lookup_object(resolved) + } + } + } + + fn downstream_next_lookup_objects(&self, resolved: &ResolvedObject) -> Vec { + table_like_id_and_name_lookup_objects(resolved) + } +} + +async fn resolve_start_object( + ctx: &Arc, + table_name_parser: &TableNameParser, + args: &GetLineageArgs, +) -> Result { + match args.object_domain { + ObjectDomain::Table | ObjectDomain::View => { + let table = resolve_table_like( + ctx, + table_name_parser, + &args.object_name, + args.object_domain, + ) + .await?; + Ok(FrontierNode { + lookup_objects: start_lookup_objects(args.direction, &table), + resolved: table, + column_filter: None, + }) + } + ObjectDomain::Column => { + let (table_name, column_name) = split_column_name(&args.object_name)?; + let table = + resolve_table_like(ctx, table_name_parser, &table_name, ObjectDomain::Column) + .await?; + let column_filter = resolve_column_filter( + &table, + ColumnRef::Name(table_name_parser.normalize_column_identifier(&column_name)), + )?; + Ok(FrontierNode { + lookup_objects: start_lookup_objects(args.direction, &table), + resolved: table, + column_filter: Some(column_filter), + }) + } + ObjectDomain::Stage => { + let stage = resolve_stage(ctx, &args.object_name).await?; + Ok(FrontierNode { + lookup_objects: start_lookup_objects(args.direction, &stage), + resolved: stage, + column_filter: None, + }) + } + } +} + +async fn resolve_table_like( + ctx: &Arc, + table_name_parser: &TableNameParser, + object_name: &str, + expected_domain: ObjectDomain, +) -> Result { + let (catalog_name, db_name, table_name) = table_name_parser.parse_table_name(object_name)?; + let table = ctx.get_table(&catalog_name, &db_name, &table_name).await?; + let table_id = table.get_table_info().ident.table_id; + let engine = table.engine(); + let columns = Some(column_map_for_table(ctx, &table).await); + match expected_domain { + ObjectDomain::View if engine != VIEW_ENGINE => { + return Err(ErrorCode::BadArguments(format!( + "object '{}' is not a VIEW", + object_name + ))); + } + ObjectDomain::Table if matches!(engine, VIEW_ENGINE | STREAM_ENGINE) => { + return Err(ErrorCode::BadArguments(format!( + "object '{}' is not a TABLE", + object_name + ))); + } + _ => {} + } + + Ok(ResolvedObject { + object: LineageObjectRef { + object_type: LineageObjectType::Table, + identity: LineageIdentity::Id { + id: table_id.to_string(), + }, + }, + domain: table_domain(engine).to_string(), + name: full_table_name(&catalog_name, &db_name, &table_name), + addressing: ObjectAddressing::TableId, + columns, + }) +} + +async fn resolve_stage(ctx: &Arc, object_name: &str) -> Result { + let stage_name = object_name.trim().trim_start_matches('@'); + if stage_name.is_empty() { + return Err(ErrorCode::BadArguments("stage name must not be empty")); + } + + let tenant = ctx.get_tenant(); + let stage_info = UserApiProvider::instance() + .get_stage(&tenant, stage_name) + .await?; + + Ok(ResolvedObject { + object: LineageObjectRef { + object_type: LineageObjectType::Stage, + identity: LineageIdentity::Name { + name: stage_info.stage_name.clone(), + }, + }, + domain: "STAGE".to_string(), + name: stage_info.stage_name, + addressing: ObjectAddressing::NameAddressed, + columns: None, + }) +} + +async fn resolve_lineage_object( + ctx: &Arc, + table_name_parser: &TableNameParser, + object: &LineageObjectRef, +) -> Result> { + match (&object.object_type, &object.identity) { + (LineageObjectType::Table, LineageIdentity::Id { id }) => { + let Ok(table_id) = id.parse::() else { + return Ok(None); + }; + // Stable table ids belong to the meta-backed catalog; external + // catalogs such as Iceberg and Hive use name-addressed lineage. + let catalog_name = ctx.get_current_catalog(); + let catalog = ctx.get_catalog(&catalog_name).await?; + 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 db_name = match catalog.get_db_name_by_id(name_entry.data.db_id).await { + Ok(db_name) => db_name, + Err(error) if error.code() == ErrorCode::UNKNOWN_DATABASE_ID => return Ok(None), + Err(error) => return Err(error), + }; + let table_name = name_entry.data.table_name; + let Ok(table) = ctx.get_table(&catalog_name, &db_name, &table_name).await else { + return Ok(None); + }; + let actual_table_id = table.get_table_info().ident.table_id; + if actual_table_id != table_id { + log::warn!( + "skip lineage object '{}.{}.{}': expected table id {}, got {}", + catalog_name, + db_name, + table_name, + table_id, + actual_table_id + ); + return Ok(None); + } + let full_name = full_table_name(&catalog_name, &db_name, &table_name); + Ok(Some(ResolvedObject { + object: object.clone(), + domain: table_domain(table.engine()).to_string(), + name: full_name, + addressing: ObjectAddressing::TableId, + columns: Some(column_map_for_table(ctx, &table).await), + })) + } + (LineageObjectType::Table, LineageIdentity::Name { name }) => { + let Ok((catalog_name, db_name, table_name)) = table_name_parser.parse_table_name(name) + else { + return Ok(None); + }; + let table = ctx + .get_table(&catalog_name, &db_name, &table_name) + .await + .ok(); + let Some(table) = table else { + return Ok(None); + }; + let table_id = table.get_table_info().ident.table_id; + Ok(Some(ResolvedObject { + object: LineageObjectRef { + object_type: LineageObjectType::Table, + identity: LineageIdentity::Id { + id: table_id.to_string(), + }, + }, + domain: table_domain(table.engine()).to_string(), + name: full_table_name(&catalog_name, &db_name, &table_name), + addressing: ObjectAddressing::TableId, + columns: Some(column_map_for_table(ctx, &table).await), + })) + } + (LineageObjectType::Stage, LineageIdentity::Name { name }) => Ok(Some(ResolvedObject { + object: object.clone(), + domain: "STAGE".to_string(), + name: name.clone(), + addressing: ObjectAddressing::NameAddressed, + columns: None, + })), + _ => Ok(None), + } +} + +async fn column_map_for_table(ctx: &Arc, table: &Arc) -> ColumnMap { + if table.engine() != VIEW_ENGINE { + return column_map_from_schema(&table.schema()); + } + + let Some(query) = table.options().get(QUERY) else { + return ColumnMap { columns: vec![] }; + }; + let mut planner = Planner::new(ctx.clone()); + match planner.plan_sql(query).await { + Ok((plan, _)) => match infer_table_schema(&plan.schema()) { + Ok(schema) => column_map_from_schema(&schema), + Err(error) => { + log::warn!( + "failed to infer view schema for lineage object '{}': {error}", + table.name() + ); + ColumnMap { columns: vec![] } + } + }, + Err(error) => { + log::warn!( + "failed to plan view schema for lineage object '{}': {error}", + table.name() + ); + ColumnMap { columns: vec![] } + } + } +} + +fn next_upstream_frontier_for_edge(edge: &ResolvedLineageEdge) -> Option { + if edge.source.object.object_type != LineageObjectType::Table + || edge.source.addressing != ObjectAddressing::TableId + { + return None; + } + let spec = LineageSpec::new(&edge.kind).expect("resolved lineage edge must have known kind"); + Some(FrontierNode { + lookup_objects: spec.upstream_next_lookup_objects(&edge.source), + resolved: edge.source.clone(), + column_filter: edge.source_column.clone().map(|column| ColumnFilter { + columns: vec![column], + }), + }) +} + +fn next_downstream_frontier_for_edge(edge: &ResolvedLineageEdge) -> Option { + if edge.target.object.object_type != LineageObjectType::Table + || edge.target.addressing != ObjectAddressing::TableId + { + return None; + } + let spec = LineageSpec::new(&edge.kind).expect("resolved lineage edge must have known kind"); + Some(FrontierNode { + lookup_objects: spec.downstream_next_lookup_objects(&edge.target), + resolved: edge.target.clone(), + column_filter: edge.target_column.clone().map(|column| ColumnFilter { + columns: vec![column], + }), + }) +} + +fn start_lookup_objects( + direction: QueryDirection, + resolved: &ResolvedObject, +) -> Vec { + match direction { + QueryDirection::Upstream => upstream_lookup_objects(resolved), + QueryDirection::Downstream => downstream_lookup_objects(resolved), + } +} + +fn upstream_lookup_objects(resolved: &ResolvedObject) -> Vec { + canonical_lookup_object(resolved) +} + +fn downstream_lookup_objects(resolved: &ResolvedObject) -> Vec { + table_like_id_and_name_lookup_objects(resolved) +} + +fn table_like_id_and_name_lookup_objects(resolved: &ResolvedObject) -> Vec { + if resolved.object.object_type == LineageObjectType::Table + && resolved.addressing == ObjectAddressing::TableId + { + vec![resolved.object.clone(), LineageObjectRef { + object_type: LineageObjectType::Table, + identity: LineageIdentity::Name { + name: resolved.name.clone(), + }, + }] + } else { + vec![resolved.object.clone()] + } +} + +fn canonical_lookup_object(resolved: &ResolvedObject) -> Vec { + vec![resolved.object.clone()] +} + +fn resolved_edge_dedup_key( + distance: u8, + source: &ResolvedObject, + target: &ResolvedObject, + source_column: Option<&ResolvedColumn>, + target_column: Option<&ResolvedColumn>, +) -> String { + format!( + "{distance}/{}/{}/{}/{}/{}/{}", + object_type_key(&source.object.object_type), + identity_key(&source.object.identity), + column_key(source_column), + object_type_key(&target.object.object_type), + identity_key(&target.object.identity), + column_key(target_column) + ) +} + +fn column_key(column: Option<&ResolvedColumn>) -> String { + match column { + Some(ResolvedColumn { id: Some(id), .. }) => format!("id/{id}"), + Some(ResolvedColumn { name, .. }) => format!("name/{name}"), + None => String::new(), + } +} + +fn split_column_name(input: &str) -> Result<(String, String)> { + let input = input.trim(); + if input.is_empty() { + return Err(ErrorCode::BadArguments( + "column object_name must not be empty", + )); + } + + let Some(index) = last_unquoted_dot(input) else { + return Err(ErrorCode::BadArguments( + "column object_name must be qualified by table name".to_string(), + )); + }; + 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".to_string(), + )); + } + Ok((table.to_string(), column.to_string())) +} + +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 +} + +fn table_domain(engine: &str) -> &'static str { + match engine { + VIEW_ENGINE => "VIEW", + STREAM_ENGINE => "STREAM", + _ => "TABLE", + } +} + +fn full_table_name(catalog: &str, database: &str, table: &str) -> String { + if catalog == "default" { + format!("{database}.{table}") + } else { + format!("{catalog}.{database}.{table}") + } +} + +fn process_json(last_query_id: Option<&str>) -> String { + match last_query_id { + Some(last_query_id) => json!({ "last_query_id": last_query_id }).to_string(), + None => "{}".to_string(), + } +} + +fn object_type_key(object_type: &LineageObjectType) -> &'static str { + match object_type { + LineageObjectType::Table => "table", + LineageObjectType::Stage => "stage", + } +} + +fn identity_key(identity: &LineageIdentity) -> String { + match identity { + LineageIdentity::Id { id } => format!("id/{id}"), + LineageIdentity::Name { name } => format!("name/{name}"), + } +} + +#[cfg(test)] +mod tests { + use chrono::TimeZone; + use chrono::Utc; + use databend_common_catalog::table_args::TableArgs; + use databend_common_expression::Scalar; + + use super::*; + + #[test] + fn test_parse_args() -> Result<()> { + let args = TableArgs::new_positioned(vec![ + Scalar::String("db.t".to_string()), + Scalar::String("table".to_string()), + Scalar::String("downSTREAM".to_string()), + ]); + let parsed = GetLineageArgs::parse(&args)?; + assert_eq!(parsed.object_name, "db.t"); + assert_eq!(parsed.object_domain, ObjectDomain::Table); + assert_eq!(parsed.direction, QueryDirection::Downstream); + assert_eq!(parsed.distance, DEFAULT_DISTANCE); + + let args = TableArgs::new_positioned(vec![ + Scalar::String("db.t".to_string()), + Scalar::String("column".to_string()), + Scalar::String("UPSTREAM".to_string()), + Scalar::Number(databend_common_expression::types::NumberScalar::Int64(2)), + ]); + assert_eq!(GetLineageArgs::parse(&args)?.distance, 2); + + let args = TableArgs::new_positioned(vec![ + Scalar::String("db.t".to_string()), + Scalar::String("stream".to_string()), + Scalar::String("UPSTREAM".to_string()), + ]); + assert!(GetLineageArgs::parse(&args).is_err()); + + for distance in [0, MAX_DISTANCE + 1] { + let args = TableArgs::new_positioned(vec![ + Scalar::String("db.t".to_string()), + Scalar::String("table".to_string()), + Scalar::String("UPSTREAM".to_string()), + Scalar::Number(databend_common_expression::types::NumberScalar::Int64( + i64::from(distance), + )), + ]); + assert!(GetLineageArgs::parse(&args).is_err()); + } + + let invalid_direction = TableArgs::new_positioned(vec![ + Scalar::String("db.t".to_string()), + Scalar::String("table".to_string()), + Scalar::String("SIDEWAYS".to_string()), + ]); + assert!(GetLineageArgs::parse(&invalid_direction).is_err()); + + let too_few = TableArgs::new_positioned(vec![ + Scalar::String("db.t".to_string()), + Scalar::String("table".to_string()), + ]); + assert!(GetLineageArgs::parse(&too_few).is_err()); + + Ok(()) + } + + #[test] + fn test_split_column_name() -> Result<()> { + assert_eq!( + split_column_name("db.t.c")?, + ("db.t".to_string(), "c".to_string()) + ); + assert_eq!( + split_column_name(r#""db.name"."t.name"."c.name""#)?, + ( + r#""db.name"."t.name""#.to_string(), + r#""c.name""#.to_string() + ) + ); + assert!(split_column_name("c").is_err()); + assert!(split_column_name("t.").is_err()); + Ok(()) + } + + #[test] + fn test_directional_edges_from_entry_table_direction() { + let key = lineage_key(LineageDirection::Downstream); + let detail = LineageDetail { + kind: LineageKind::DataMovement, + last_query_id: Some("q1".to_string()), + updated_on: lineage_time(), + column_lineage: vec![], + }; + + let downstream = downstream_edges_from_entry(&key, &detail, None); + assert_eq!(identity_key(&downstream[0].source_object.identity), "id/1"); + assert_eq!(identity_key(&downstream[0].target_object.identity), "id/2"); + + let upstream = upstream_edges_from_entry(&key, &detail, None); + assert_eq!(identity_key(&upstream[0].source_object.identity), "id/2"); + assert_eq!(identity_key(&upstream[0].target_object.identity), "id/1"); + } + + #[test] + fn test_directional_edges_from_entry_column_filter() { + let key = lineage_key(LineageDirection::Downstream); + let detail = LineageDetail { + kind: LineageKind::DataMovement, + last_query_id: None, + updated_on: lineage_time(), + column_lineage: vec![ + LineageColumn { + upstream: ColumnRef::Id(1), + downstream: ColumnRef::Id(10), + }, + LineageColumn { + upstream: ColumnRef::Id(2), + downstream: ColumnRef::Id(10), + }, + ], + }; + + let downstream = downstream_edges_from_entry( + &key, + &detail, + Some(&ColumnFilter { + columns: vec![resolved_column("a", Some(1))], + }), + ); + assert_eq!(downstream.len(), 1); + assert_eq!( + downstream[0] + .source_column + .as_ref() + .map(column_ref_display) + .as_deref(), + Some("1") + ); + assert_eq!( + downstream[0] + .target_column + .as_ref() + .map(column_ref_display) + .as_deref(), + Some("10") + ); + + let upstream = upstream_edges_from_entry( + &key, + &detail, + Some(&ColumnFilter { + columns: vec![resolved_column("x", Some(10))], + }), + ); + assert_eq!(upstream.len(), 2); + } + + #[test] + fn test_table_lookup_objects_by_direction() { + let resolved = resolved_table("1", "db.t"); + + let downstream = downstream_lookup_objects(&resolved); + assert_eq!(downstream.len(), 2); + assert_eq!(identity_key(&downstream[0].identity), "id/1"); + assert_eq!(identity_key(&downstream[1].identity), "name/db.t"); + + let upstream = upstream_lookup_objects(&resolved); + assert_eq!(upstream.len(), 1); + assert_eq!(identity_key(&upstream[0].identity), "id/1"); + + let start_upstream = start_lookup_objects(QueryDirection::Upstream, &resolved); + assert_eq!(start_upstream, upstream); + + let start_downstream = start_lookup_objects(QueryDirection::Downstream, &resolved); + assert_eq!(start_downstream, downstream); + } + + #[test] + fn test_next_frontier_keeps_resolved_current_object() { + let source = resolved_table("1", "db.source"); + let target = resolved_table("2", "db.target"); + let edge = ResolvedLineageEdge { + kind: LineageKind::DataMovement, + last_query_id: None, + source_column: Some(resolved_column("a", Some(1))), + target_column: Some(resolved_column("b", Some(2))), + source: source.clone(), + target: target.clone(), + }; + + let upstream = next_upstream_frontier_for_edge(&edge).unwrap(); + assert_eq!(upstream.resolved, source); + assert_eq!(upstream.lookup_objects.len(), 1); + assert_eq!(identity_key(&upstream.lookup_objects[0].identity), "id/1"); + assert_eq!( + upstream + .column_filter + .as_ref() + .and_then(|filter| filter.columns.first()) + .map(|column| column.name.as_str()), + Some("a") + ); + + let downstream = next_downstream_frontier_for_edge(&edge).unwrap(); + assert_eq!(downstream.resolved, target); + assert_eq!(downstream.lookup_objects.len(), 2); + assert_eq!(identity_key(&downstream.lookup_objects[0].identity), "id/2"); + assert_eq!( + identity_key(&downstream.lookup_objects[1].identity), + "name/db.target" + ); + assert_eq!( + downstream + .column_filter + .as_ref() + .and_then(|filter| filter.columns.first()) + .map(|column| column.name.as_str()), + Some("b") + ); + } + + #[test] + fn test_upstream_view_kind_expands_next_lookup_to_id_and_name() { + let source = resolved_table("1", "db.source"); + let target = resolved_table("2", "db.target"); + let edge = ResolvedLineageEdge { + kind: LineageKind::View, + last_query_id: None, + source_column: None, + target_column: None, + source: source.clone(), + target, + }; + + let upstream = next_upstream_frontier_for_edge(&edge).unwrap(); + assert_eq!(upstream.lookup_objects.len(), 2); + assert_eq!(identity_key(&upstream.lookup_objects[0].identity), "id/1"); + assert_eq!( + identity_key(&upstream.lookup_objects[1].identity), + "name/db.source" + ); + } + + #[test] + fn test_upstream_stage_source_does_not_continue() { + let source = ResolvedObject { + object: LineageObjectRef { + object_type: LineageObjectType::Stage, + identity: LineageIdentity::Name { + name: "s".to_string(), + }, + }, + domain: "STAGE".to_string(), + name: "s".to_string(), + addressing: ObjectAddressing::NameAddressed, + columns: None, + }; + let target = resolved_table("2", "db.target"); + let edge = ResolvedLineageEdge { + kind: LineageKind::DataMovement, + last_query_id: None, + source_column: None, + target_column: None, + source, + target, + }; + + assert!(next_upstream_frontier_for_edge(&edge).is_none()); + } + + fn resolved_table(id: &str, name: &str) -> ResolvedObject { + ResolvedObject { + object: LineageObjectRef { + object_type: LineageObjectType::Table, + identity: LineageIdentity::Id { id: id.to_string() }, + }, + domain: "TABLE".to_string(), + name: name.to_string(), + addressing: ObjectAddressing::TableId, + columns: Some(ColumnMap { + columns: vec![ + resolved_column("a", Some(1)), + resolved_column("b", Some(2)), + resolved_column("c", Some(3)), + resolved_column("x", Some(10)), + ], + }), + } + } + + fn resolved_column(name: &str, id: Option) -> ResolvedColumn { + ResolvedColumn { + name: name.to_string(), + id, + } + } + + fn lineage_key(direction: LineageDirection) -> LineageKey { + LineageKey { + direction, + object: LineageObjectRef { + object_type: LineageObjectType::Table, + identity: LineageIdentity::Id { + id: "1".to_string(), + }, + }, + related_object: LineageObjectRef { + object_type: LineageObjectType::Table, + identity: LineageIdentity::Id { + id: "2".to_string(), + }, + }, + } + } + + fn lineage_time() -> chrono::DateTime { + Utc.with_ymd_and_hms(2026, 7, 23, 0, 0, 0).unwrap() + } +} diff --git a/src/query/service/src/table_functions/mod.rs b/src/query/service/src/table_functions/mod.rs index 0f490612c5453..654a44c1c67b0 100644 --- a/src/query/service/src/table_functions/mod.rs +++ b/src/query/service/src/table_functions/mod.rs @@ -16,11 +16,13 @@ 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; mod list_stage; mod numbers; +pub(crate) mod object_name; mod others; mod policy_references; #[cfg(feature = "task-support")] @@ -42,6 +44,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/object_name.rs b/src/query/service/src/table_functions/object_name.rs new file mode 100644 index 0000000000000..952da14c3d086 --- /dev/null +++ b/src/query/service/src/table_functions/object_name.rs @@ -0,0 +1,84 @@ +// 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_ast::parser::Dialect; +use databend_common_ast::parser::parse_table_ref; +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; +use databend_common_sql::planner::NameResolutionContext; +use databend_common_sql::planner::normalize_identifier; + +use crate::sessions::TableContext; + +pub(crate) struct TableNameParser { + current_catalog: String, + current_database: String, + dialect: Dialect, + name_resolution_ctx: NameResolutionContext, +} + +impl TableNameParser { + pub(crate) fn new(ctx: &Arc) -> Result { + let settings = ctx.get_settings(); + let name_resolution_ctx = NameResolutionContext::try_from(settings.as_ref())?; + let dialect = settings.get_sql_dialect().unwrap_or_default(); + + Ok(Self { + current_catalog: ctx.get_current_catalog(), + current_database: ctx.get_current_database(), + dialect, + name_resolution_ctx, + }) + } + + /// Parse table name in format "table", "db.table", or "catalog.db.table". + /// Correctly handles quoted identifiers and normalizes them according to session settings. + pub(crate) fn parse_table_name(&self, name: &str) -> Result<(String, String, String)> { + let trimmed = name.trim(); + if trimmed.is_empty() { + return Err(ErrorCode::BadArguments("object_name must not be empty")); + } + + let table_ref = parse_table_ref(trimmed, self.dialect).map_err(|e| { + ErrorCode::BadArguments(format!("Invalid table name '{}': {}", name, e.1)) + })?; + + let catalog = table_ref + .catalog + .map(|i| normalize_identifier(&i, &self.name_resolution_ctx).name) + .unwrap_or_else(|| self.current_catalog.clone()); + let database = table_ref + .database + .map(|i| normalize_identifier(&i, &self.name_resolution_ctx).name) + .unwrap_or_else(|| self.current_database.clone()); + let table = normalize_identifier(&table_ref.table, &self.name_resolution_ctx).name; + + Ok((catalog, database, table)) + } + + pub(crate) fn normalize_column_identifier(&self, name: &str) -> String { + let trimmed = name.trim(); + if trimmed.starts_with('"') && trimmed.ends_with('"') && trimmed.len() >= 2 { + return trimmed[1..trimmed.len() - 1].replace("\"\"", "\""); + } + + if self.name_resolution_ctx.unquoted_ident_case_sensitive { + trimmed.to_string() + } else { + trimmed.to_lowercase() + } + } +} 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..3b6c12780e9ba 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; @@ -449,6 +450,10 @@ impl TableFunctionFactory { "copy_history".to_string(), (next_id(), Arc::new(CopyHistoryTable::create)), ); + creators.insert( + "get_lineage".to_string(), + (next_id(), Arc::new(GetLineageTable::create)), + ); creators.insert( "stream_backlog".to_string(), (next_id(), Arc::new(StreamBacklogTable::create)), diff --git a/src/query/service/src/table_functions/tag_references/tag_references_table.rs b/src/query/service/src/table_functions/tag_references/tag_references_table.rs index 8d3d156fb35e1..7392d2e7a72a1 100644 --- a/src/query/service/src/table_functions/tag_references/tag_references_table.rs +++ b/src/query/service/src/table_functions/tag_references/tag_references_table.rs @@ -17,7 +17,6 @@ use std::sync::Arc; use databend_common_ast::parser::parse_database_ref; use databend_common_ast::parser::parse_procedure_ref; -use databend_common_ast::parser::parse_table_ref; use databend_common_ast::parser::parse_udf_ref; use databend_common_catalog::plan::DataSourcePlan; use databend_common_catalog::plan::PartStatistics; @@ -67,6 +66,7 @@ use databend_common_users::UserApiProvider; use crate::meta_service_error; use crate::sessions::TableContext; +use crate::table_functions::object_name::TableNameParser; const TAG_REFERENCES_FUNC: &str = "tag_references"; const TAG_REFERENCES_ENGINE: &str = "TAG_REFERENCES"; @@ -258,7 +258,9 @@ async fn collect_tag_references( ) } "TABLE" => { - let (catalog_name, db_name, table_name) = parse_table_name(&ctx, &object_name)?; + let table_name_parser = TableNameParser::new(&ctx)?; + let (catalog_name, db_name, table_name) = + table_name_parser.parse_table_name(&object_name)?; let catalog = ctx.get_catalog(&catalog_name).await?; let db = catalog.get_database(&tenant, &db_name).await?; let db_id = db.get_db_info().database_id.db_id; @@ -411,7 +413,9 @@ async fn collect_tag_references( ) } "VIEW" => { - let (catalog_name, db_name, view_name) = parse_table_name(&ctx, &object_name)?; + let table_name_parser = TableNameParser::new(&ctx)?; + let (catalog_name, db_name, view_name) = + table_name_parser.parse_table_name(&object_name)?; let catalog = ctx.get_catalog(&catalog_name).await?; let db = catalog.get_database(&tenant, &db_name).await?; let db_id = db.get_db_info().database_id.db_id; @@ -445,7 +449,9 @@ async fn collect_tag_references( ) } "STREAM" => { - let (catalog_name, db_name, stream_name) = parse_table_name(&ctx, &object_name)?; + let table_name_parser = TableNameParser::new(&ctx)?; + let (catalog_name, db_name, stream_name) = + table_name_parser.parse_table_name(&object_name)?; let catalog = ctx.get_catalog(&catalog_name).await?; let db = catalog.get_database(&tenant, &db_name).await?; let db_id = db.get_db_info().database_id.db_id; @@ -649,32 +655,3 @@ fn parse_database_name(ctx: &Arc, name: &str) -> Result<(Strin Ok((catalog, database)) } - -/// Parse table name in format "table", "db.table", or "catalog.db.table". -/// Correctly handles quoted identifiers and normalizes them according to session settings. -/// Returns (catalog, database, table). -fn parse_table_name(ctx: &Arc, name: &str) -> Result<(String, String, String)> { - let trimmed = name.trim(); - if trimmed.is_empty() { - return Err(ErrorCode::BadArguments("object_name must not be empty")); - } - - let settings = ctx.get_settings(); - let name_resolution_ctx = NameResolutionContext::try_from(settings.as_ref())?; - let dialect = settings.get_sql_dialect().unwrap_or_default(); - - let table_ref = parse_table_ref(trimmed, dialect) - .map_err(|e| ErrorCode::BadArguments(format!("Invalid table name '{}': {}", name, e.1)))?; - - let catalog = table_ref - .catalog - .map(|i| normalize_identifier(&i, &name_resolution_ctx).name) - .unwrap_or_else(|| ctx.get_current_catalog()); - let database = table_ref - .database - .map(|i| normalize_identifier(&i, &name_resolution_ctx).name) - .unwrap_or_else(|| ctx.get_current_database()); - let table = normalize_identifier(&table_ref.table, &name_resolution_ctx).name; - - Ok((catalog, database, table)) -} diff --git a/src/query/service/tests/it/catalogs/database_catalog.rs b/src/query/service/tests/it/catalogs/database_catalog.rs index 2a1045329ec2b..cf2f0c09e9e3d 100644 --- a/src/query/service/tests/it/catalogs/database_catalog.rs +++ b/src/query/service/tests/it/catalogs/database_catalog.rs @@ -166,6 +166,7 @@ async fn test_catalogs_table() -> anyhow::Result<()> { created_on, ..TableMeta::default() }, + lineage_updates: vec![], as_dropped: false, table_properties: None, table_partition: None, diff --git a/src/query/service/tests/it/catalogs/test_error_handling.rs b/src/query/service/tests/it/catalogs/test_error_handling.rs index 569f16983dbc1..b8258fb22effd 100644 --- a/src/query/service/tests/it/catalogs/test_error_handling.rs +++ b/src/query/service/tests/it/catalogs/test_error_handling.rs @@ -88,6 +88,7 @@ async fn test_get_table_error_handling() -> anyhow::Result<()> { engine: "MEMORY".to_string(), ..TableMeta::default() }, + lineage_updates: vec![], as_dropped: false, table_properties: None, table_partition: None, @@ -222,6 +223,7 @@ async fn test_get_table_by_info_error_handling() -> anyhow::Result<()> { engine: "MEMORY".to_string(), ..TableMeta::default() }, + lineage_updates: vec![], as_dropped: false, table_properties: None, table_partition: None, diff --git a/src/query/service/tests/it/storages/fuse/operations/mutation/block_compact_mutator.rs b/src/query/service/tests/it/storages/fuse/operations/mutation/block_compact_mutator.rs index 84de6451b1d48..8c82f809e9330 100644 --- a/src/query/service/tests/it/storages/fuse/operations/mutation/block_compact_mutator.rs +++ b/src/query/service/tests/it/storages/fuse/operations/mutation/block_compact_mutator.rs @@ -145,6 +145,7 @@ async fn do_compact(ctx: Arc, table: Arc) -> Result anyhow::Res })), ..TableMeta::default() }, + lineage_updates: vec![], as_dropped: false, table_properties: None, table_partition: None, diff --git a/src/query/service/tests/it/storages/testdata/configs_table_basic.txt b/src/query/service/tests/it/storages/testdata/configs_table_basic.txt index 4bb1201e0ed03..58e69214926db 100644 --- a/src/query/service/tests/it/storages/testdata/configs_table_basic.txt +++ b/src/query/service/tests/it/storages/testdata/configs_table_basic.txt @@ -207,6 +207,7 @@ DB.Table: 'system'.'configs', Table: configs-table_id:1, ver:0, Engine: SystemCo | 'query' | 'jwks_refresh_timeout' | '10' | '' | | 'query' | 'jwt_key_file' | '' | '' | | 'query' | 'jwt_key_files' | '' | '' | +| 'query' | 'lineage.capture_enabled' | 'false' | '' | | 'query' | 'management_mode' | 'false' | '' | | 'query' | 'max_active_sessions' | '256' | '' | | 'query' | 'max_cached_queries_profiles' | '50' | '' | 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..aae8797d78a15 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 @@ -25,12 +25,15 @@ use databend_common_catalog::table::TimeNavigation; use databend_common_catalog::table_with_options::check_with_opt_valid; 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_config::GlobalConfig; use databend_common_exception::ErrorCode; use databend_common_exception::Result; use databend_common_storages_basic::view_table::QUERY; use databend_storages_common_table_meta::table::get_change_type; use crate::BindContext; +use crate::LineageSourceRelation; +use crate::ViewLineageSourceColumn; use crate::binder::Binder; use crate::binder::ViewIdent; use crate::binder::util::TableIdentifier; @@ -189,17 +192,25 @@ impl Binder { { let change_type = get_change_type(&table_name_alias); if change_type.is_some() { - let table_index = self.metadata.write().add_table( - catalog, - database.clone(), - table_meta.clone(), - branch_name, - table_name_alias, - !bind_context.binding_views.is_empty(), - bind_context.planning_agg_index, - false, - cte_suffix_name, - ); + let stream_lineage_source = stream_lineage_source_relation(&table_meta); + let table_index = { + let mut metadata = self.metadata.write(); + let table_index = metadata.add_table( + catalog, + database.clone(), + table_meta.clone(), + branch_name, + table_name_alias, + !bind_context.binding_views.is_empty(), + bind_context.planning_agg_index, + false, + cte_suffix_name, + ); + if let Some(stream_lineage_source) = stream_lineage_source { + metadata.set_stream_lineage_source(table_index, stream_lineage_source); + } + table_index + }; let (s_expr, mut bind_context) = self.bind_base_table( bind_context, database.as_str(), @@ -280,11 +291,11 @@ impl Binder { new_bind_context.binding_views.insert(view_ident); if let Statement::Query(query) = &stmt { self.metadata.write().add_table( - catalog, + catalog.clone(), database.clone(), - table_meta, - branch_name, - table_name_alias, + table_meta.clone(), + branch_name.clone(), + table_name_alias.clone(), false, false, false, @@ -302,6 +313,13 @@ impl Binder { column.table_name = Some(self.normalize_identifier(table).name); } } + self.add_view_lineage_source_columns( + &new_bind_context, + catalog.as_str(), + database.as_str(), + table_name.as_str(), + &table_meta, + ); // Restore binding_views to the outer scope's value so the // current view does not leak into sibling/parent contexts. new_bind_context.binding_views = bind_context.binding_views.clone(); @@ -343,4 +361,64 @@ impl Binder { } } } + + fn add_view_lineage_source_columns( + &mut self, + bind_context: &BindContext, + catalog: &str, + database: &str, + view_name: &str, + view: &std::sync::Arc, + ) { + if !GlobalConfig::instance() + .query + .common + .lineage + .capture_enabled + { + return; + } + + 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 schema = view.schema(); + let mut metadata = self.metadata.write(); + for (idx, column) in bind_context.columns.iter().enumerate() { + let Some(field) = schema.fields().get(idx) else { + continue; + }; + metadata.add_view_lineage_source_column(column.index, ViewLineageSourceColumn { + relation: relation.clone(), + name: field.name().clone(), + id: field.column_id, + }); + } + } +} + +fn stream_lineage_source_relation( + table: &std::sync::Arc, +) -> Option { + if !GlobalConfig::instance() + .query + .common + .lineage + .capture_enabled + { + return None; + } + + 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..8243eabe0c663 100644 --- a/src/query/sql/src/planner/binder/ddl/view.rs +++ b/src/query/sql/src/planner/binder/ddl/view.rs @@ -18,9 +18,11 @@ 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; +use databend_common_config::GlobalConfig; use databend_common_exception::Result; use databend_common_expression::DataField; use databend_common_expression::DataSchemaRefExt; @@ -65,6 +67,16 @@ impl Binder { }; query.walk_mut(&mut visitor)?; let subquery = format!("{}", query); + let query_plan = if GlobalConfig::instance() + .query + .common + .lineage + .capture_enabled + { + Some(Box::new(self.view_query_plan(&query).await?)) + } else { + None + }; let plan = CreateViewPlan { create_option: create_option.clone().into(), @@ -74,6 +86,7 @@ impl Binder { view_name, column_names, subquery, + query_plan, }; Ok(Plan::CreateView(plan.into())) } @@ -225,4 +238,10 @@ 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..362dd3c934b86 100644 --- a/src/query/sql/src/planner/binder/insert.rs +++ b/src/query/sql/src/planner/binder/insert.rs @@ -20,6 +20,7 @@ use databend_common_ast::ast::InsertSource; use databend_common_ast::ast::InsertStmt; use databend_common_ast::ast::Statement; use databend_common_catalog::session_type::SessionType; +use databend_common_config::GlobalConfig; use databend_common_exception::ErrorCode; use databend_common_exception::Result; use databend_common_expression::DataSchemaRef; @@ -311,6 +312,12 @@ impl Binder { overwrite: *overwrite, source: input_source?, table_info: None, + lineage_target_table_id: GlobalConfig::instance() + .query + .common + .lineage + .capture_enabled + .then_some(table.get_table_info().ident.table_id), }; Ok(Plan::Insert(Box::new(plan))) diff --git a/src/query/sql/src/planner/lineage.rs b/src/query/sql/src/planner/lineage.rs new file mode 100644 index 0000000000000..c6942a542ed7d --- /dev/null +++ b/src/query/sql/src/planner/lineage.rs @@ -0,0 +1,1649 @@ +// 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_exception::ErrorCode; +use databend_common_exception::Result; +use databend_common_expression::ColumnId; +use databend_common_expression::FieldIndex; +use databend_common_expression::TableField; + +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, + pub downstreams: Vec, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum QueryLineageKind { + Ctas, + Dml, + CreateView, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct LineageDownstream { + pub relation: QueryLineageRelation, + pub upstreams: Vec, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct LineageUpstream { + pub relation: QueryLineageRelation, + pub columns: Vec, +} + +#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)] +pub struct QueryLineageColumnEdge { + pub upstream: QueryLineageColumn, + pub downstream: QueryLineageColumn, +} + +#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)] +pub struct QueryLineageRelation { + pub catalog: String, + pub database: String, + pub name: String, + pub id: Option, + pub kind: QueryLineageRelationKind, +} + +#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)] +pub enum QueryLineageRelationKind { + Table, + View, +} + +#[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 downstreams = 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 { + upstream: source_column, + downstream: column.target_column.clone(), + } + })); + } + } + + LineageDownstream { + relation: target.target, + upstreams: sources + .into_iter() + .map(|(relation, mut columns)| { + columns.sort(); + columns.dedup(); + LineageUpstream { relation, columns } + }) + .collect(), + } + }) + .collect(); + + QueryLineage { kind, downstreams } + } +} + +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, + 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, + 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, + 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), + kind: QueryLineageRelationKind::Table, + }, + plan.schema.fields(), + )?; + (QueryLineageKind::Dml, target_bindings) + } + 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, + 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 relation = QueryLineageRelation { + catalog: mutation.catalog_name.clone(), + database: mutation.database_name.clone(), + name: mutation.table_name.clone(), + id: target_metadata + .tables() + .get(mutation.target_table_index) + .map(|table| table.table().get_table_info().ident.table_id), + 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 + } +} + +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), + 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 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), + kind, + })); + } + Ok(Some(QueryLineageRelation { + catalog: table.catalog().to_string(), + database: table.database().to_string(), + name: table.name().to_string(), + id: Some(table.table().get_table_info().ident.table_id), + 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::CreateOption; + use databend_common_meta_app::schema::DatabaseType; + 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, + } + + #[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() + } + } + + #[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_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_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 { + metadata.write().add_table( + "default".to_string(), + "default".to_string(), + fake_table(table_id, table_name, columns), + 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(table_id: u64, table_name: &str, columns: &[&str]) -> Arc { + Arc::new(FakeTable { + table_info: table_info(table_id, table_name, columns, "FUSE"), + stream_source_table_info: None, + stream_columns: vec![], + }) + } + + 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, + )], + }) + } + + fn table_info(table_id: u64, table_name: &str, columns: &[&str], engine: &str) -> 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: Arc::new(CatalogInfo::default()), + db_type: DatabaseType::NormalDB, + } + } + + 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, + 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 + .downstreams + .iter() + .find(|table| table.relation.name == target_table) + .unwrap_or_else(|| panic!("missing target table {target_table}: {lineage:?}")) + .upstreams + .iter() + .flat_map(|table| { + table + .columns + .iter() + .filter(|column| column.downstream.name == target_column) + .map(|column| format!("{}.{}", table.relation.name, column.upstream.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 abfe3e9ec64a7..1706317bac6fd 100644 --- a/src/query/sql/src/planner/optimizer/optimizer.rs +++ b/src/query/sql/src/planner/optimizer/optimizer.rs @@ -221,7 +221,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..8089eb93e62e2 100644 --- a/src/query/sql/src/planner/plans/insert.rs +++ b/src/query/sql/src/planner/plans/insert.rs @@ -89,6 +89,9 @@ 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, } impl PartialEq for Insert { @@ -126,6 +129,7 @@ impl Insert { overwrite, // table_info only used create table as select. table_info: _, + lineage_target_table_id: _, 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..3c4e832939f03 100644 --- a/src/query/sql/test-support/src/lite_context.rs +++ b/src/query/sql/test-support/src/lite_context.rs @@ -74,6 +74,7 @@ use databend_common_expression::Scalar; use databend_common_expression::TableDataType; use databend_common_expression::TableField; use databend_common_expression::TableSchema; +use databend_common_expression::infer_schema_type; use databend_common_expression::types::NumberDataType; use databend_common_io::prelude::InputFormatSettings; use databend_common_io::prelude::OutputFormatSettings; @@ -146,14 +147,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 +169,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 +265,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 +329,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 +354,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 +905,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 +1122,79 @@ 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"); + }; + let fields = bind_context + .columns + .iter() + .enumerate() + .map(|(idx, column)| { + Ok(TableField::new_from_column_id( + &column.column_name, + infer_schema_type(&column.data_type)?, + idx as ColumnId, + )) + }) + .collect::>>()?; + + self.register_view( + &ReplayView { + catalog: self.current_catalog.clone(), + database: database.to_string(), + view: view_name.to_string(), + query: query.to_string(), + }, + Arc::new(TableSchema::new(fields)), + ) + } + + 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 +1205,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 +1213,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..4318f4bae6e5e --- /dev/null +++ b/src/query/sql/tests/it/planner/lineage.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::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_sql::LineageDownstream; +use databend_common_sql::LineageUpstream; +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 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 lineage = query_lineage_from_sql(&ctx, sql).await?; + 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_does_not_bind_query_when_lineage_capture_is_disabled() -> 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![table_column(&ctx, "v", "vx").await?], + ); + + assert_eq!(lineage, expected); + Ok(()) + }) + }) + .map_err(|err| ErrorCode::Internal(err.to_string()))? + .join() + .map_err(|_| ErrorCode::Internal("lineage 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.query.common.lineage.capture_enabled = true; + 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 + .downstreams + .iter() + .find(|downstream| downstream.relation.name == target_table) + .unwrap_or_else(|| panic!("missing target table {target_table}: {lineage:?}")) + .upstreams + .iter() + .flat_map(|from_relation| { + from_relation + .columns + .iter() + .filter(|edge| edge.downstream.name == target_column) + .map(|edge| format!("{}.{}", from_relation.relation.name, edge.upstream.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, + downstream: QueryLineageRelation, + upstream: QueryLineageRelation, + downstream_column: QueryLineageColumn, + upstream_columns: Vec, +) -> QueryLineage { + QueryLineage { + kind, + downstreams: vec![LineageDownstream { + relation: downstream, + upstreams: vec![LineageUpstream { + relation: upstream, + columns: upstream_columns + .into_iter() + .map(|upstream| QueryLineageColumnEdge { + upstream, + downstream: downstream_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, + 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/sqllogictests/suites/ee/09_ee_lineage/get_lineage.test b/tests/sqllogictests/suites/ee/09_ee_lineage/get_lineage.test new file mode 100644 index 0000000000000..0eafdad885b1d --- /dev/null +++ b/tests/sqllogictests/suites/ee/09_ee_lineage/get_lineage.test @@ -0,0 +1,778 @@ +## Copyright 2023 Databend Cloud +## +## Licensed under the Elastic 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 +## +## https://www.elastic.co/licensing/elastic-license +## +## 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. + +# Requires databend-query config: +# [query.lineage] +# capture_enabled = true + +# Lineage notation used by this test: +# - `upstream -> downstream` always denotes the data-flow direction, independent +# of whether GET_LINEAGE is queried with UPSTREAM or DOWNSTREAM. +# - CTAS and DML edges are `TABLE[id] -> TABLE[id]`; stable table and column ids +# keep these edges valid across renames. +# - View dependencies are `TABLE[name] -> VIEW[id]`; the referenced table and +# its columns are matched by name, while the view side uses stable ids. +# - Columns used only by joins or filters do not form column-lineage edges. + +# ----------------------------------------------------------------------------- +# Test setup +# ----------------------------------------------------------------------------- + +statement ok +DROP DATABASE IF EXISTS lineage_slt; + +statement ok +DROP DATABASE IF EXISTS lineage_slt_other; + +statement ok +CREATE DATABASE lineage_slt; + +statement ok +USE lineage_slt; + +statement ok +SET enable_distributed_merge_into = 1; + +statement ok +SET enable_distributed_replace_into = 1; + +statement ok +CREATE TABLE src(id INT, a INT, b INT, filter_key INT, filter_flag INT); + +statement ok +INSERT INTO src VALUES + (1, 10, 100, 7, 1), + (2, 20, 200, 8, 1); + +# ----------------------------------------------------------------------------- +# CTAS: src[id] -> ctas[id] +# +# Query the edge in the upstream direction from the target, then query column +# lineage in both directions. `filter_flag` is filter-only and must be absent. +# ----------------------------------------------------------------------------- + +statement ok +CREATE TABLE ctas AS +SELECT id AS dst_id, a + b AS payload +FROM src +WHERE filter_flag > 0; + +query ITTTTT +SELECT distance, source_object_domain, source_object_name, target_object_domain, target_object_name, target_status +FROM get_lineage('lineage_slt.ctas', 'TABLE', 'UPSTREAM', 1) +ORDER BY distance, source_object_name, target_object_name; +---- +1 TABLE lineage_slt.src TABLE lineage_slt.ctas ACTIVE + +query ITTTT +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_slt.ctas.payload', 'COLUMN', 'UPSTREAM', 1) +ORDER BY distance, source_column_name, target_column_name; +---- +1 lineage_slt.src a lineage_slt.ctas payload +1 lineage_slt.src b lineage_slt.ctas payload + +query I +SELECT count(*) +FROM get_lineage('lineage_slt.src.filter_flag', 'COLUMN', 'DOWNSTREAM', 2); +---- +0 + +# VACUUM TABLE only removes obsolete table data. It must not remove lineage for +# the active table id. + +statement ok +VACUUM TABLE ctas; + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_slt.ctas', 'TABLE', 'UPSTREAM', 1) +ORDER BY distance, source_object_name, target_object_name; +---- +1 lineage_slt.src lineage_slt.ctas + +# ----------------------------------------------------------------------------- +# Cross-database by-id resolution +# +# Both databases contain a table named `src`. Resolving the CTAS upstream id +# must return `lineage_slt_other.src`, not the same-named table in the current +# database and not an empty result. +# ----------------------------------------------------------------------------- + +statement ok +CREATE DATABASE lineage_slt_other; + +statement ok +CREATE TABLE lineage_slt_other.src(id INT); + +statement ok +CREATE TABLE lineage_slt_other.ctas AS +SELECT id +FROM lineage_slt_other.src; + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_slt_other.ctas', 'TABLE', 'UPSTREAM', 1) +ORDER BY distance, source_object_name, target_object_name; +---- +1 lineage_slt_other.src lineage_slt_other.ctas + +# ----------------------------------------------------------------------------- +# Dropped source database invalidates by-id lineage +# +# Keep the target in `lineage_slt` while dropping the source database. The +# target remains queryable, but its source table is no longer ACTIVE. +# ----------------------------------------------------------------------------- + +statement ok +CREATE TABLE drop_db_dst AS +SELECT id +FROM lineage_slt_other.src; + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_slt.drop_db_dst', 'TABLE', 'UPSTREAM', 1) +ORDER BY distance, source_object_name, target_object_name; +---- +1 lineage_slt_other.src lineage_slt.drop_db_dst + +statement ok +DROP DATABASE lineage_slt_other; + +query I +SELECT count(*) +FROM get_lineage('lineage_slt.drop_db_dst', 'TABLE', 'UPSTREAM', 1); +---- +0 + +# ----------------------------------------------------------------------------- +# INSERT SELECT: ctas[id] -> insert_dst[id] +# +# The object query walks downstream two hops from `src`; the column query walks +# one hop downstream from `ctas.payload`. +# ----------------------------------------------------------------------------- + +statement ok +CREATE TABLE insert_dst(dst_id INT, payload INT); + +statement ok +INSERT INTO insert_dst +SELECT dst_id, payload +FROM ctas +WHERE payload > 0; + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_slt.src', 'TABLE', 'DOWNSTREAM', 2) +WHERE target_object_name = 'lineage_slt.insert_dst' +ORDER BY distance, source_object_name, target_object_name; +---- +2 lineage_slt.ctas lineage_slt.insert_dst + +query ITTTT +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_slt.ctas.payload', 'COLUMN', 'DOWNSTREAM', 1) +ORDER BY distance, source_column_name, target_column_name; +---- +1 lineage_slt.ctas payload lineage_slt.insert_dst payload + +# ----------------------------------------------------------------------------- +# Multi-table INSERT: src[id] -> multi_dst1[id], multi_dst2[id] +# +# Both targets must persist object and column lineage. Source column `b` fans +# out to a different destination column in each target. +# ----------------------------------------------------------------------------- + +statement ok +CREATE TABLE multi_dst1(x INT, y INT); + +statement ok +CREATE TABLE multi_dst2(x INT, y INT); + +statement ok +INSERT ALL + INTO multi_dst1 VALUES(a, b) + INTO multi_dst2(x, y) VALUES(b, id) +SELECT id, a, b FROM src; + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_slt.src', 'TABLE', 'DOWNSTREAM', 1) +WHERE target_object_name IN ('lineage_slt.multi_dst1', 'lineage_slt.multi_dst2') +ORDER BY distance, source_object_name, target_object_name; +---- +1 lineage_slt.src lineage_slt.multi_dst1 +1 lineage_slt.src lineage_slt.multi_dst2 + +query ITTTT +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_slt.src.b', 'COLUMN', 'DOWNSTREAM', 1) +WHERE target_object_name IN ('lineage_slt.multi_dst1', 'lineage_slt.multi_dst2') +ORDER BY distance, source_column_name, target_object_name, target_column_name; +---- +1 lineage_slt.src b lineage_slt.multi_dst1 y +1 lineage_slt.src b lineage_slt.multi_dst2 x + +query ITTTT +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_slt.multi_dst1.x', 'COLUMN', 'UPSTREAM', 1) +ORDER BY distance, source_column_name, target_column_name; +---- +1 lineage_slt.src a lineage_slt.multi_dst1 x + +query ITTTT +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_slt.multi_dst2.y', 'COLUMN', 'UPSTREAM', 1) +ORDER BY distance, source_column_name, target_column_name; +---- +1 lineage_slt.src id lineage_slt.multi_dst2 y + +# ----------------------------------------------------------------------------- +# CREATE VIEW: src[name] -> src_view[id] +# +# The view can be queried upstream through its stable id. A downstream query +# from the source must also find the by-name dependency and map source column +# names to stable view column ids. +# ----------------------------------------------------------------------------- + +statement ok +CREATE VIEW src_view AS +SELECT a + b AS v_payload +FROM src +WHERE filter_key > 0; + +query ITTTTT +SELECT distance, source_object_domain, source_object_name, target_object_domain, target_object_name, target_status +FROM get_lineage('lineage_slt.src_view', 'VIEW', 'UPSTREAM', 1) +ORDER BY distance, source_object_name, target_object_name; +---- +1 TABLE lineage_slt.src VIEW lineage_slt.src_view ACTIVE + +query ITTTT +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_slt.src.a', 'COLUMN', 'DOWNSTREAM', 1) +WHERE target_object_name = 'lineage_slt.src_view' +ORDER BY distance, source_column_name, target_column_name; +---- +1 lineage_slt.src a lineage_slt.src_view v_payload + +# ----------------------------------------------------------------------------- +# Rename and drop behavior +# +# Renaming the by-name source temporarily invalidates the view edge; renaming +# it back restores the edge. By-id CTAS/DML object and column edges follow +# renames. Dropped targets and explicitly dropped views are not returned. +# ----------------------------------------------------------------------------- + +statement ok +ALTER TABLE src RENAME TO src_renamed; + +query I +SELECT count(*) +FROM get_lineage('lineage_slt.src_view', 'VIEW', 'UPSTREAM', 1); +---- +0 + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_slt.ctas', 'TABLE', 'UPSTREAM', 1) +ORDER BY distance, source_object_name, target_object_name; +---- +1 lineage_slt.src_renamed lineage_slt.ctas + +statement ok +ALTER TABLE src_renamed RENAME TO src; + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_slt.src_view', 'VIEW', 'UPSTREAM', 1) +ORDER BY distance, source_object_name, target_object_name; +---- +1 lineage_slt.src lineage_slt.src_view + +# Data Movement column refs use stable ids, so renaming a source column keeps +# the CTAS edge. View source columns are name-addressed, so the matching column +# edge disappears until the old name is restored. + +statement ok +ALTER TABLE src RENAME COLUMN a TO a_renamed; + +query ITTTT +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_slt.ctas.payload', 'COLUMN', 'UPSTREAM', 1) +ORDER BY distance, source_column_name, target_column_name; +---- +1 lineage_slt.src a_renamed lineage_slt.ctas payload +1 lineage_slt.src b lineage_slt.ctas payload + +query I +SELECT count(*) +FROM get_lineage('lineage_slt.src.a_renamed', 'COLUMN', 'DOWNSTREAM', 1) +WHERE target_object_name = 'lineage_slt.src_view'; +---- +0 + +statement ok +ALTER TABLE src RENAME COLUMN a_renamed TO a; + +query ITTTT +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_slt.src.a', 'COLUMN', 'DOWNSTREAM', 1) +WHERE target_object_name = 'lineage_slt.src_view' +ORDER BY distance, source_column_name, target_column_name; +---- +1 lineage_slt.src a lineage_slt.src_view v_payload + +statement ok +ALTER TABLE ctas RENAME COLUMN payload TO payload2; + +query ITTTT +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_slt.ctas.payload2', 'COLUMN', 'UPSTREAM', 1) +ORDER BY distance, source_column_name, target_column_name; +---- +1 lineage_slt.src a lineage_slt.ctas payload2 +1 lineage_slt.src b lineage_slt.ctas payload2 + +statement ok +ALTER TABLE ctas RENAME TO ctas_renamed; + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_slt.ctas_renamed', 'TABLE', 'DOWNSTREAM', 1) +ORDER BY distance, source_object_name, target_object_name; +---- +1 lineage_slt.ctas_renamed lineage_slt.insert_dst + +statement ok +ALTER TABLE insert_dst RENAME TO insert_dst_renamed; + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_slt.ctas_renamed', 'TABLE', 'DOWNSTREAM', 1) +ORDER BY distance, source_object_name, target_object_name; +---- +1 lineage_slt.ctas_renamed lineage_slt.insert_dst_renamed + +statement ok +DROP TABLE insert_dst_renamed; + +query I +SELECT count(*) +FROM get_lineage('lineage_slt.ctas_renamed', 'TABLE', 'DOWNSTREAM', 1); +---- +0 + +statement ok +UNDROP TABLE insert_dst_renamed; + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_slt.ctas_renamed', 'TABLE', 'DOWNSTREAM', 1) +ORDER BY distance, source_object_name, target_object_name; +---- +1 lineage_slt.ctas_renamed lineage_slt.insert_dst_renamed + +statement ok +DROP VIEW src_view; + +query I +SELECT count(*) +FROM get_lineage('lineage_slt.src.a', 'COLUMN', 'DOWNSTREAM', 1) +WHERE target_object_name = 'lineage_slt.src_view'; +---- +0 + +# ----------------------------------------------------------------------------- +# REPLACE INTO: replace_src[id] -> replace_dst[id] +# +# Expression inputs `a` and `b` flow to `x`; the filter-only column does not. +# ----------------------------------------------------------------------------- + +statement ok +CREATE TABLE replace_src(id INT NOT NULL, a INT, b INT, filter_flag INT); + +statement ok +CREATE TABLE replace_dst(id INT NOT NULL, x INT); + +statement ok +INSERT INTO replace_src VALUES (1, 3, 30, 1), (2, 4, 40, 1); + +statement ok +REPLACE INTO replace_dst(id, x) ON(id) +SELECT id, a + b +FROM replace_src +WHERE filter_flag > 0; + +query ITTTT +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_slt.replace_dst.x', 'COLUMN', 'UPSTREAM', 1) +ORDER BY distance, source_column_name, target_column_name; +---- +1 lineage_slt.replace_src a lineage_slt.replace_dst x +1 lineage_slt.replace_src b lineage_slt.replace_dst x + +query I +SELECT count(*) +FROM get_lineage('lineage_slt.replace_src.filter_flag', 'COLUMN', 'DOWNSTREAM', 1); +---- +0 + +# ----------------------------------------------------------------------------- +# UPDATE FROM: update_src[id] -> update_dst[id] +# +# Only columns contributing to the assigned value form lineage. Join and +# filter predicates are excluded. +# ----------------------------------------------------------------------------- + +statement ok +CREATE TABLE update_src(id INT, a INT, b INT, filter_flag INT); + +statement ok +CREATE TABLE update_dst(id INT, x INT); + +statement ok +INSERT INTO update_src VALUES (1, 5, 50, 1); + +statement ok +INSERT INTO update_dst VALUES (1, 0); + +statement ok +UPDATE update_dst +SET x = s.a + s.b +FROM update_src AS s +WHERE update_dst.id = s.id AND s.filter_flag > 0; + +query ITTTT +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_slt.update_dst.x', 'COLUMN', 'UPSTREAM', 1) +ORDER BY distance, source_column_name, target_column_name; +---- +1 lineage_slt.update_src a lineage_slt.update_dst x +1 lineage_slt.update_src b lineage_slt.update_dst x + +query I +SELECT count(*) +FROM get_lineage('lineage_slt.update_src.filter_flag', 'COLUMN', 'DOWNSTREAM', 1); +---- +0 + +# ----------------------------------------------------------------------------- +# Self-write suppression +# +# `self_write[id] -> self_write[id]` is not a useful dependency and must not be +# persisted at either object or column level. +# ----------------------------------------------------------------------------- + +statement ok +CREATE TABLE self_write(id INT, value INT); + +statement ok +INSERT INTO self_write VALUES (1, 10); + +statement ok +INSERT INTO self_write +SELECT id + 1, value +FROM self_write; + +query I +SELECT count(*) +FROM get_lineage('lineage_slt.self_write', 'TABLE', 'UPSTREAM', 1); +---- +0 + +query I +SELECT count(*) +FROM get_lineage('lineage_slt.self_write.value', 'COLUMN', 'UPSTREAM', 1); +---- +0 + +# ----------------------------------------------------------------------------- +# COUNT(*) +# +# COUNT(*) depends on row cardinality rather than a source column and does not +# create lineage. +# +# TODO: COUNT(column) should depend on the named column. The optimizer can fold +# it to a metadata-derived constant and currently loses the original argument; +# preserve lineage metadata across that rewrite before enabling the behavior. +# ----------------------------------------------------------------------------- + +statement ok +CREATE TABLE agg_src(id INT, value INT); + +statement ok +INSERT INTO agg_src VALUES (1, 10), (2, NULL); + +statement ok +CREATE TABLE count_star_dst(cnt UINT64); + +statement ok +INSERT INTO count_star_dst +SELECT count(*) FROM agg_src; + +query I +SELECT count(*) +FROM get_lineage('lineage_slt.count_star_dst', 'TABLE', 'UPSTREAM', 1); +---- +0 + +query I +SELECT count(*) +FROM get_lineage('lineage_slt.count_star_dst.cnt', 'COLUMN', 'UPSTREAM', 1); +---- +0 + +# ----------------------------------------------------------------------------- +# SUM(column) +# +# `SUM(value)` retains its aggregate argument in the plan, so column lineage +# passes through the aggregate operator. +# ----------------------------------------------------------------------------- + +statement ok +CREATE TABLE sum_dst(total BIGINT); + +statement ok +INSERT INTO sum_dst +SELECT sum(value) FROM agg_src; + +query ITT +SELECT distance, source_object_name, target_object_name +FROM get_lineage('lineage_slt.sum_dst', 'TABLE', 'UPSTREAM', 1) +ORDER BY distance, source_object_name, target_object_name; +---- +1 lineage_slt.agg_src lineage_slt.sum_dst + +query ITTTT +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_slt.sum_dst.total', 'COLUMN', 'UPSTREAM', 1) +ORDER BY distance, source_column_name, target_column_name; +---- +1 lineage_slt.agg_src value lineage_slt.sum_dst total + +query I +SELECT count(*) +FROM get_lineage('lineage_slt.agg_src.id', 'COLUMN', 'DOWNSTREAM', 1) +WHERE target_object_name = 'lineage_slt.sum_dst'; +---- +0 + +# ----------------------------------------------------------------------------- +# Automatic materialized CTE +# +# Auto-materialized CTE consumers resolve through MaterializedCTERef nodes to +# the base table. Columns used only to join consumers remain excluded. +# +# TODO: Explicit `AS MATERIALIZED` is implemented through a query-scoped +# temporary table. Preserve its column lineage outside persistent meta before +# adding explicit MCTE coverage here. +# ----------------------------------------------------------------------------- + +statement ok +CREATE TABLE mcte_src(a INT, b INT, k INT); + +statement ok +INSERT INTO mcte_src VALUES (1, 10, 1), (2, 20, 2); + +statement ok +SET enable_auto_materialize_cte = 1; + +statement ok +CREATE TABLE mcte_auto_dst(x INT); + +statement ok +INSERT INTO mcte_auto_dst +WITH q AS ( + SELECT a, b, k FROM mcte_src +) +SELECT q1.a + q2.b +FROM q AS q1 +JOIN q AS q2 ON q1.k = q2.k; + +statement ok +UNSET enable_auto_materialize_cte; + +query ITTTT +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_slt.mcte_auto_dst.x', 'COLUMN', 'UPSTREAM', 1) +ORDER BY distance, source_column_name, target_column_name; +---- +1 lineage_slt.mcte_src a lineage_slt.mcte_auto_dst x +1 lineage_slt.mcte_src b lineage_slt.mcte_auto_dst x + +query I +SELECT count(*) +FROM get_lineage('lineage_slt.mcte_src.k', 'COLUMN', 'DOWNSTREAM', 1) +WHERE target_object_name = 'lineage_slt.mcte_auto_dst'; +---- +0 + +# ----------------------------------------------------------------------------- +# Scalar subquery +# +# Values from both the outer query and scalar subquery contribute to `x`. +# Correlation and filter-only columns do not form column-lineage edges. +# ----------------------------------------------------------------------------- + +statement ok +CREATE TABLE subquery_left(id INT, a INT); + +statement ok +CREATE TABLE subquery_right(id INT, b INT, filter_flag INT); + +statement ok +CREATE TABLE subquery_dst(x INT); + +statement ok +INSERT INTO subquery_left VALUES (1, 10); + +statement ok +INSERT INTO subquery_right VALUES (1, 20, 1); + +statement ok +INSERT INTO subquery_dst +SELECT l.a + ( + SELECT max(r.b) + FROM subquery_right AS r + WHERE r.id = l.id AND r.filter_flag > 0 +) +FROM subquery_left AS l; + +query ITTTT +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_slt.subquery_dst.x', 'COLUMN', 'UPSTREAM', 1) +ORDER BY distance, source_object_name, source_column_name; +---- +1 lineage_slt.subquery_left a lineage_slt.subquery_dst x +1 lineage_slt.subquery_right b lineage_slt.subquery_dst x + +query I +SELECT count(*) +FROM get_lineage('lineage_slt.subquery_right.filter_flag', 'COLUMN', 'DOWNSTREAM', 1); +---- +0 + +# ----------------------------------------------------------------------------- +# JOIN USING +# +# The unqualified USING column is represented by the left input. The right id +# participates only in the join condition and must not become a value source. +# ----------------------------------------------------------------------------- + +statement ok +CREATE TABLE using_left(id INT, a INT); + +statement ok +CREATE TABLE using_right(id INT, b INT); + +statement ok +CREATE TABLE using_dst(id INT); + +statement ok +INSERT INTO using_left VALUES (1, 10); + +statement ok +INSERT INTO using_right VALUES (1, 20); + +statement ok +INSERT INTO using_dst +SELECT id +FROM using_left +JOIN using_right USING (id); + +query ITTTT +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_slt.using_dst.id', 'COLUMN', 'UPSTREAM', 1) +ORDER BY distance, source_object_name, source_column_name; +---- +1 lineage_slt.using_left id lineage_slt.using_dst id + +query I +SELECT count(*) +FROM get_lineage('lineage_slt.using_right.id', 'COLUMN', 'DOWNSTREAM', 1); +---- +0 + +# ----------------------------------------------------------------------------- +# MERGE: merge_src[id] -> merge_dst[id] +# +# Verify column mapping for matched updates and the not-matched insert. Columns +# used only in match predicates (`c` and `d`) must not appear downstream. +# ----------------------------------------------------------------------------- + +statement ok +CREATE TABLE merge_src(id INT, a INT, b INT, c INT, d INT); + +statement ok +CREATE TABLE merge_dst(id INT, x INT, y INT); + +statement ok +INSERT INTO merge_src VALUES (1, 6, 60, 1, 0), (2, 7, 70, 0, 1); + +statement ok +INSERT INTO merge_dst VALUES (1, 0, 0); + +query II +MERGE INTO merge_dst AS d +USING merge_src AS s +ON d.id = s.id +WHEN MATCHED AND s.c > 0 THEN UPDATE SET x = s.a +WHEN MATCHED AND s.d > 0 THEN UPDATE SET y = s.b +WHEN NOT MATCHED THEN INSERT (id, x, y) VALUES (s.id, s.a, s.b); +---- +1 1 + +query ITTTT +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_slt.merge_dst.id', 'COLUMN', 'UPSTREAM', 1) +ORDER BY distance, source_column_name, target_column_name; +---- +1 lineage_slt.merge_src id lineage_slt.merge_dst id + +query ITTTT +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_slt.merge_dst.x', 'COLUMN', 'UPSTREAM', 1) +ORDER BY distance, source_column_name, target_column_name; +---- +1 lineage_slt.merge_src a lineage_slt.merge_dst x + +query ITTTT +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_slt.merge_dst.y', 'COLUMN', 'UPSTREAM', 1) +ORDER BY distance, source_column_name, target_column_name; +---- +1 lineage_slt.merge_src b lineage_slt.merge_dst y + +query I +SELECT count(*) +FROM get_lineage('lineage_slt.merge_src.c', 'COLUMN', 'DOWNSTREAM', 1); +---- +0 + +query I +SELECT count(*) +FROM get_lineage('lineage_slt.merge_src.d', 'COLUMN', 'DOWNSTREAM', 1); +---- +0 + +# ----------------------------------------------------------------------------- +# Cleanup +# ----------------------------------------------------------------------------- + +statement ok +DROP DATABASE IF EXISTS lineage_slt_other; + +statement ok +DROP DATABASE lineage_slt; diff --git a/tests/sqllogictests/suites/ee/09_ee_lineage/get_lineage_stream.test b/tests/sqllogictests/suites/ee/09_ee_lineage/get_lineage_stream.test new file mode 100644 index 0000000000000..fde9547ffa066 --- /dev/null +++ b/tests/sqllogictests/suites/ee/09_ee_lineage/get_lineage_stream.test @@ -0,0 +1,60 @@ +## Copyright 2023 Databend Cloud +## +## Licensed under the Elastic 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 +## +## https://www.elastic.co/licensing/elastic-license +## +## 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. + +# Requires databend-query config: +# [query.lineage] +# capture_enabled = true + +statement ok +DROP DATABASE IF EXISTS lineage_stream_slt; + +statement ok +CREATE DATABASE lineage_stream_slt; + +statement ok +USE lineage_stream_slt; + +statement ok +CREATE TABLE stream_src(id INT, value INT); + +statement ok +CREATE STREAM stream_changes ON TABLE stream_src APPEND_ONLY = TRUE; + +statement ok +INSERT INTO stream_src VALUES (1, 10); + +statement ok +CREATE TABLE stream_dst(id INT, value INT); + +statement ok +INSERT INTO stream_dst +SELECT id, value +FROM stream_changes; + +query ITTTT +SELECT distance, source_object_domain, source_object_name, target_object_domain, target_object_name +FROM get_lineage('lineage_stream_slt.stream_dst', 'TABLE', 'UPSTREAM', 1) +ORDER BY distance, source_object_name, target_object_name; +---- +1 TABLE lineage_stream_slt.stream_src TABLE lineage_stream_slt.stream_dst + +query ITTTT +SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name +FROM get_lineage('lineage_stream_slt.stream_dst.value', 'COLUMN', 'UPSTREAM', 1) +ORDER BY distance, source_column_name, target_column_name; +---- +1 lineage_stream_slt.stream_src value lineage_stream_slt.stream_dst value + +statement ok +DROP DATABASE lineage_stream_slt;