Merge branch 'main' into er/feat/chunk_metrics

src/commands/sql.rs

@@ -17,8 +17,14 @@ mod repl_command;
 #[derive(Debug, StructOpt)]
 pub struct Config {
     // TODO add an option to avoid saving history
-// TODO add an option to specify the default database (rather than having to set it via USE DATABASE)
-// TODO add an option to specify output formatting (now it is hard coded to use pretty printing)
+
+    // TODO add an option to specify the default database (rather than having to set it via USE DATABASE)
+    /// Format to use for output. Can be overridden using
+    /// `SET FORMAT` command
+    ///
+    /// Optional format ('pretty', 'json', or 'csv')
+    #[structopt(short, long, default_value = "pretty")]
+    format: String,
 }
 
 #[derive(Debug, Snafu)]
@@ -52,7 +58,11 @@ pub async fn command(url: String, config: Config) -> Result<()> {
     println!("Connected to IOx Server at {}", url);
     check_health(connection.clone()).await?;
 
-    repl::Repl::new(connection).run().await.context(Repl)
+    let mut repl = repl::Repl::new(connection);
+
+    repl.set_output_format(config.format).context(Repl)?;
+
+    repl.run().await.context(Repl)
 }
 
 async fn check_health(connection: Connection) -> Result<()> {

src/commands/sql/repl.rs

@@ -27,6 +27,12 @@ pub enum Error {
         source: influxdb_iox_client::format::Error,
     },
 
+    #[snafu(display("Error setting format to '{}': {}", requested_format, source))]
+    SettingFormat {
+        requested_format: String,
+        source: influxdb_iox_client::format::Error,
+    },
+
     #[snafu(display("Error parsing command: {}", message))]
     ParsingCommand { message: String },
 
@@ -88,6 +94,9 @@ pub struct Repl {
 
     /// database name against which SQL commands are run
     query_engine: Option<QueryEngine>,
+
+    /// Formatter to use to format query results
+    output_format: QueryOutputFormat,
 }
 
 impl Repl {
@@ -109,6 +118,8 @@ impl Repl {
 
         let prompt = "> ".to_string();
 
+        let output_format = QueryOutputFormat::Pretty;
+
         Self {
             rl,
             prompt,
@@ -116,6 +127,7 @@ impl Repl {
             management_client,
             flight_client,
             query_engine: None,
+            output_format,
         }
     }
 
@@ -151,6 +163,9 @@ impl Repl {
                     info!("exiting at user request");
                     return Ok(());
                 }
+                ReplCommand::SetFormat { format } => {
+                    self.set_output_format(format)?;
+                }
             }
         }
     }
@@ -276,6 +291,17 @@ impl Repl {
         self.query_engine = Some(query_engine)
     }
 
+    /// Sets the output format to the specified format
+    pub fn set_output_format<S: AsRef<str>>(&mut self, requested_format: S) -> Result<()> {
+        let requested_format = requested_format.as_ref();
+
+        self.output_format = requested_format
+            .parse()
+            .context(SettingFormat { requested_format })?;
+        println!("Set output format format to {}", self.output_format);
+        Ok(())
+    }
+
     // TODO make a setting for changing if we cache remote state or not
     async fn remote_state(&mut self) -> Result<RemoteState> {
         let state = RemoteState::try_new(&mut self.management_client).await?;
@@ -284,9 +310,10 @@ impl Repl {
 
     /// Prints to the specified output format
     fn print_results(&self, batches: &[RecordBatch]) -> Result<()> {
-        // TODO make query output format configurable
-        let output_format = QueryOutputFormat::Pretty;
-        let formatted_results = output_format.format(batches).context(FormattingResults)?;
+        let formatted_results = self
+            .output_format
+            .format(batches)
+            .context(FormattingResults)?;
         println!("{}", formatted_results);
         Ok(())
     }

src/commands/sql/repl_command.rs

@@ -8,6 +8,7 @@ pub enum ReplCommand {
     Help,
     ShowDatabases,
     Observer,
+    SetFormat { format: String },
     UseDatabase { db_name: String },
     SqlCommand { sql: String },
     Exit,
@@ -71,6 +72,10 @@ impl TryInto<ReplCommand> for String {
             })
         } else if commands.len() == 2 && commands[0] == "show" && commands[1] == "databases" {
             Ok(ReplCommand::ShowDatabases)
+        } else if commands.len() == 3 && commands[0] == "set" && commands[1] == "format" {
+            Ok(ReplCommand::SetFormat {
+                format: raw_commands[2].to_string(),
+            })
         } else {
             // Default is to treat the entire string like SQL
             Ok(ReplCommand::SqlCommand { sql: self })
@@ -89,6 +94,8 @@ SHOW DATABASES: List databases available on the server
 
 USE [DATABASE] <name>: Set the current remote database to name
 
+SET FORMAT <format>: Set the output format to Pretty, csv or json
+
 OBSERVER: Locally query unified queryable views of remote system tables
 
 [EXIT | QUIT]: Quit this session and exit the program
@@ -221,6 +228,22 @@ mod tests {
         assert_eq!("use database foo BAR".try_into(), expected);
     }
 
+    #[test]
+    fn set_format() {
+        let expected = Ok(ReplCommand::SetFormat {
+            format: "csv".to_string(),
+        });
+        assert_eq!(" set format csv".try_into(), expected);
+        assert_eq!("SET format   csv;".try_into(), expected);
+        assert_eq!("set  format csv".try_into(), expected);
+        assert_eq!("set format csv;".try_into(), expected);
+
+        let expected = Ok(ReplCommand::SetFormat {
+            format: "Hmm".to_string(),
+        });
+        assert_eq!("set format Hmm".try_into(), expected);
+    }
+
     #[test]
     fn sql_command() {
         let expected = sql_cmd("SELECT * from foo");

tests/end_to_end_cases/sql_cli.rs

@@ -128,6 +128,51 @@ async fn test_sql_use_database() {
         .stdout(predicate::str::contains(expected_output).and(predicate::str::contains("1 row")));
 }
 
+#[tokio::test]
+async fn test_sql_format() {
+    let fixture = ServerFixture::create_shared().await;
+    let addr = fixture.grpc_base();
+
+    let db_name = rand_name();
+    create_two_partition_database(&db_name, fixture.grpc_channel()).await;
+
+    let expected_output = r#"
+host,running,sleeping,time,total
+foo,4,514,2020-06-23T06:38:30.000000000,519
+"#
+    .trim();
+
+    Command::cargo_bin("influxdb_iox")
+        .unwrap()
+        .arg("sql")
+        .arg("--host")
+        .arg(addr)
+        // add flag to
+        .arg("--format")
+        .arg("csv")
+        .write_stdin(format!("use {};\n\nselect * from cpu;", db_name))
+        .assert()
+        .success()
+        .stdout(predicate::str::contains(expected_output));
+
+    // Same command but use `set format` command rather than cli flag
+    Command::cargo_bin("influxdb_iox")
+        .unwrap()
+        .arg("sql")
+        .arg("--host")
+        .arg(addr)
+        .write_stdin(format!(
+            "use {};\n\nset format csv;\n\nselect * from cpu;",
+            db_name
+        ))
+        .assert()
+        .success()
+        .stdout(
+            predicate::str::contains(expected_output)
+                .and(predicate::str::contains("Set output format format to csv")),
+        );
+}
+
 #[tokio::test]
 async fn test_sql_observer() {
     let fixture = ServerFixture::create_shared().await;