Skip to content

Commit 9a9a3dd

Browse files
youngsofunyoungsofun
committed
test(query): cover get_lineage persisted metadata
1 parent 17f744d commit 9a9a3dd

3 files changed

Lines changed: 232 additions & 0 deletions

File tree

‎src/query/ee/tests/it/lineage/get_lineage.rs‎

Lines changed: 216 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,216 @@
1+
// Copyright 2023 Databend Cloud
2+
//
3+
// Licensed under the Elastic License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.elastic.co/licensing/elastic-license
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
use databend_common_exception::Result;
16+
use databend_common_expression::DataBlock;
17+
use databend_common_expression::ScalarRef;
18+
use databend_common_expression::types::number::NumberScalar;
19+
use databend_enterprise_query::test_kits::context::EESetup;
20+
use databend_query::interpreters::InterpreterFactory;
21+
use databend_query::interpreters::interpreter_plan_sql;
22+
use databend_query::test_kits::TestFixture;
23+
use futures_util::TryStreamExt;
24+
25+
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
26+
async fn test_get_lineage_reads_persisted_meta() -> anyhow::Result<()> {
27+
let mut setup = EESetup::new();
28+
setup.config_mut().query.common.lineage.capture_enabled = true;
29+
let fixture = TestFixture::setup_with_custom(setup).await?;
30+
31+
let prefix = fixture.default_table_name();
32+
let src = format!("{prefix}_src");
33+
let ctas = format!("{prefix}_ctas");
34+
let dst = format!("{prefix}_dst");
35+
let view = format!("{prefix}_view");
36+
let cross_db = format!("{prefix}_cross_db");
37+
let cross_ctas = format!("{prefix}_cross_ctas");
38+
39+
execute_sql(
40+
&fixture,
41+
&format!("CREATE TABLE default.{src}(a INT, b INT)"),
42+
)
43+
.await?;
44+
execute_sql(
45+
&fixture,
46+
&format!("INSERT INTO default.{src} VALUES (1, 10), (2, 20)"),
47+
)
48+
.await?;
49+
execute_sql(
50+
&fixture,
51+
&format!("CREATE TABLE default.{ctas} AS SELECT a AS x FROM default.{src} WHERE b > 0"),
52+
)
53+
.await?;
54+
execute_sql(&fixture, &format!("CREATE TABLE default.{dst}(x INT)")).await?;
55+
execute_sql(
56+
&fixture,
57+
&format!("INSERT INTO default.{dst} SELECT x FROM default.{ctas}"),
58+
)
59+
.await?;
60+
execute_sql(
61+
&fixture,
62+
&format!("CREATE VIEW default.{view} AS SELECT a AS va FROM default.{src}"),
63+
)
64+
.await?;
65+
execute_sql(&fixture, &format!("CREATE DATABASE {cross_db}")).await?;
66+
execute_sql(
67+
&fixture,
68+
&format!("CREATE TABLE {cross_db}.{src}(a INT, b INT)"),
69+
)
70+
.await?;
71+
execute_sql(
72+
&fixture,
73+
&format!("CREATE TABLE {cross_db}.{cross_ctas} AS SELECT a AS x FROM {cross_db}.{src}"),
74+
)
75+
.await?;
76+
77+
let table_upstream = query_rows(
78+
&fixture,
79+
&format!(
80+
"SELECT distance, source_object_domain, source_object_name, target_object_domain, target_object_name, target_status \
81+
FROM get_lineage('default.{ctas}', 'TABLE', 'UPSTREAM', 1)"
82+
),
83+
)
84+
.await?;
85+
assert!(table_upstream.contains(&vec![
86+
Some("1".to_string()),
87+
Some("TABLE".to_string()),
88+
Some(format!("default.{src}")),
89+
Some("TABLE".to_string()),
90+
Some(format!("default.{ctas}")),
91+
Some("ACTIVE".to_string()),
92+
]));
93+
94+
let column_upstream = query_rows(
95+
&fixture,
96+
&format!(
97+
"SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name \
98+
FROM get_lineage('default.{ctas}.x', 'COLUMN', 'UPSTREAM', 1)"
99+
),
100+
)
101+
.await?;
102+
assert!(column_upstream.contains(&vec![
103+
Some("1".to_string()),
104+
Some(format!("default.{src}")),
105+
Some("a".to_string()),
106+
Some(format!("default.{ctas}")),
107+
Some("x".to_string()),
108+
]));
109+
110+
let ctas_downstream = query_rows(
111+
&fixture,
112+
&format!(
113+
"SELECT distance, source_object_name, target_object_name \
114+
FROM get_lineage('default.{ctas}', 'TABLE', 'DOWNSTREAM', 1)"
115+
),
116+
)
117+
.await?;
118+
assert!(ctas_downstream.contains(&vec![
119+
Some("1".to_string()),
120+
Some(format!("default.{ctas}")),
121+
Some(format!("default.{dst}")),
122+
]));
123+
124+
let view_upstream = query_rows(
125+
&fixture,
126+
&format!(
127+
"SELECT distance, source_object_name, target_object_name \
128+
FROM get_lineage('default.{view}', 'VIEW', 'UPSTREAM', 1)"
129+
),
130+
)
131+
.await?;
132+
assert!(view_upstream.contains(&vec![
133+
Some("1".to_string()),
134+
Some(format!("default.{src}")),
135+
Some(format!("default.{view}")),
136+
]));
137+
138+
let view_column_downstream = query_rows(
139+
&fixture,
140+
&format!(
141+
"SELECT distance, source_object_name, source_column_name, target_object_name, target_column_name \
142+
FROM get_lineage('default.{src}.a', 'COLUMN', 'DOWNSTREAM', 1) \
143+
WHERE target_object_name = 'default.{view}'"
144+
),
145+
)
146+
.await?;
147+
assert!(view_column_downstream.contains(&vec![
148+
Some("1".to_string()),
149+
Some(format!("default.{src}")),
150+
Some("a".to_string()),
151+
Some(format!("default.{view}")),
152+
Some("va".to_string()),
153+
]));
154+
155+
let cross_database_upstream = query_rows(
156+
&fixture,
157+
&format!(
158+
"SELECT source_object_name, target_object_name \
159+
FROM get_lineage('{cross_db}.{cross_ctas}', 'TABLE', 'UPSTREAM', 1)"
160+
),
161+
)
162+
.await?;
163+
assert_eq!(cross_database_upstream, vec![vec![
164+
Some(format!("{cross_db}.{src}")),
165+
Some(format!("{cross_db}.{cross_ctas}")),
166+
]]);
167+
168+
Ok(())
169+
}
170+
171+
async fn execute_sql(fixture: &TestFixture, sql: &str) -> Result<()> {
172+
let stream = execute_query(fixture, sql).await?;
173+
stream.try_collect::<Vec<DataBlock>>().await?;
174+
Ok(())
175+
}
176+
177+
async fn query_rows(fixture: &TestFixture, sql: &str) -> Result<Vec<Vec<Option<String>>>> {
178+
let stream = execute_query(fixture, sql).await?;
179+
let blocks = stream.try_collect::<Vec<DataBlock>>().await?;
180+
let block = DataBlock::concat(&blocks)?;
181+
182+
let mut rows = Vec::with_capacity(block.num_rows());
183+
for row in 0..block.num_rows() {
184+
let mut values = Vec::with_capacity(block.num_columns());
185+
for col in 0..block.num_columns() {
186+
let scalar = block
187+
.get_by_offset(col)
188+
.index(row)
189+
.unwrap_or(ScalarRef::Null);
190+
values.push(scalar_to_string(scalar));
191+
}
192+
rows.push(values);
193+
}
194+
Ok(rows)
195+
}
196+
197+
fn scalar_to_string(scalar: ScalarRef<'_>) -> Option<String> {
198+
match scalar {
199+
ScalarRef::Null => None,
200+
ScalarRef::String(v) => Some(v.to_string()),
201+
ScalarRef::Number(NumberScalar::UInt64(v)) => Some(v.to_string()),
202+
ScalarRef::Number(NumberScalar::UInt32(v)) => Some(v.to_string()),
203+
ScalarRef::Number(NumberScalar::Int64(v)) => Some(v.to_string()),
204+
other => panic!("unexpected scalar in get_lineage test: {other:?}"),
205+
}
206+
}
207+
208+
async fn execute_query(
209+
fixture: &TestFixture,
210+
sql: &str,
211+
) -> Result<databend_common_expression::SendableDataBlockStream> {
212+
let ctx = fixture.new_query_ctx().await?;
213+
let (plan, _extras, _guard) = interpreter_plan_sql(ctx.clone(), sql, false, None).await?;
214+
let interpreter = InterpreterFactory::get(ctx.clone(), &plan).await?;
215+
interpreter.execute(ctx).await
216+
}

‎src/query/ee/tests/it/lineage/mod.rs‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
// Copyright 2023 Databend Cloud
2+
//
3+
// Licensed under the Elastic License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.elastic.co/licensing/elastic-license
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
mod get_lineage;

‎src/query/ee/tests/it/main.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,5 +14,6 @@
1414

1515
#![feature(unwrap_infallible)]
1616
mod license;
17+
mod lineage;
1718
mod storages;
1819
mod stream;

0 commit comments

Comments
 (0)