add command to list columns

datafusion-cli/src/command.rs

@@ -18,31 +18,47 @@
 //! Command within CLI
 
 use crate::context::Context;
+use crate::print_options::PrintOptions;
 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;
+use std::time::Instant;
 
 /// Command
 #[derive(Debug)]
 pub enum Command {
     Quit,
     Help,
     ListTables,
+    DescribeTable(String),
 }
 
 impl Command {
-    pub async fn execute(&self, ctx: &mut Context) -> Result<()> {
+    pub async fn execute(
+        &self,
+        ctx: &mut Context,
+        print_options: &PrintOptions,
+    ) -> Result<()> {
+        let now = Instant::now();
         match self {
-            Self::Help => pretty::print_batches(&[all_commands_info()])
+            Self::Help => print_options
+                .print_batches(&[all_commands_info()], now)
                 .map_err(|e| DataFusionError::Execution(e.to_string())),
             Self::ListTables => {
                 let df = ctx.sql("SHOW TABLES").await?;
                 let batches = df.collect().await?;
-                pretty::print_batches(&batches)
+                print_options
+                    .print_batches(&batches, now)
+                    .map_err(|e| DataFusionError::Execution(e.to_string()))
+            }
+            Self::DescribeTable(name) => {
+                let df = ctx.sql(&format!("SHOW COLUMNS FROM {}", name)).await?;
+                let batches = df.collect().await?;
+                print_options
+                    .print_batches(&batches, now)
                     .map_err(|e| DataFusionError::Execution(e.to_string()))
             }
             Self::Quit => Err(DataFusionError::Execution(
@@ -51,24 +67,30 @@ impl Command {
         }
     }
 
-    fn get_name_and_description(&self) -> (&str, &str) {
+    fn get_name_and_description(&self) -> (&'static str, &'static str) {
         match self {
             Self::Quit => ("\\q", "quit datafusion-cli"),
             Self::ListTables => ("\\d", "list tables"),
+            Self::DescribeTable(_) => ("\\d name", "describe table"),
             Self::Help => ("\\?", "help"),
         }
     }
 }
 
-const ALL_COMMANDS: [Command; 3] = [Command::ListTables, Command::Quit, Command::Help];
+const ALL_COMMANDS: [Command; 4] = [
+    Command::ListTables,
+    Command::DescribeTable(String::new()),
+    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()
+        .into_iter()
         .map(|c| c.get_name_and_description())
         .unzip();
     RecordBatch::try_new(
@@ -85,10 +107,16 @@ impl FromStr for Command {
     type Err = ();
 
     fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
-        Ok(match s {
-            "q" => Self::Quit,
-            "d" => Self::ListTables,
-            "?" => Self::Help,
+        let (c, arg) = if let Some((a, b)) = s.split_once(' ') {
+            (a, Some(b))
+        } else {
+            (s, None)
+        };
+        Ok(match (c, arg) {
+            ("q", None) => Self::Quit,
+            ("d", None) => Self::ListTables,
+            ("d", Some(name)) => Self::DescribeTable(name.into()),
+            ("?", None) => Self::Help,
             _ => return Err(()),
         })
     }

datafusion-cli/src/exec.rs

@@ -38,7 +38,7 @@ use std::time::Instant;
 pub async fn exec_from_lines(
     ctx: &mut Context,
     reader: &mut BufReader<File>,
-    print_options: PrintOptions,
+    print_options: &PrintOptions,
 ) {
     let mut query = "".to_owned();
 
@@ -51,7 +51,7 @@ pub async fn exec_from_lines(
                 let line = line.trim_end();
                 query.push_str(line);
                 if line.ends_with(';') {
-                    match exec_and_print(ctx, print_options.clone(), query).await {
+                    match exec_and_print(ctx, print_options, query).await {
                         Ok(_) => {}
                         Err(err) => println!("{:?}", err),
                     }
@@ -76,19 +76,20 @@ pub async fn exec_from_lines(
 }
 
 /// run and execute SQL statements and commands against a context with the given print options
-pub async fn exec_from_repl(ctx: &mut Context, print_options: PrintOptions) {
+pub async fn exec_from_repl(ctx: &mut Context, print_options: &PrintOptions) {
     let mut rl = Editor::<()>::new();
     rl.load_history(".history").ok();
 
     let mut query = "".to_owned();
     loop {
         match rl.readline("> ") {
             Ok(line) if line.starts_with('\\') => {
+                rl.add_history_entry(line.trim_end());
                 if let Ok(cmd) = &line[1..].parse::<Command>() {
                     match cmd {
                         Command::Quit => break,
-                        others => {
-                            if let Err(e) = others.execute(ctx).await {
+                        _ => {
+                            if let Err(e) = cmd.execute(ctx, print_options).await {
                                 eprintln!("{}", e)
                             }
                         }
@@ -103,7 +104,7 @@ pub async fn exec_from_repl(ctx: &mut Context, print_options: PrintOptions) {
             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 {
+                match exec_and_print(ctx, print_options, query).await {
                     Ok(_) => {}
                     Err(err) => eprintln!("{:?}", err),
                 }
@@ -124,7 +125,7 @@ pub async fn exec_from_repl(ctx: &mut Context, print_options: PrintOptions) {
 
 async fn exec_and_print(
     ctx: &mut Context,
-    print_options: PrintOptions,
+    print_options: &PrintOptions,
     sql: String,
 ) -> Result<()> {
     let now = Instant::now();

datafusion-cli/src/main.rs

@@ -145,10 +145,10 @@ pub async fn main() -> Result<()> {
             .collect::<Vec<_>>();
         for file in files {
             let mut reader = BufReader::new(file);
-            exec::exec_from_lines(&mut ctx, &mut reader, print_options.clone()).await;
+            exec::exec_from_lines(&mut ctx, &mut reader, &print_options).await;
         }
     } else {
-        exec::exec_from_repl(&mut ctx, print_options).await;
+        exec::exec_from_repl(&mut ctx, &print_options).await;
     }
 
     Ok(())