feat: add -c option to pass the SQL query directly as an argument on datafusion-cli (#6765)

* feat: add `-c option` to pass the SQL query directly as an argument on datafusion-cli

* feat: add integration tests of datafusion-cli

* chore: exclude test data files of datafusion-cli from RAT check

Author: r4ntix

datafusion-cli/Cargo.lock

@@ -43,3 +43,9 @@ parking_lot = { version = "0.12" }
 rustyline = "11.0"
 tokio = { version = "1.24", features = ["macros", "rt", "rt-multi-thread", "sync", "parking_lot"] }
 url = "2.2"
+
+[dev-dependencies]
+assert_cmd = "2.0"
+ctor = "0.2.0"
+predicates = "3.0"
+rstest = "0.17"

datafusion-cli/Cargo.toml

@@ -41,6 +41,20 @@ use std::time::Instant;
 use std::{fs::File, sync::Arc};
 use url::Url;
 
+/// run and execute SQL statements and commands, against a context with the given print options
+pub async fn exec_from_commands(
+    ctx: &mut SessionContext,
+    print_options: &PrintOptions,
+    commands: Vec<String>,
+) {
+    for sql in commands {
+        match exec_and_print(ctx, print_options, sql).await {
+            Ok(_) => {}
+            Err(err) => println!("{err}"),
+        }
+    }
+}
+
 /// run and execute SQL statements and commands from a file, against a context with the given print options
 pub async fn exec_from_lines(
     ctx: &mut SessionContext,
@@ -58,11 +72,8 @@ pub async fn exec_from_lines(
                 let line = line.trim_end();
                 query.push_str(line);
                 if line.ends_with(';') {
-                    match unescape_input(line) {
-                        Ok(sql) => match exec_and_print(ctx, print_options, sql).await {
-                            Ok(_) => {}
-                            Err(err) => eprintln!("{err}"),
-                        },
+                    match exec_and_print(ctx, print_options, query).await {
+                        Ok(_) => {}
                         Err(err) => eprintln!("{err}"),
                     }
                     query = "".to_owned();
@@ -149,11 +160,8 @@ pub async fn exec_from_repl(
             }
             Ok(line) => {
                 rl.add_history_entry(line.trim_end())?;
-                match unescape_input(&line) {
-                    Ok(sql) => match exec_and_print(ctx, &print_options, sql).await {
-                        Ok(_) => {}
-                        Err(err) => eprintln!("{err}"),
-                    },
+                match exec_and_print(ctx, &print_options, line).await {
+                    Ok(_) => {}
                     Err(err) => eprintln!("{err}"),
                 }
             }
@@ -182,6 +190,7 @@ async fn exec_and_print(
 ) -> Result<()> {
     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)) => {

datafusion-cli/src/exec.rs

@@ -44,13 +44,21 @@ struct Args {
     data_path: Option<String>,
 
     #[clap(
-        short = 'c',
+        short = 'b',
         long,
         help = "The batch size of each query, or use DataFusion default",
         validator(is_valid_batch_size)
     )]
     batch_size: Option<usize>,
 
+    #[clap(
+        short = 'c',
+        long,
+        multiple_values = true,
+        help = "Execute the given command string(s), then exit"
+    )]
+    command: Vec<String>,
+
     #[clap(
         short,
         long,
@@ -116,6 +124,7 @@ pub async fn main() -> Result<()> {
         quiet: args.quiet,
     };
 
+    let commands = args.command;
     let files = args.file;
     let rc = match args.rc {
         Some(file) => file,
@@ -132,18 +141,25 @@ pub async fn main() -> Result<()> {
         }
     };
 
-    if !files.is_empty() {
-        exec::exec_from_files(files, &mut ctx, &print_options).await;
-        Ok(())
-    } else {
+    if commands.is_empty() && files.is_empty() {
         if !rc.is_empty() {
             exec::exec_from_files(rc, &mut ctx, &print_options).await
         }
         // TODO maybe we can have thiserror for cli but for now let's keep it simple
-        exec::exec_from_repl(&mut ctx, &mut print_options)
+        return exec::exec_from_repl(&mut ctx, &mut print_options)
             .await
-            .map_err(|e| DataFusionError::External(Box::new(e)))
+            .map_err(|e| DataFusionError::External(Box::new(e)));
     }
+
+    if !files.is_empty() {
+        exec::exec_from_files(files, &mut ctx, &print_options).await;
+    }
+
+    if !commands.is_empty() {
+        exec::exec_from_commands(&mut ctx, &print_options, commands).await;
+    }
+
+    Ok(())
 }
 
 fn create_runtime_env() -> Result<RuntimeEnv> {