Add json output to sql queries (#1561)

* Add json output to sql queries. Un-hide the sql command.

cli/Cargo.toml

@@ -21,7 +21,7 @@ restate-service-protocol = { workspace = true, features = ["awakeable-id"] }
 restate-types = { workspace = true }
 
 anyhow = { workspace = true }
-arrow = { version = "51.0.0", features = ["ipc", "prettyprint"] }
+arrow = { version = "51.0.0", features = ["ipc", "prettyprint", "json"] }
 arrow_convert = { version = "0.6.6" }
 axum = { workspace = true, default-features = false, features = ["http1", "http2", "query", "tokio"] }
 bytes = { workspace = true }

cli/src/app.rs

@@ -91,7 +91,6 @@ pub enum Command {
     #[clap(subcommand)]
     Invocations(invocations::Invocations),
     /// Runs SQL queries against the data fusion service
-    #[clap(hide = true)]
     Sql(sql::Sql),
     /// Download one of Restate's examples in this directory.
     #[clap(name = "example", alias = "examples")]

cli/src/commands/sql.rs

@@ -9,8 +9,6 @@
 // by the Apache License, Version 2.0.
 //
 
-use std::time::Instant;
-
 use anyhow::Result;
 use arrow::error::ArrowError;
 use arrow::util::display::ArrayFormatter;
@@ -20,6 +18,8 @@ use comfy_table::Cell;
 use comfy_table::Table;
 use serde::Deserialize;
 use serde::Serialize;
+use std::io;
+use std::time::Instant;
 
 use crate::c_eprintln;
 use crate::c_println;
@@ -37,6 +37,14 @@ pub struct Sql {
 
     #[clap(flatten)]
     watch: Watch,
+
+    /// Print result as line delimited json instead of using the tabular format
+    #[arg(long, alias = "ldjson")]
+    pub jsonl: bool,
+
+    /// Print result as json array instead of using the tabular format
+    #[arg(long)]
+    pub json: bool,
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -61,27 +69,41 @@ async fn run_query(env: &CliEnv, sql_opts: &Sql) -> Result<()> {
     }
     table.set_styled_header(headers);
 
-    let format_options = FormatOptions::default().with_display_error(true);
-    for batch in resp.batches {
-        let formatters = batch
-            .columns()
-            .iter()
-            .map(|c| ArrayFormatter::try_new(c.as_ref(), &format_options))
-            .collect::<Result<Vec<_>, ArrowError>>()?;
+    if sql_opts.json {
+        let mut writer = arrow::json::ArrayWriter::new(io::stdout());
+        for batch in resp.batches {
+            writer.write_batches(&[&batch])?;
+        }
+        writer.finish()?;
+    } else if sql_opts.jsonl {
+        let mut writer = arrow::json::LineDelimitedWriter::new(io::stdout());
+        for batch in resp.batches {
+            writer.write_batches(&[&batch])?;
+        }
+        writer.finish()?;
+    } else {
+        let format_options = FormatOptions::default().with_display_error(true);
+        for batch in resp.batches {
+            let formatters = batch
+                .columns()
+                .iter()
+                .map(|c| ArrayFormatter::try_new(c.as_ref(), &format_options))
+                .collect::<Result<Vec<_>, ArrowError>>()?;
 
-        for row in 0..batch.num_rows() {
-            let mut cells = Vec::new();
-            for formatter in &formatters {
-                cells.push(Cell::new(formatter.value(row)));
+            for row in 0..batch.num_rows() {
+                let mut cells = Vec::new();
+                for formatter in &formatters {
+                    cells.push(Cell::new(formatter.value(row)));
+                }
+                table.add_row(cells);
             }
-            table.add_row(cells);
         }
-    }
 
-    // Only print if there are actual results.
-    if table.row_count() > 0 {
-        c_println!("{}", table);
-        c_println!();
+        // Only print if there are actual results.
+        if table.row_count() > 0 {
+            c_println!("{}", table);
+            c_println!();
+        }
     }
 
     c_eprintln!(