feat(cli): add new param quote_style (#277)

databendlabs · Oct 26, 2023 · 0d0dd21 · 0d0dd21
1 parent a3daa57
commit 0d0dd21
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 2 deletions.
diff --git a/cli/src/config.rs b/cli/src/config.rs
@@ -83,6 +83,8 @@ pub struct Settings {
     pub max_width: usize,
     /// Output format is set by the flag.
     pub output_format: OutputFormat,
+    // Output Quote Style.
+    pub quote_style: OutputQuoteStyle,
     /// Expand table format display, default off, could be on/off/auto.
     /// only works with output format `table`.
     pub expand: ExpandMode,
@@ -105,6 +107,14 @@ pub enum OutputFormat {
     Null,
 }
 
+#[derive(ValueEnum, Clone, Debug, PartialEq, Deserialize)]
+pub enum OutputQuoteStyle {
+    Always,
+    Necessary,
+    NonNumeric,
+    Never,
+}
+
 #[derive(ValueEnum, Clone, Debug, PartialEq)]
 pub enum TimeOption {
     Local,
@@ -157,6 +167,15 @@ impl Settings {
                     _ => return Err(anyhow!("Unknown output format: {}", cmd_value)),
                 }
             }
+            "quote_style" => {
+                self.quote_style = match cmd_value.to_ascii_lowercase().as_str() {
+                    "necessary" => OutputQuoteStyle::Necessary,
+                    "always" => OutputQuoteStyle::Always,
+                    "never" => OutputQuoteStyle::Never,
+                    "nonnumeric" => OutputQuoteStyle::NonNumeric,
+                    _ => return Err(anyhow!("Unknown quote style: {}", cmd_value)),
+                }
+            }
             "expand" => self.expand = cmd_value.into(),
             "time" => self.time = Some(cmd_value.try_into()?),
             "multi_line" => self.multi_line = cmd_value.parse()?,
@@ -217,6 +236,7 @@ impl Default for Settings {
             progress_color: "cyan".to_string(),
             prompt: "{user}@{warehouse}/{database}> ".to_string(),
             output_format: OutputFormat::Table,
+            quote_style: OutputQuoteStyle::Necessary,
             expand: ExpandMode::Off,
             show_progress: false,
             max_display_rows: 40,

diff --git a/cli/src/display.rs b/cli/src/display.rs
@@ -27,6 +27,7 @@ use tokio_stream::StreamExt;
 
 use indicatif::{HumanBytes, ProgressBar, ProgressState, ProgressStyle};
 
+use crate::config::OutputQuoteStyle;
 use crate::{
     ast::format_query,
     config::{ExpandMode, OutputFormat, Settings},
@@ -175,8 +176,14 @@ impl<'a> FormatDisplay<'a> {
     }
 
     async fn display_csv(&mut self) -> Result<()> {
+        let quote_style = match self.settings.quote_style {
+            OutputQuoteStyle::Always => csv::QuoteStyle::Always,
+            OutputQuoteStyle::Necessary => csv::QuoteStyle::Necessary,
+            OutputQuoteStyle::NonNumeric => csv::QuoteStyle::NonNumeric,
+            OutputQuoteStyle::Never => csv::QuoteStyle::Never,
+        };
         let mut wtr = csv::WriterBuilder::new()
-            .quote_style(csv::QuoteStyle::Necessary)
+            .quote_style(quote_style)
             .from_writer(std::io::stdout());
         while let Some(line) = self.data.next().await {
             match line {
@@ -198,9 +205,16 @@ impl<'a> FormatDisplay<'a> {
     }
 
     async fn display_tsv(&mut self) -> Result<()> {
+        let quote_style = match self.settings.quote_style {
+            OutputQuoteStyle::Always => csv::QuoteStyle::Always,
+            OutputQuoteStyle::Necessary => csv::QuoteStyle::Necessary,
+            OutputQuoteStyle::NonNumeric => csv::QuoteStyle::NonNumeric,
+            OutputQuoteStyle::Never => csv::QuoteStyle::Never,
+        };
         let mut wtr = csv::WriterBuilder::new()
             .delimiter(b'\t')
-            .quote_style(csv::QuoteStyle::Necessary)
+            .quote(b'"')
+            .quote_style(quote_style)
             .from_writer(std::io::stdout());
         while let Some(line) = self.data.next().await {
             match line {

diff --git a/cli/src/main.rs b/cli/src/main.rs
@@ -25,6 +25,7 @@ use std::{
     io::{stdin, IsTerminal},
 };
 
+use crate::config::OutputQuoteStyle;
 use anyhow::{anyhow, Result};
 use clap::{ArgAction, CommandFactory, Parser, ValueEnum};
 use config::{Config, OutputFormat, Settings, TimeOption};
@@ -149,6 +150,9 @@ struct Args {
     #[clap(short = 'o', long, help = "Output format")]
     output: Option<OutputFormat>,
 
+    #[clap(short = 's', long, help = "Output quote style")]
+    quote_style: Option<OutputQuoteStyle>,
+
     #[clap(
         long,
         help = "Show progress for query execution in stderr, only works with output format `table` and `null`."
@@ -314,6 +318,9 @@ pub async fn main() -> Result<()> {
     if let Some(output) = args.output {
         settings.output_format = output;
     }
+    if let Some(quote_style) = args.quote_style {
+        settings.quote_style = quote_style
+    }
     if args.progress {
         settings.show_progress = true;
     }