Skip to content

Commit 1491c63

Browse files
youngsofunyoungsofun
committed
feat(query): extract query lineage
1 parent 1e875e4 commit 1491c63

18 files changed

Lines changed: 2536 additions & 21 deletions

File tree

‎src/query/catalog/src/table.rs‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,15 @@ pub trait Table: Sync + Send {
109109

110110
fn get_table_info(&self) -> &TableInfo;
111111

112+
/// Returns the source table whose data columns a stream exposes.
113+
///
114+
/// Lineage intentionally passes through a stream to its source table.
115+
/// Views are lineage boundaries and must not use this relation-level hook;
116+
/// their output columns are annotated separately by the planner.
117+
fn stream_source_table_info(&self) -> Option<&TableInfo> {
118+
None
119+
}
120+
112121
fn get_data_source_info(&self) -> DataSourceInfo {
113122
DataSourceInfo::TableSource(self.get_table_info().clone())
114123
}

‎src/query/config/src/config.rs‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1959,6 +1959,10 @@ pub struct QueryConfig {
19591959
/// Max number of async table hook jobs running concurrently.
19601960
#[clap(long, value_name = "VALUE", default_value = "2")]
19611961
pub table_hook_async_max_concurrency: usize,
1962+
1963+
/// Lineage collection config.
1964+
#[clap(skip)]
1965+
pub lineage: LineageConfig,
19621966
}
19631967

19641968
impl Default for QueryConfig {
@@ -1974,6 +1978,27 @@ impl Default for QueryConfig {
19741978
}
19751979
}
19761980

1981+
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Args)]
1982+
#[serde(default)]
1983+
pub struct LineageConfig {
1984+
/// Whether to collect lineage for successful DDL/DML commits.
1985+
#[clap(
1986+
long,
1987+
value_name = "VALUE",
1988+
value_parser = clap::value_parser!(bool),
1989+
default_value = "false"
1990+
)]
1991+
pub capture_enabled: bool,
1992+
}
1993+
1994+
impl Default for LineageConfig {
1995+
fn default() -> Self {
1996+
Self {
1997+
capture_enabled: false,
1998+
}
1999+
}
2000+
}
2001+
19772002
impl TryInto<InnerQueryConfig> for QueryConfig {
19782003
type Error = ErrorCode;
19792004

‎src/query/service/src/interpreters/interpreter_table_create.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -250,6 +250,7 @@ impl CreateTableInterpreter {
250250
overwrite: false,
251251
source: InsertInputSource::SelectPlan(select_plan),
252252
table_info: Some(table_info),
253+
lineage_target_table_id: None,
253254
};
254255

255256
let mut pipeline = InsertInterpreter::try_create(self.ctx.clone(), insert_plan)?

‎src/query/service/tests/it/storages/testdata/configs_table_basic.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -207,6 +207,7 @@ DB.Table: 'system'.'configs', Table: configs-table_id:1, ver:0, Engine: SystemCo
207207
| 'query' | 'jwks_refresh_timeout' | '10' | '' |
208208
| 'query' | 'jwt_key_file' | '' | '' |
209209
| 'query' | 'jwt_key_files' | '' | '' |
210+
| 'query' | 'lineage.capture_enabled' | 'false' | '' |
210211
| 'query' | 'management_mode' | 'false' | '' |
211212
| 'query' | 'max_active_sessions' | '256' | '' |
212213
| 'query' | 'max_cached_queries_profiles' | '50' | '' |

‎src/query/sql/src/planner/binder/bind_table_reference/bind_table.rs‎

Lines changed: 93 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -25,12 +25,15 @@ use databend_common_catalog::table::TimeNavigation;
2525
use databend_common_catalog::table_with_options::check_with_opt_valid;
2626
use databend_common_catalog::table_with_options::get_with_opt_consume;
2727
use databend_common_catalog::table_with_options::get_with_opt_max_batch_size;
28+
use databend_common_config::GlobalConfig;
2829
use databend_common_exception::ErrorCode;
2930
use databend_common_exception::Result;
3031
use databend_common_storages_basic::view_table::QUERY;
3132
use databend_storages_common_table_meta::table::get_change_type;
3233

3334
use crate::BindContext;
35+
use crate::LineageSourceRelation;
36+
use crate::ViewLineageSourceColumn;
3437
use crate::binder::Binder;
3538
use crate::binder::ViewIdent;
3639
use crate::binder::util::TableIdentifier;
@@ -189,17 +192,25 @@ impl Binder {
189192
{
190193
let change_type = get_change_type(&table_name_alias);
191194
if change_type.is_some() {
192-
let table_index = self.metadata.write().add_table(
193-
catalog,
194-
database.clone(),
195-
table_meta.clone(),
196-
branch_name,
197-
table_name_alias,
198-
!bind_context.binding_views.is_empty(),
199-
bind_context.planning_agg_index,
200-
false,
201-
cte_suffix_name,
202-
);
195+
let stream_lineage_source = stream_lineage_source_relation(&table_meta);
196+
let table_index = {
197+
let mut metadata = self.metadata.write();
198+
let table_index = metadata.add_table(
199+
catalog,
200+
database.clone(),
201+
table_meta.clone(),
202+
branch_name,
203+
table_name_alias,
204+
!bind_context.binding_views.is_empty(),
205+
bind_context.planning_agg_index,
206+
false,
207+
cte_suffix_name,
208+
);
209+
if let Some(stream_lineage_source) = stream_lineage_source {
210+
metadata.set_stream_lineage_source(table_index, stream_lineage_source);
211+
}
212+
table_index
213+
};
203214
let (s_expr, mut bind_context) = self.bind_base_table(
204215
bind_context,
205216
database.as_str(),
@@ -280,11 +291,11 @@ impl Binder {
280291
new_bind_context.binding_views.insert(view_ident);
281292
if let Statement::Query(query) = &stmt {
282293
self.metadata.write().add_table(
283-
catalog,
294+
catalog.clone(),
284295
database.clone(),
285-
table_meta,
286-
branch_name,
287-
table_name_alias,
296+
table_meta.clone(),
297+
branch_name.clone(),
298+
table_name_alias.clone(),
288299
false,
289300
false,
290301
false,
@@ -302,6 +313,13 @@ impl Binder {
302313
column.table_name = Some(self.normalize_identifier(table).name);
303314
}
304315
}
316+
self.add_view_lineage_source_columns(
317+
&new_bind_context,
318+
catalog.as_str(),
319+
database.as_str(),
320+
table_name.as_str(),
321+
&table_meta,
322+
);
305323
// Restore binding_views to the outer scope's value so the
306324
// current view does not leak into sibling/parent contexts.
307325
new_bind_context.binding_views = bind_context.binding_views.clone();
@@ -343,4 +361,64 @@ impl Binder {
343361
}
344362
}
345363
}
364+
365+
fn add_view_lineage_source_columns(
366+
&mut self,
367+
bind_context: &BindContext,
368+
catalog: &str,
369+
database: &str,
370+
view_name: &str,
371+
view: &std::sync::Arc<dyn databend_common_catalog::table::Table>,
372+
) {
373+
if !GlobalConfig::instance()
374+
.query
375+
.common
376+
.lineage
377+
.capture_enabled
378+
{
379+
return;
380+
}
381+
382+
let relation = LineageSourceRelation {
383+
catalog: catalog.to_string(),
384+
database: database.to_string(),
385+
name: view_name.to_string(),
386+
id: view.get_table_info().ident.table_id,
387+
};
388+
let schema = view.schema();
389+
let mut metadata = self.metadata.write();
390+
for (idx, column) in bind_context.columns.iter().enumerate() {
391+
let Some(field) = schema.fields().get(idx) else {
392+
continue;
393+
};
394+
metadata.add_view_lineage_source_column(column.index, ViewLineageSourceColumn {
395+
relation: relation.clone(),
396+
name: field.name().clone(),
397+
id: field.column_id,
398+
});
399+
}
400+
}
401+
}
402+
403+
fn stream_lineage_source_relation(
404+
table: &std::sync::Arc<dyn databend_common_catalog::table::Table>,
405+
) -> Option<LineageSourceRelation> {
406+
if !GlobalConfig::instance()
407+
.query
408+
.common
409+
.lineage
410+
.capture_enabled
411+
{
412+
return None;
413+
}
414+
415+
table.stream_source_table_info().and_then(|table_info| {
416+
let database = table_info.database_name().ok()?.to_string();
417+
Some(LineageSourceRelation {
418+
catalog: table_info.catalog().to_string(),
419+
database,
420+
name: table_info.name.clone(),
421+
id: table_info.ident.table_id,
422+
})
423+
})
346424
}

‎src/query/sql/src/planner/binder/ddl/view.rs‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,11 @@ use databend_common_ast::ast::DescribeViewStmt;
1818
use databend_common_ast::ast::DropViewStmt;
1919
use databend_common_ast::ast::ShowLimit;
2020
use databend_common_ast::ast::ShowViewsStmt;
21+
use databend_common_ast::ast::Statement;
2122
use databend_common_ast::ast::quote::QuotedIdent;
2223
use databend_common_ast::ast::quote::QuotedString;
2324
use databend_common_ast::visit::WalkMut;
25+
use databend_common_config::GlobalConfig;
2426
use databend_common_exception::Result;
2527
use databend_common_expression::DataField;
2628
use databend_common_expression::DataSchemaRefExt;
@@ -65,6 +67,16 @@ impl Binder {
6567
};
6668
query.walk_mut(&mut visitor)?;
6769
let subquery = format!("{}", query);
70+
let query_plan = if GlobalConfig::instance()
71+
.query
72+
.common
73+
.lineage
74+
.capture_enabled
75+
{
76+
Some(Box::new(self.view_query_plan(&query).await?))
77+
} else {
78+
None
79+
};
6880

6981
let plan = CreateViewPlan {
7082
create_option: create_option.clone().into(),
@@ -74,6 +86,7 @@ impl Binder {
7486
view_name,
7587
column_names,
7688
subquery,
89+
query_plan,
7790
};
7891
Ok(Plan::CreateView(plan.into()))
7992
}
@@ -225,4 +238,10 @@ impl Binder {
225238
schema,
226239
})))
227240
}
241+
242+
async fn view_query_plan(&mut self, query: &databend_common_ast::ast::Query) -> Result<Plan> {
243+
let stmt = Statement::Query(Box::new(query.clone()));
244+
let mut bind_context = BindContext::new();
245+
self.bind_statement(&mut bind_context, &stmt).await
246+
}
228247
}

‎src/query/sql/src/planner/binder/insert.rs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ use databend_common_ast::ast::InsertSource;
2020
use databend_common_ast::ast::InsertStmt;
2121
use databend_common_ast::ast::Statement;
2222
use databend_common_catalog::session_type::SessionType;
23+
use databend_common_config::GlobalConfig;
2324
use databend_common_exception::ErrorCode;
2425
use databend_common_exception::Result;
2526
use databend_common_expression::DataSchemaRef;
@@ -311,6 +312,12 @@ impl Binder {
311312
overwrite: *overwrite,
312313
source: input_source?,
313314
table_info: None,
315+
lineage_target_table_id: GlobalConfig::instance()
316+
.query
317+
.common
318+
.lineage
319+
.capture_enabled
320+
.then_some(table.get_table_info().ident.table_id),
314321
};
315322

316323
Ok(Plan::Insert(Box::new(plan)))

0 commit comments

Comments
 (0)