feat(cli): support multi line and output format & add test for ontime…

…200 (#83)

* fix(driver): message for unimplemented in FlightSQL

* feat(cli): support multi line & add test for ontime200

* fix(cli): format progress message

* fix: output csv for ontime

* feat(cli): support multiple output

cli/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "bendsql"
-version = "0.2.7"
+version = "0.3.0"
 edition = "2021"
 license = "Apache-2.0"
 description = "Databend Native Command Line Tool"
@@ -16,7 +16,7 @@ chrono = { version = "0.4.24", default-features = false, features = ["clock"] }
 clap = { version = "4.1.0", features = ["derive", "env"] }
 comfy-table = "6.1.4"
 csv = "1.2.1"
-databend-driver = { path = "../driver", version = "0.2.17", features = ["rustls", "flight-sql"] }
+databend-driver = { path = "../driver", version = "0.2.18", features = ["rustls", "flight-sql"] }
 futures = { version = "0.3", default-features = false, features = ["alloc"] }
 humantime-serde = "1.1.1"
 indicatif = "0.17.3"

cli/src/config.rs

@@ -16,6 +16,7 @@
 
 use std::{collections::BTreeMap, path::Path};
 
+use clap::ValueEnum;
 use serde::Deserialize;
 
 #[derive(Clone, Debug, Default, Deserialize)]
@@ -32,6 +33,17 @@ pub struct Settings {
     pub display_pretty_sql: bool,
     pub prompt: String,
     pub progress_color: String,
+    pub output_format: OutputFormat,
+    /// Show progress [bar] when executing queries.
+    /// Only works in non-interactive mode.
+    pub show_progress: bool,
+}
+
+#[derive(ValueEnum, Clone, Debug, PartialEq, Deserialize)]
+pub enum OutputFormat {
+    Table,
+    CSV,
+    TSV,
 }
 
 #[derive(Clone, Debug, Deserialize)]
@@ -72,6 +84,8 @@ impl Default for Settings {
             display_pretty_sql: true,
             progress_color: "cyan".to_string(),
             prompt: "bendsql> ".to_string(),
+            output_format: OutputFormat::Table,
+            show_progress: false,
         }
     }
 }

cli/src/display.rs

@@ -25,7 +25,11 @@ use tokio::time::Instant;
 
 use indicatif::{HumanBytes, ProgressBar, ProgressState, ProgressStyle};
 
