This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 1901d6fd41 feat: allow `datafusion-cli` to accept multiple statements
(#7138)
1901d6fd41 is described below
commit 1901d6fd41b85656ff334b515ef5afc00ddb7355
Author: Niwaka <[email protected]>
AuthorDate: Thu Aug 3 21:44:15 2023 +0900
feat: allow `datafusion-cli` to accept multiple statements (#7138)
* feat: allow cli to accept multiple statements
* add empty line between results
* use current context dialect
* move empty line to print_timing_info
---
datafusion-cli/src/exec.rs | 34 ++++++++++++++++++++++-----------
datafusion-cli/src/helper.rs | 3 ---
datafusion-cli/src/print_options.rs | 2 +-
datafusion-cli/tests/cli_integration.rs | 4 ++++
4 files changed, 28 insertions(+), 15 deletions(-)
diff --git a/datafusion-cli/src/exec.rs b/datafusion-cli/src/exec.rs
index 88553648ff..adbdb06e52 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, sqlparser::dialect::dialect_from_str};
use datafusion::{
datasource::listing::ListingTableUrl,
error::{DataFusionError, Result},
@@ -192,18 +193,29 @@ 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 results = df.collect().await?;
- print_options.print_batches(&results, now)?;
+ 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 {
+ 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)?;
+ }
Ok(())
}
diff --git a/datafusion-cli/src/helper.rs b/datafusion-cli/src/helper.rs
index 15464eec13..981c4b5aa3 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}",
diff --git a/datafusion-cli/src/print_options.rs
b/datafusion-cli/src/print_options.rs
index 5e3792634a..33ba7ef086 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 c6bee274e9..28344ffa94 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"
