From 772856d7d5ff69f34c4a103f10aa34caf8d39909 Mon Sep 17 00:00:00 2001 From: NiwakaDev Date: Sun, 30 Jul 2023 22:48:50 +0900 Subject: [PATCH 1/4] feat: allow cli to accept multiple statements --- datafusion-cli/src/exec.rs | 24 ++++++++++++++---------- datafusion-cli/src/helper.rs | 3 --- 2 files changed, 14 insertions(+), 13 deletions(-) diff --git a/datafusion-cli/src/exec.rs b/datafusion-cli/src/exec.rs index 88553648ff5d..7129d517a0e4 100644 --- a/datafusion-cli/src/exec.rs +++ b/datafusion-cli/src/exec.rs @@ -26,6 +26,7 @@ use crate::{ }, print_options::PrintOptions, }; +use datafusion::sql::parser::DFParser; use datafusion::{ datasource::listing::ListingTableUrl, error::{DataFusionError, Result}, @@ -192,17 +193,20 @@ async fn exec_and_print( let now = Instant::now(); let sql = unescape_input(&sql)?; - let plan = ctx.state().create_logical_plan(&sql).await?; - let df = match &plan { - LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) => { - create_external_table(ctx, cmd).await?; - ctx.execute_logical_plan(plan).await? - } - _ => ctx.execute_logical_plan(plan).await?, - }; + let statements = DFParser::parse_sql(&sql)?; + for statement in statements { + let plan = ctx.state().statement_to_plan(statement).await?; + let df = match &plan { + LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) => { + create_external_table(ctx, cmd).await?; + ctx.execute_logical_plan(plan).await? + } + _ => ctx.execute_logical_plan(plan).await?, + }; - let results = df.collect().await?; - print_options.print_batches(&results, now)?; + let results = df.collect().await?; + print_options.print_batches(&results, now)?; + } Ok(()) } diff --git a/datafusion-cli/src/helper.rs b/datafusion-cli/src/helper.rs index 15464eec13a0..981c4b5aa3f3 100644 --- a/datafusion-cli/src/helper.rs +++ b/datafusion-cli/src/helper.rs @@ -54,9 +54,6 @@ impl CliHelper { Ok(statements) if statements.is_empty() => Ok(ValidationResult::Invalid( Some(" 🤔 You entered an empty statement".to_string()), )), - Ok(statements) if statements.len() > 1 => Ok(ValidationResult::Invalid( - Some(" 🤔 You entered more than one statement".to_string()), - )), Ok(_statements) => Ok(ValidationResult::Valid(None)), Err(err) => Ok(ValidationResult::Invalid(Some(format!( " 🤔 Invalid statement: {err}", From f3c1a328c1af57ce4d497372938fcb8ebaf0abff Mon Sep 17 00:00:00 2001 From: NiwakaDev Date: Tue, 1 Aug 2023 06:06:42 +0900 Subject: [PATCH 2/4] add empty line between results --- datafusion-cli/src/exec.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/datafusion-cli/src/exec.rs b/datafusion-cli/src/exec.rs index 7129d517a0e4..a1f504e510c1 100644 --- a/datafusion-cli/src/exec.rs +++ b/datafusion-cli/src/exec.rs @@ -206,6 +206,7 @@ async fn exec_and_print( let results = df.collect().await?; print_options.print_batches(&results, now)?; + println!(); } Ok(()) From 800f29646e6ec9f0666b55d841e49425df03c72d Mon Sep 17 00:00:00 2001 From: NiwakaDev Date: Tue, 1 Aug 2023 06:58:25 +0900 Subject: [PATCH 3/4] use current context dialect --- datafusion-cli/src/exec.rs | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/datafusion-cli/src/exec.rs b/datafusion-cli/src/exec.rs index a1f504e510c1..c8e32c11f6a5 100644 --- a/datafusion-cli/src/exec.rs +++ b/datafusion-cli/src/exec.rs @@ -26,7 +26,7 @@ use crate::{ }, print_options::PrintOptions, }; -use datafusion::sql::parser::DFParser; +use datafusion::sql::{parser::DFParser, sqlparser::dialect::dialect_from_str}; use datafusion::{ datasource::listing::ListingTableUrl, error::{DataFusionError, Result}, @@ -193,7 +193,16 @@ async fn exec_and_print( let now = Instant::now(); let sql = unescape_input(&sql)?; - let statements = DFParser::parse_sql(&sql)?; + let task_ctx = ctx.task_ctx(); + let dialect = &task_ctx.session_config().options().sql_parser.dialect; + let dialect = dialect_from_str(dialect).ok_or_else(|| { + DataFusionError::Plan(format!( + "Unsupported SQL dialect: {dialect}. Available dialects: \ + Generic, MySQL, PostgreSQL, Hive, SQLite, Snowflake, Redshift, \ + MsSQL, ClickHouse, BigQuery, Ansi." + )) + })?; + let statements = DFParser::parse_sql_with_dialect(&sql, dialect.as_ref())?; for statement in statements { let plan = ctx.state().statement_to_plan(statement).await?; let df = match &plan { From c7f3f01bf4f4482a2e537f0b6621fb590f5b69d8 Mon Sep 17 00:00:00 2001 From: NiwakaDev Date: Wed, 2 Aug 2023 07:55:01 +0900 Subject: [PATCH 4/4] move empty line to print_timing_info --- datafusion-cli/src/exec.rs | 2 -- datafusion-cli/src/print_options.rs | 2 +- datafusion-cli/tests/cli_integration.rs | 4 ++++ 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/datafusion-cli/src/exec.rs b/datafusion-cli/src/exec.rs index c8e32c11f6a5..adbdb06e5299 100644 --- a/datafusion-cli/src/exec.rs +++ b/datafusion-cli/src/exec.rs @@ -215,9 +215,7 @@ async fn exec_and_print( let results = df.collect().await?; print_options.print_batches(&results, now)?; - println!(); } - Ok(()) } diff --git a/datafusion-cli/src/print_options.rs b/datafusion-cli/src/print_options.rs index 5e3792634a4e..33ba7ef0863e 100644 --- a/datafusion-cli/src/print_options.rs +++ b/datafusion-cli/src/print_options.rs @@ -28,7 +28,7 @@ pub struct PrintOptions { fn print_timing_info(row_count: usize, now: Instant) { println!( - "{} {} in set. Query took {:.3} seconds.", + "{} {} in set. Query took {:.3} seconds.\n", row_count, if row_count == 1 { "row" } else { "rows" }, now.elapsed().as_secs_f64() diff --git a/datafusion-cli/tests/cli_integration.rs b/datafusion-cli/tests/cli_integration.rs index c6bee274e93c..28344ffa94f8 100644 --- a/datafusion-cli/tests/cli_integration.rs +++ b/datafusion-cli/tests/cli_integration.rs @@ -33,6 +33,10 @@ fn init() { ["--command", "select 1", "--format", "json", "-q"], "[{\"Int64(1)\":1}]\n" )] +#[case::exec_multiple_statements( + ["--command", "select 1; select 2;", "--format", "json", "-q"], + "[{\"Int64(1)\":1}]\n[{\"Int64(2)\":2}]\n" +)] #[case::exec_from_files( ["--file", "tests/data/sql.txt", "--format", "json", "-q"], "[{\"Int64(1)\":1}]\n"