-use crate::{ast::format_query, config::Settings, helper::CliHelper};
+use crate::{
+    ast::format_query,
+    config::{OutputFormat, Settings},
+    helper::CliHelper,
+};
 
 #[async_trait::async_trait]
 pub trait ChunkDisplay {
@@ -97,9 +101,10 @@ impl<'a> ChunkDisplay for ReplDisplay<'a> {
         if let Some(pb) = self.progress.take() {
             pb.finish_and_clear();
         }
-        print_rows(self.schema.clone(), &rows, self.settings)?;
-
-        println!();
+        if !rows.is_empty() {
+            println!("{}", create_table(self.schema.clone(), &rows)?);
+            println!();
+        }
 
         let rows_str = if self.rows > 1 { "rows" } else { "row" };
         println!(
@@ -123,8 +128,8 @@ impl<'a> ChunkDisplay for ReplDisplay<'a> {
 }
 
 pub struct FormatDisplay<'a> {
-    _settings: &'a Settings,
-    _schema: SchemaRef,
+    settings: &'a Settings,
+    schema: SchemaRef,
     data: RowProgressIterator,
 
     rows: usize,
@@ -140,8 +145,8 @@ impl<'a> FormatDisplay<'a> {
         data: RowProgressIterator,
     ) -> Self {
         Self {
-            _settings: settings,
-            _schema: schema,
+            settings,
+            schema,
             data,
             rows: 0,
             _progress: None,
@@ -167,13 +172,32 @@ impl<'a> ChunkDisplay for FormatDisplay<'a> {
                 }
             }
         }
-        let mut wtr = csv::WriterBuilder::new()
-            .delimiter(b'\t')
-            .quote_style(csv::QuoteStyle::Necessary)
-            .from_writer(std::io::stdout());
-        for row in rows {
-            let values: Vec<String> = row.values().iter().map(|v| v.to_string()).collect();
-            wtr.write_record(values)?;
+        if rows.is_empty() {
+            return Ok(());
+        }
+        match self.settings.output_format {
+            OutputFormat::Table => {
+                println!("{}", create_table(self.schema.clone(), &rows)?);
+            }
+            OutputFormat::CSV => {
+                let mut wtr = csv::WriterBuilder::new()
+                    .quote_style(csv::QuoteStyle::Necessary)
+                    .from_writer(std::io::stdout());
+                for row in rows {
+                    let record = row.into_iter().map(|v| v.to_string()).collect::<Vec<_>>();
+                    wtr.write_record(record)?;
+                }
+            }
+            OutputFormat::TSV => {
+                let mut wtr = csv::WriterBuilder::new()
+                    .delimiter(b'\t')
+                    .quote_style(csv::QuoteStyle::Necessary)
+                    .from_writer(std::io::stdout());
+                for row in rows {
+                    let record = row.into_iter().map(|v| v.to_string()).collect::<Vec<_>>();
+                    wtr.write_record(record)?;
+                }
+            }
         }
         Ok(())
     }
@@ -183,36 +207,26 @@ impl<'a> ChunkDisplay for FormatDisplay<'a> {
     }
 }
 
-fn print_rows(schema: SchemaRef, results: &[Row], _settings: &Settings) -> Result<()> {
-    if !results.is_empty() {
-        println!("{}", create_table(schema, results)?);
-    }
-    Ok(())
-}
-
 fn format_read_progress(progress: &QueryProgress, elapsed: f64) -> String {
-    let mut s = String::new();
-    s.push_str(&format!(
-        "{} rows, {} processed, ({} rows/s, {}/s)",
+    format!(
+        "Processing {}/{} ({} rows/s), {}/{} ({}/s)",
+        humanize_count(progress.read_rows as f64),
         humanize_count(progress.total_rows as f64),
+        humanize_count(progress.read_rows as f64 / elapsed),
+        HumanBytes(progress.read_bytes as u64),
         HumanBytes(progress.total_bytes as u64),
-        humanize_count(progress.total_rows as f64 / elapsed),
-        HumanBytes((progress.total_bytes as f64 / elapsed) as u64)
-    ));
-    s
+        HumanBytes((progress.read_bytes as f64 / elapsed) as u64)
+    )
 }
 
-// TODO:(everpcpc)
-fn _format_write_progress(progress: &QueryProgress, elapsed: f64) -> String {
-    let mut s = String::new();
-    s.push_str(&format!(
-        "{} rows, {} written, ({} rows/s., {}/s.)",
+pub fn format_write_progress(progress: &QueryProgress, elapsed: f64) -> String {
+    format!(
+        "Written {} ({} rows/s), {} ({}/s)",
         humanize_count(progress.write_rows as f64),
-        HumanBytes(progress.write_bytes as u64),
         humanize_count(progress.write_rows as f64 / elapsed),
+        HumanBytes(progress.write_bytes as u64),
         HumanBytes((progress.write_bytes as f64 / elapsed) as u64)
-    ));
-    s
+    )
 }
 
 fn display_read_progress(pb: Option<ProgressBar>, current: &QueryProgress) -> ProgressBar {

cli/src/main.rs

@@ -24,12 +24,12 @@ use std::collections::BTreeMap;
 
 use anyhow::{anyhow, Result};
 use clap::{CommandFactory, Parser, ValueEnum};
-use config::Config;
+use config::{Config, OutputFormat};
 
 /// Supported file format and options:
 /// https://databend.rs/doc/sql-reference/file-format-options
 #[derive(ValueEnum, Clone, Debug, PartialEq)]
-enum InputFormat {
+pub enum InputFormat {
     CSV,
     TSV,
     NDJSON,
@@ -71,13 +71,6 @@ impl InputFormat {
     }
 }
 
-#[derive(ValueEnum, Clone, Debug, PartialEq)]
-enum OutputFormat {
-    Table,
-    CSV,
-    TSV,
-}
-
 #[derive(Debug, Parser, PartialEq)]
 // disable default help flag since it would conflict with --host
 #[command(author, version, about, disable_help_flag = true)]
@@ -235,6 +228,8 @@ pub async fn main() -> Result<()> {
             conn_args.get_dsn()?
         }
     };
+    config.settings.output_format = args.output;
+    config.settings.show_progress = args.progress;
 
     let is_repl = atty::is(atty::Stream::Stdin) && !args.non_interactive && args.query.is_none();
     let mut session = session::Session::try_new(dsn, config.settings, is_repl).await?;