Support output format option

Cargo.toml

@@ -7,14 +7,15 @@ repository = "https://github.com/sundy-li/arrow_cli"
 edition = "2024"
 license = "Apache-2.0"
 name = "arrow_cli"
-version = "0.4.1"
+version = "0.5.0"
 
 
 [dependencies]
 arrow-array = "57"
 arrow-cast = { version = "57", features = ["prettyprint"] }
 arrow-csv = "57"
 arrow-flight = { version = "57", features = ["flight-sql-experimental"] }
+arrow-json = "57"
 arrow-schema = "57"
 atty = "0.2"
 clap = { version = "4.5", features = ["derive"] }

README.md

@@ -7,7 +7,7 @@
 
 ## Overview
 
-arrow_cli is a CLI tool for interacting with server in Flight SQL protocol.
+arrow_cli is a CLI tool for interacting with a server that speaks the Flight SQL protocol.
 
 ## Install 
 
@@ -17,7 +17,7 @@ cargo install arrow_cli
 
 ## Usage
 
-```
+```text
 > arrow_cli --help
 Usage: arrow_cli [OPTIONS]
 
@@ -29,59 +29,87 @@ Options:
       --tls
       --timeout <TIMEOUT>    Request timeout in seconds [default: 180]
       --prepared             Execute query using prepared statement
+      --print-schema         Print resultset schema
+      --output <OUTPUT>      Result output format [default: table] [possible values: table, json, csv, tsv, psv]
+  -c, --command <COMMAND>    Execute SQL command and exit
   -h, --help                 Print help
 ```
 
 ## Examples
 
-### REPL
-```sql
-❯ arrow_cli -h arch -u sundy -p abc --port 8900
-Welcome to Arrow CLI.
-Connecting to http://arch:8900/ as user sundy.
-
-arch :) select avg(number) from numbers(10);
-
-select avg(number) from numbers(10);
-
+### Single command with table output
+```bash
+❯ arrow_cli -h arch -u sundy -p abc --port 8900 --output table --command "select avg(number) from numbers(10);"
 +-------------+
 | avg(number) |
 +-------------+
 | 4.5         |
 +-------------+
 
-1 rows in set (0.036 sec)
-
-arch :) show tables like 'c%';
-
-show tables like 'c%';
-
-+-------------------+
-| tables_in_default |
-+-------------------+
-| customer          |
-+-------------------+
+1 rows in set (tickets received in 0.036 sec, rows received in 0.036 sec)
+```
 
-1 rows in set (0.030 sec)
+### Single command with JSON output
 
-arch :) exit
-Bye
+```bash
+❯ arrow_cli -h arch -u sundy -p abc --port 8900 --output json --command "select number from numbers(3)"
+{"number":0}
+{"number":1}
+{"number":2}
 ```
 
-### StdIn Pipe
+### StdIn pipe with CSV output
 
 ```bash
-❯ echo "select number from numbers(3)" | arrow_cli -h arch -u sundy -p abc --port 8900
+❯ echo "select number from numbers(3)" | arrow_cli -h arch -u sundy -p abc --port 8900 --output csv
 0
 1
 2
 ```
 
+### StdIn pipe with TSV output
+
+```bash
+❯ echo "select number, concat('v', to_string(number)) from numbers(3)" | arrow_cli -h arch -u sundy -p abc --port 8900 --output tsv
+0	v0
+1	v1
+2	v2
+```
+
+### Single command with PSV output
+
+```bash
+❯ arrow_cli -h arch -u sundy -p abc --port 8900 --output psv --command "select number, concat('v', to_string(number)) from numbers(3)"
+0|v0
+1|v1
+2|v2
+```
+
+### Interactive session with JSON output
+
+```text
+❯ arrow_cli -h arch -u sundy -p abc --port 8900 --output json
+Welcome to Arrow CLI v0.4.1.
+Connecting to http://arch:8900/ as user sundy.
+
+arch :) select number from numbers(2);
+
+select number from numbers(2);
+
+{"number":0}
+{"number":1}
+
+arch :) exit
+Bye
+```
+
 ## Features
 
 - basic keywords highlight
 - basic auto-completion
 - select query support
+- output formats: table, json, csv, tsv, psv
+- delimited formats use: csv=`,`, tsv=`\t`, psv=`|`
 - TBD
 
 #### License

SKILL.md

@@ -45,16 +45,10 @@ Output:
 echo "select 1" | arrow_cli --host localhost --port 8900 --user admin --password abc
 ```
 
-Output:
-
-```text
-1
-```
-
 ### Option: print the result schema
 
 ```bash
-arrow_cli --host localhost --port 8900 --user admin --password abc --print-schema --command "select number from numbers(3);"
+arrow_cli --host localhost --port 8900 --user admin --password abc --print-schema --command "select 1"
 ```
 
 Output:
@@ -79,5 +73,12 @@ Schema {
 1 rows in set (tickets received in 0.008 sec, rows received in 0.010 sec)
 ```
 
-## Other notes
-- Run `arrow_cli --help` for more details
+### Option: output formats
+
+- `--output table`: pretty table output with row count and timing summary
+- `--output json`: line-delimited JSON rows
+- `--output csv`: comma-separated rows without headers
+- `--output tsv`: tab-separated rows without headers
+- `--output psv`: pipe-separated rows without headers
+
+Run `arrow_cli --help` for more usage details.

