Skip to content

Commit c260b44

Browse files
YaroslavLitvinovYaroslavLitvinov
committed
add insert affected count test
1 parent 38667e1 commit c260b44

File tree

4 files changed

+182
-56
lines changed

4 files changed

+182
-56
lines changed

‎crates/runtime/src/tests/queries.rs‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18-
// use crate::tests::utils::macros::test_query;
18+
use crate::tests::utils::macros::test_query;
1919

2020
// // SELECT
2121
// test_query!(select_star, "SELECT * FROM employee_table");
@@ -507,3 +507,7 @@
507507
// type_map_contains_key,
508508
// "SELECT MAP_CONTAINS_KEY('key_to_find', {'my_key':'my_value'}::MAP(VARCHAR,VARCHAR));"
509509
// );
510+
test_query!(
511+
type_insert_get_count,
512+
"INSERT INTO insert_count(a) VALUES (1),(2),(3),(4),(5);"
513+
);

‎crates/runtime/src/tests/queries/table_setup.sql‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -153,4 +153,5 @@ CREATE OR REPLACE TABLE d2 (
153153
INSERT INTO d2 (id, value) VALUES
154154
(1,'xx'),
155155
(2,'yy'),
156-
(5,'zz');
156+
(5,'zz');
157+
CREATE TABLE insert_count(a INT);

‎crates/runtime/src/tests/snapshots/query_type_insert_get_count.snap‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
---
2+
source: crates/runtime/src/tests/queries.rs
3+
description: "\"INSERT INTO insert_count(a) VALUES (1),(2),(3),(4),(5);\""
4+
snapshot_kind: text
5+
---
6+
(
7+
Ok(
8+
Statement(
9+
Insert(
10+
Insert {
11+
or: None,
12+
ignore: false,
13+
into: true,
14+
table_name: ObjectName(
15+
[
16+
Ident {
17+
value: "insert_count",
18+
quote_style: None,
19+
span: Span(Location(1,13)..Location(1,25)),
20+
},
21+
],
22+
),
23+
table_alias: None,
24+
columns: [
25+
Ident {
26+
value: "a",
27+
quote_style: None,
28+
span: Span(Location(1,26)..Location(1,27)),
29+
},
30+
],
31+
overwrite: false,
32+
source: Some(
33+
Query {
34+
with: None,
35+
body: Values(
36+
Values {
37+
explicit_row: false,
38+
rows: [
39+
[
40+
Value(
41+
Number(
42+
"1",
43+
false,
44+
),
45+
),
46+
],
47+
[
48+
Value(
49+
Number(
50+
"2",
51+
false,
52+
),
53+
),
54+
],
55+
[
56+
Value(
57+
Number(
58+
"3",
59+
false,
60+
),
61+
),
62+
],
63+
[
64+
Value(
65+
Number(
66+
"4",
67+
false,
68+
),
69+
),
70+
],
71+
[
72+
Value(
73+
Number(
74+
"5",
75+
false,
76+
),
77+
),
78+
],
79+
],
80+
},
81+
),
82+
order_by: None,
83+
limit: None,
84+
limit_by: [],
85+
offset: None,
86+
fetch: None,
87+
locks: [],
88+
for_clause: None,
89+
settings: None,
90+
format_clause: None,
91+
},
92+
),
93+
partitioned: None,
94+
after_columns: [],
95+
table: false,
96+
on: None,
97+
returning: None,
98+
replace_into: false,
99+
priority: None,
100+
insert_alias: None,
101+
},
102+
),
103+
),
104+
),
105+
Ok(
106+
[
107+
"Dml: op=[Insert Into] table=[insert_count]",
108+
" Projection: CAST(column1 AS Int32) AS a",
109+
" Values: (Int64(1)), (Int64(2)), (Int64(3)), (Int64(4)), (Int64(5))",
110+
],
111+
),
112+
Ok(
113+
[
114+
"+-------+",
115+
"| count |",
116+
"+-------+",
117+
"| 5 |",
118+
"+-------+",
119+
],
120+
),
121+
)

‎crates/runtime/src/tests/utils.rs‎

Lines changed: 54 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -15,65 +15,65 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18-
// use crate::datafusion::functions::register_udfs;
19-
// use datafusion::prelude::{SessionConfig, SessionContext};
18+
use crate::datafusion::functions::register_udfs;
19+
use datafusion::prelude::{SessionConfig, SessionContext};
2020

21-
// static TABLE_SETUP: &str = include_str!(r"./queries/table_setup.sql");
21+
static TABLE_SETUP: &str = include_str!(r"./queries/table_setup.sql");
2222

23-
// #[allow(clippy::unwrap_used)]
24-
// pub async fn create_df_session() -> SessionContext {
25-
// let mut config = SessionConfig::new();
26-
// config.options_mut().catalog.information_schema = true;
27-
// let mut ctx = SessionContext::new_with_config(config);
23+
#[allow(clippy::unwrap_used)]
24+
pub async fn create_df_session() -> SessionContext {
25+
let mut config = SessionConfig::new();
26+
config.options_mut().catalog.information_schema = true;
27+
let mut ctx = SessionContext::new_with_config(config);
2828

29-
// register_udfs(&mut ctx).unwrap();
29+
register_udfs(&mut ctx).unwrap();
3030

31-
// for query in TABLE_SETUP.split(';') {
32-
// if !query.is_empty() {
33-
// dbg!("Running query: ", query);
34-
// ctx.sql(query).await.unwrap().collect().await.unwrap();
35-
// }
36-
// }
37-
// ctx
38-
// }
31+
for query in TABLE_SETUP.split(';') {
32+
if !query.is_empty() {
33+
dbg!("Running query: ", query);
34+
ctx.sql(query).await.unwrap().collect().await.unwrap();
35+
}
36+
}
37+
ctx
38+
}
3939

40-
// pub mod macros {
41-
// macro_rules! test_query {
42-
// ($test_fn_name:ident, $query:expr) => {
43-
// paste::paste! {
44-
// #[tokio::test]
45-
// async fn [< query_ $test_fn_name >]() {
46-
// let ctx = crate::tests::utils::create_df_session().await;
47-
// let statement = ctx.state().sql_to_statement($query, "snowflake");
40+
pub mod macros {
41+
macro_rules! test_query {
42+
($test_fn_name:ident, $query:expr) => {
43+
paste::paste! {
44+
#[tokio::test]
45+
async fn [< query_ $test_fn_name >]() {
46+
let ctx = crate::tests::utils::create_df_session().await;
47+
let statement = ctx.state().sql_to_statement($query, "snowflake");
4848

49-
// let plan = ctx.state().create_logical_plan($query)
50-
// .await;
49+
let plan = ctx.state().create_logical_plan($query)
50+
.await;
5151

52-
// let df = match &plan {
53-
// Ok(plan) => {
54-
// match ctx.execute_logical_plan(plan.clone()).await {
55-
// Ok(df) => {
56-
// let record_batches = df.collect().await.unwrap();
57-
// Ok(datafusion::arrow::util::pretty::pretty_format_batches(&record_batches).unwrap().to_string())
58-
// },
59-
// Err(e) => Err(e)
60-
// }
61-
// },
62-
// _ => Err(datafusion::error::DataFusionError::Execution("Failed to create logical plan".to_string()))
63-
// };
64-
// insta::with_settings!({
65-
// description => stringify!($query),
66-
// omit_expression => true,
67-
// prepend_module_to_snapshot => false
68-
// }, {
69-
// let plan = plan.map(|plan| plan.to_string().split("\n").map(|s| s.to_string()).collect::<Vec<String>>());
70-
// let df = df.map(|df| df.split("\n").map(|s| s.to_string()).collect::<Vec<String>>());
71-
// insta::assert_debug_snapshot!((statement, plan, df));
72-
// })
73-
// }
74-
// }
75-
// }
76-
// }
52+
let df = match &plan {
53+
Ok(plan) => {
54+
match ctx.execute_logical_plan(plan.clone()).await {
55+
Ok(df) => {
56+
let record_batches = df.collect().await.unwrap();
57+
Ok(datafusion::arrow::util::pretty::pretty_format_batches(&record_batches).unwrap().to_string())
58+
},
59+
Err(e) => Err(e)
60+
}
61+
},
62+
_ => Err(datafusion::error::DataFusionError::Execution("Failed to create logical plan".to_string()))
63+
};
64+
insta::with_settings!({
65+
description => stringify!($query),
66+
omit_expression => true,
67+
prepend_module_to_snapshot => false
68+
}, {
69+
let plan = plan.map(|plan| plan.to_string().split("\n").map(|s| s.to_string()).collect::<Vec<String>>());
70+
let df = df.map(|df| df.split("\n").map(|s| s.to_string()).collect::<Vec<String>>());
71+
insta::assert_debug_snapshot!((statement, plan, df));
72+
})
73+
}
74+
}
75+
}
76+
}
7777

78-
// pub(crate) use test_query;
79-
// }
78+
pub(crate) use test_query;
79+
}

0 commit comments

Comments
 (0)