Skip to content

Commit d3bc3c8

Browse files
youngsofunyoungsofun
committed
feat(query): extract query lineage
1 parent 7117b1f commit d3bc3c8

17 files changed

Lines changed: 2384 additions & 32 deletions

File tree

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

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

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

112+
/// Returns the underlying table that should be treated as this table's
113+
/// lineage source. Virtual wrappers such as streams may override this.
114+
fn lineage_source_table_info(&self) -> Option<&TableInfo> {
115+
None
116+
}
117+
112118
fn get_data_source_info(&self) -> DataSourceInfo {
113119
DataSourceInfo::TableSource(self.get_table_info().clone())
114120
}

‎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: 113 additions & 26 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::LineageSourceColumn;
36+
use crate::LineageSourceRelation;
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 lineage_source = 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(lineage_source) = lineage_source {
210+
metadata.set_table_lineage_source_relation(table_index, 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();
@@ -315,17 +333,25 @@ impl Binder {
315333
}
316334
}
317335
_ => {
318-
let table_index = self.metadata.write().add_table(
319-
catalog.clone(),
320-
database.clone(),
321-
table_meta.clone(),
322-
branch_name,
323-
table_name_alias,
324-
!bind_context.binding_views.is_empty(),
325-
bind_context.planning_agg_index,
326-
false,
327-
cte_suffix_name,
328-
);
336+
let lineage_source = lineage_source_relation(&table_meta);
337+
let table_index = {
338+
let mut metadata = self.metadata.write();
339+
let table_index = metadata.add_table(
340+
catalog.clone(),
341+
database.clone(),
342+
table_meta.clone(),
343+
branch_name,
344+
table_name_alias,
345+
!bind_context.binding_views.is_empty(),
346+
bind_context.planning_agg_index,
347+
false,
348+
cte_suffix_name,
349+
);
350+
if let Some(lineage_source) = lineage_source {
351+
metadata.set_table_lineage_source_relation(table_index, lineage_source);
352+
}
353+
table_index
354+
};
329355

330356
let (s_expr, mut bind_context) = self.bind_base_table(
331357
bind_context,
@@ -343,4 +369,65 @@ impl Binder {
343369
}
344370
}
345371
}
372+
373+
fn add_view_lineage_source_columns(
374+
&mut self,
375+
bind_context: &BindContext,
376+
catalog: &str,
377+
database: &str,
378+
view_name: &str,
379+
view: &std::sync::Arc<dyn databend_common_catalog::table::Table>,
380+
) {
381+
if !GlobalConfig::instance()
382+
.query
383+
.common
384+
.lineage
385+
.capture_enabled
386+
{
387+
return;
388+
}
389+
390+
let relation = LineageSourceRelation {
391+
catalog: catalog.to_string(),
392+
database: database.to_string(),
393+
name: view_name.to_string(),
394+
id: view.get_table_info().ident.table_id,
395+
};
396+
let schema = view.schema();
397+
let mut metadata = self.metadata.write();
398+
for (idx, column) in bind_context.columns.iter().enumerate() {
399+
let Some(field) = schema.fields().get(idx) else {
400+
continue;
401+
};
402+
metadata.add_lineage_source_column(column.index, LineageSourceColumn {
403+
relation: relation.clone(),
404+
relation_is_view: true,
405+
name: field.name().clone(),
406+
id: field.column_id,
407+
});
408+
}
409+
}
410+
}
411+
412+
fn lineage_source_relation(
413+
table: &std::sync::Arc<dyn databend_common_catalog::table::Table>,
414+
) -> Option<LineageSourceRelation> {
415+
if !GlobalConfig::instance()
416+
.query
417+
.common
418+
.lineage
419+
.capture_enabled
420+
{
421+
return None;
422+
}
423+
424+
table.lineage_source_table_info().and_then(|table_info| {
425+
let database = table_info.database_name().ok()?.to_string();
426+
Some(LineageSourceRelation {
427+
catalog: table_info.catalog().to_string(),
428+
database,
429+
name: table_info.name.clone(),
430+
id: table_info.ident.table_id,
431+
})
432+
})
346433
}

‎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)