[arrow-datafusion] branch master updated: Introduce `\i` command to execute from a file (#3136)

[agrove]

This is an automated email from the ASF dual-hosted git repository.

agrove pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/master by this push:
     new 2aa0a98d6 Introduce  `\i` command to execute from a file (#3136)
2aa0a98d6 is described below

commit 2aa0a98d6d95cc2641216ee3ccfe954e2b2eec57
Author: gorkem <yurtseven.gor...@gmail.com>
AuthorDate: Wed Aug 17 04:19:36 2022 +0300

    Introduce  `\i` command to execute from a file (#3136)
    
    * Introduce  command
    
    * Better error message
---
 datafusion-cli/src/command.rs | 28 +++++++++++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

diff --git a/datafusion-cli/src/command.rs b/datafusion-cli/src/command.rs
index aaebf64fb..f1b6f67e3 100644
--- a/datafusion-cli/src/command.rs
+++ b/datafusion-cli/src/command.rs
@@ -17,6 +17,7 @@
 
 //! Command within CLI
 
+use crate::exec::exec_from_lines;
 use crate::functions::{display_all_functions, Function};
 use crate::print_format::PrintFormat;
 use crate::print_options::PrintOptions;
@@ -26,6 +27,8 @@ use datafusion::arrow::datatypes::{DataType, Field, Schema};
 use datafusion::arrow::record_batch::RecordBatch;
 use datafusion::error::{DataFusionError, Result};
 use datafusion::prelude::SessionContext;
+use std::fs::File;
+use std::io::BufReader;
 use std::str::FromStr;
 use std::sync::Arc;
 use std::time::Instant;
@@ -38,6 +41,7 @@ pub enum Command {
     ListTables,
     DescribeTable(String),
     ListFunctions,
+    Include(Option<String>),
     SearchFunctions(String),
     QuietMode(Option<bool>),
     OutputFormat(Option<String>),
@@ -72,6 +76,22 @@ impl Command {
                     .print_batches(&batches, now)
                     .map_err(|e| DataFusionError::Execution(e.to_string()))
             }
+            Self::Include(filename) => {
+                if let Some(filename) = filename {
+                    let file = File::open(filename).map_err(|e| {
+                        DataFusionError::Execution(format!(
+                            "Error opening {:?} {}",
+                            filename, e
+                        ))
+                    })?;
+                    exec_from_lines(ctx, &mut BufReader::new(file), 
print_options).await;
+                    Ok(())
+                } else {
+                    Err(DataFusionError::Execution(
+                        "Required filename argument is missing".into(),
+                    ))
+                }
+            }
             Self::QuietMode(quiet) => {
                 if let Some(quiet) = quiet {
                     print_options.quiet = *quiet;
@@ -113,6 +133,9 @@ impl Command {
             Self::ListTables => ("\\d", "list tables"),
             Self::DescribeTable(_) => ("\\d name", "describe table"),
             Self::Help => ("\\?", "help"),
+            Self::Include(_) => {
+                ("\\i filename", "reads input from the specified filename")
+            }
             Self::ListFunctions => ("\\h", "function list"),
             Self::SearchFunctions(_) => ("\\h function", "search function"),
             Self::QuietMode(_) => ("\\quiet (true|false)?", "print or set 
quiet mode"),
@@ -123,11 +146,12 @@ impl Command {
     }
 }
 
-const ALL_COMMANDS: [Command; 8] = [
+const ALL_COMMANDS: [Command; 9] = [
     Command::ListTables,
     Command::DescribeTable(String::new()),
     Command::Quit,
     Command::Help,
+    Command::Include(Some(String::new())),
     Command::ListFunctions,
     Command::SearchFunctions(String::new()),
     Command::QuietMode(None),
@@ -169,6 +193,8 @@ impl FromStr for Command {
             ("?", None) => Self::Help,
             ("h", None) => Self::ListFunctions,
             ("h", Some(function)) => Self::SearchFunctions(function.into()),
+            ("i", None) => Self::Include(None),
+            ("i", Some(filename)) => Self::Include(Some(filename.to_owned())),
             ("quiet", Some("true" | "t" | "yes" | "y" | "on")) => {
                 Self::QuietMode(Some(true))
             }