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

Author: NiwakaDev

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(())
 }
 

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}",

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

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"