apache · jimexist · Nov 3, 2021 · Nov 3, 2021
diff --git a/datafusion-cli/src/command.rs b/datafusion-cli/src/command.rs
@@ -0,0 +1,85 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Command within CLI
+
+use datafusion::arrow::array::{ArrayRef, StringArray};
+use datafusion::arrow::datatypes::{DataType, Field, Schema};
+use datafusion::arrow::record_batch::RecordBatch;
+use datafusion::arrow::util::pretty;
+use datafusion::error::{DataFusionError, Result};
+use std::str::FromStr;
+use std::sync::Arc;
+
+/// Command
+#[derive(Debug)]
+pub enum Command {
+    Quit,
+    Help,
+}
+
+impl Command {
+    pub fn execute(&self) -> Result<()> {
+        match self {
+            Self::Help => pretty::print_batches(&[all_commands_info()])
+                .map_err(|e| DataFusionError::Execution(e.to_string())),
+            Self::Quit => Err(DataFusionError::Execution(
+                "Unexpected quit, this should be handled outside".into(),
+            )),
+        }
+    }
+
+    fn get_name_and_description(&self) -> (&str, &str) {
+        match self {
+            Self::Quit => ("\\q", "quit datafusion-cli"),
+            Self::Help => ("\\?", "help"),
+        }
+    }
+}
+
+const ALL_COMMANDS: [Command; 2] = [Command::Quit, Command::Help];
+
+fn all_commands_info() -> RecordBatch {
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("Command", DataType::Utf8, false),
+        Field::new("Description", DataType::Utf8, false),
+    ]));
+    let (names, description): (Vec<&str>, Vec<&str>) = ALL_COMMANDS
+        .iter()
+        .map(|c| c.get_name_and_description())
+        .unzip();
+    RecordBatch::try_new(
+        schema,
+        [names, description]
+            .into_iter()
+            .map(|i| Arc::new(StringArray::from(i)) as ArrayRef)
+            .collect::<Vec<_>>(),
+    )
+    .expect("This should not fail")
+}
+
+impl FromStr for Command {
+    type Err = ();
+
+    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
+        match s {
+            "q" => Ok(Self::Quit),
+            "?" => Ok(Self::Help),
+            _ => Err(()),
+        }
+    }
+}
diff --git a/datafusion-cli/src/exec.rs b/datafusion-cli/src/exec.rs
@@ -18,15 +18,20 @@
 //! Execution functions
 
 use crate::{
+    command::Command,
     context::Context,
     print_format::{all_print_formats, PrintFormat},
     print_options::PrintOptions,
 };
+use datafusion::arrow::record_batch::RecordBatch;
+use datafusion::arrow::util::pretty;
 use datafusion::error::{DataFusionError, Result};
 use rustyline::Editor;
 use std::fs::File;
 use std::io::prelude::*;
 use std::io::BufReader;
+use std::str::FromStr;
+use std::sync::Arc;
 use std::time::Instant;
 
 /// run and execute SQL statements and commands from a file, against a context with the given print options
@@ -78,23 +83,34 @@ pub async fn exec_from_repl(ctx: &mut Context, print_options: PrintOptions) {
     let mut query = "".to_owned();
     loop {
         match rl.readline("> ") {
-            Ok(ref line) if is_exit_command(line) && query.is_empty() => {
-                break;
+            Ok(line) if line.starts_with('\\') => {
+                if let Ok(cmd) = &line[1..].parse::<Command>() {
+                    match cmd {
+                        Command::Quit => break,
+                        others => {
+                            if let Err(e) = others.execute() {
+                                eprintln!("{}", e)
+                            }
+                        }
+                    }
+                } else {
+                    eprintln!("'\\{}' is not a valid command", &line[1..]);
+                }
             }
-            Ok(ref line) if line.starts_with("--") => {
+            Ok(line) if line.starts_with("--") => {
                 continue;
             }
-            Ok(ref line) if line.trim_end().ends_with(';') => {
+            Ok(line) if line.trim_end().ends_with(';') => {
                 query.push_str(line.trim_end());
                 rl.add_history_entry(query.clone());
                 match exec_and_print(ctx, print_options.clone(), query).await {
                     Ok(_) => {}
-                    Err(err) => println!("{:?}", err),
+                    Err(err) => eprintln!("{:?}", err),
                 }
                 query = "".to_owned();
             }
-            Ok(ref line) => {
-                query.push_str(line);
+            Ok(line) => {
+                query.push_str(&line);
                 query.push('\n');
             }
             Err(_) => {
@@ -118,8 +134,3 @@ async fn exec_and_print(
 
     Ok(())
 }
-
-fn is_exit_command(line: &str) -> bool {
-    let line = line.trim_end().to_lowercase();
-    line == "quit" || line == "exit"
-}
diff --git a/datafusion-cli/src/lib.rs b/datafusion-cli/src/lib.rs
@@ -19,6 +19,7 @@
 #![allow(unused_imports)]
 pub const DATAFUSION_CLI_VERSION: &str = env!("CARGO_PKG_VERSION");
 
+pub mod command;
 pub mod context;
 pub mod exec;
 pub mod print_format;