From 8fbcb4e6a861b9051e09db4b7b1b9dfa4add3d80 Mon Sep 17 00:00:00 2001
From: slinkydeveloper <francescoguard@gmail.com>
Date: Tue, 28 May 2024 16:11:25 +0200
Subject: [PATCH 1/2] Add json output to sql queries. Un-hide the sql command.

---
 cli/Cargo.toml          |  2 +-
 cli/src/app.rs          |  1 -
 cli/src/commands/sql.rs | 60 ++++++++++++++++++++++++++++-------------
 3 files changed, 42 insertions(+), 21 deletions(-)

diff --git a/cli/Cargo.toml b/cli/Cargo.toml
index c12cfa20e..c55cfcb6d 100644
--- a/cli/Cargo.toml
+++ b/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 }
diff --git a/cli/src/app.rs b/cli/src/app.rs
index 0dd1d06cb..c513838f0 100644
--- a/cli/src/app.rs
+++ b/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")]
diff --git a/cli/src/commands/sql.rs b/cli/src/commands/sql.rs
index a61afdd49..0a0ede57f 100644
--- a/cli/src/commands/sql.rs
+++ b/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]).unwrap();
+        }
+        writer.finish().unwrap();
+    } else if sql_opts.jsonl {
+        let mut writer = arrow::json::LineDelimitedWriter::new(io::stdout());
+        for batch in resp.batches {
+            writer.write_batches(&[&batch]).unwrap();
+        }
+        writer.finish().unwrap();
+    } 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!(

From c66ace75c9dfe0f651487363b72c58b6d584ea2d Mon Sep 17 00:00:00 2001
From: slinkydeveloper <francescoguard@gmail.com>
Date: Wed, 29 May 2024 11:49:34 +0200
Subject: [PATCH 2/2] Return the error

---
 cli/src/commands/sql.rs | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/cli/src/commands/sql.rs b/cli/src/commands/sql.rs
index 0a0ede57f..acbbff138 100644
--- a/cli/src/commands/sql.rs
+++ b/cli/src/commands/sql.rs
@@ -72,15 +72,15 @@ async fn run_query(env: &CliEnv, sql_opts: &Sql) -> Result<()> {
     if sql_opts.json {
         let mut writer = arrow::json::ArrayWriter::new(io::stdout());
         for batch in resp.batches {
-            writer.write_batches(&[&batch]).unwrap();
+            writer.write_batches(&[&batch])?;
         }
-        writer.finish().unwrap();
+        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]).unwrap();
+            writer.write_batches(&[&batch])?;
         }
-        writer.finish().unwrap();
+        writer.finish()?;
     } else {
         let format_options = FormatOptions::default().with_display_error(true);
         for batch in resp.batches {