src/main.rs

@@ -1,4 +1,5 @@
 mod helper;
+mod output;
 mod session;
 
 use std::time::Duration;
@@ -7,6 +8,7 @@ use arrow_schema::ArrowError;
 use atty::Stream;
 
 use clap::Parser;
+use output::Output;
 use tonic::transport::{ClientTlsConfig, Endpoint};
 
 #[derive(Debug, Parser, PartialEq)]
@@ -44,6 +46,14 @@ struct Args {
     #[clap(long, default_value = "false", help = "Print resultset schema")]
     print_schema: bool,
 
+    #[clap(
+        long,
+        value_enum,
+        default_value_t = Output::Table,
+        help = "Result output format"
+    )]
+    output: Output,
+
     #[clap(short = 'c', long, help = "Execute SQL command and exit")]
     command: Option<String>,
 }

src/output.rs

@@ -0,0 +1,132 @@
+use arrow_array::RecordBatch;
+use arrow_cast::pretty::pretty_format_batches;
+use arrow_csv::WriterBuilder as CsvWriterBuilder;
+use arrow_json::LineDelimitedWriter;
+use arrow_schema::ArrowError;
+use clap::ValueEnum;
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq, ValueEnum)]
+pub enum Output {
+    Table,
+    Json,
+    Csv,
+    Tsv,
+    Psv,
+}
+
+pub fn print_batches(batches: &[RecordBatch], output: Output) -> Result<(), ArrowError> {
+    let formatted = format_batches(batches, output)?;
+    print!("{formatted}");
+    Ok(())
+}
+
+pub fn format_batches(batches: &[RecordBatch], output: Output) -> Result<String, ArrowError> {
+    match output {
+        Output::Table => format_table(batches),
+        Output::Json => format_json(batches),
+        Output::Csv => format_delimited(batches, b','),
+        Output::Tsv => format_delimited(batches, b'\t'),
+        Output::Psv => format_delimited(batches, b'|'),
+    }
+}
+
+fn format_table(batches: &[RecordBatch]) -> Result<String, ArrowError> {
+    Ok(format!("{}\n", pretty_format_batches(batches)?))
+}
+
+fn format_json(batches: &[RecordBatch]) -> Result<String, ArrowError> {
+    let mut bytes = vec![];
+    {
+        let mut writer = LineDelimitedWriter::new(&mut bytes);
+        for batch in batches {
+            writer.write(batch)?;
+        }
+        writer.finish()?;
+    }
+
+    let formatted = String::from_utf8(bytes).map_err(|e| ArrowError::JsonError(e.to_string()))?;
+    Ok(formatted)
+}
+
+fn format_delimited(batches: &[RecordBatch], delimiter: u8) -> Result<String, ArrowError> {
+    let mut bytes = vec![];
+    {
+        let mut writer = CsvWriterBuilder::new()
+            .with_header(false)
+            .with_delimiter(delimiter)
+            .build(&mut bytes);
+        for batch in batches {
+            writer.write(batch)?;
+        }
+    }
+
+    let formatted = String::from_utf8(bytes).map_err(|e| ArrowError::CsvError(e.to_string()))?;
+    Ok(formatted)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow_array::{Int64Array, StringArray};
+    use arrow_schema::{DataType, Field, Schema};
+    use std::sync::Arc;
+
+    fn sample_batch() -> RecordBatch {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("number", DataType::Int64, false),
+            Field::new("label", DataType::Utf8, false),
+        ]));
+
+        RecordBatch::try_new(
+            schema,
+            vec![
+                Arc::new(Int64Array::from(vec![1, 2])),
+                Arc::new(StringArray::from(vec!["one", "two"])),
+            ],
+        )
+        .unwrap()
+    }
+
+    #[test]
+    fn formats_table_output() {
+        let batches = vec![sample_batch()];
+        let result = format_table(&batches);
+        assert!(result.is_ok());
+        assert!(result.unwrap().contains("| number |"));
+    }
+
+    #[test]
+    fn formats_json_output() {
+        let batches = vec![sample_batch()];
+        let result = format_json(&batches);
+        assert!(result.is_ok());
+        assert_eq!(
+            result.unwrap(),
+            "{\"number\":1,\"label\":\"one\"}\n{\"number\":2,\"label\":\"two\"}\n"
+        );
+    }
+
+    #[test]
+    fn formats_csv_output() {
+        let batches = vec![sample_batch()];
+        let result = format_batches(&batches, Output::Csv);
+        assert!(result.is_ok());
+        assert_eq!(result.unwrap(), "1,one\n2,two\n");
+    }
+
+    #[test]
+    fn formats_tsv_output() {
+        let batches = vec![sample_batch()];
+        let result = format_batches(&batches, Output::Tsv);
+        assert!(result.is_ok());
+        assert_eq!(result.unwrap(), "1\tone\n2\ttwo\n");
+    }
+
+    #[test]
+    fn formats_psv_output() {
+        let batches = vec![sample_batch()];
+        let result = format_batches(&batches, Output::Psv);
+        assert!(result.is_ok());
+        assert_eq!(result.unwrap(), "1|one\n2|two\n");
+    }
+}