From 0ccdcc6bb340f22880ee67b9d3f7ebafd3bad0df Mon Sep 17 00:00:00 2001 From: Jiayu Liu Date: Wed, 3 Nov 2021 08:55:36 +0800 Subject: [PATCH] add \q as quit command and add \? for help --- datafusion-cli/src/command.rs | 85 +++++++++++++++++++++++++++++++++++ datafusion-cli/src/exec.rs | 35 ++++++++++----- datafusion-cli/src/lib.rs | 1 + 3 files changed, 109 insertions(+), 12 deletions(-) create mode 100644 datafusion-cli/src/command.rs diff --git a/datafusion-cli/src/command.rs b/datafusion-cli/src/command.rs new file mode 100644 index 0000000000000..f99130532cbd8 --- /dev/null +++ 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::>(), + ) + .expect("This should not fail") +} + +impl FromStr for Command { + type Err = (); + + fn from_str(s: &str) -> std::result::Result { + match s { + "q" => Ok(Self::Quit), + "?" => Ok(Self::Help), + _ => Err(()), + } + } +} diff --git a/datafusion-cli/src/exec.rs b/datafusion-cli/src/exec.rs index ed3ed98bcf933..979a222601937 100644 --- 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::() { + 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 index a116c35644ef8..89776f09baa4c 100644 --- 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;