feat: add file output

Author: esadek

src/cli.rs

@@ -22,6 +22,7 @@ pub struct DatabaseConfig {
     pub options: Vec<(String, String)>,
     pub query_source: QuerySource,
     pub table_mode: TableMode,
+    pub output_path: Option<PathBuf>,
 }
 
 fn is_stdin_piped() -> bool {
@@ -60,6 +61,10 @@ pub fn parse_args() -> DatabaseConfig {
             .long("file")
             .help("Read and execute file and exit")
             .conflicts_with("query"),
+        Arg::new("output")
+            .long("output")
+            .help("Write result to file")
+            .value_name("file"),
     ];
     let command = Command::new("adbcli")
         .version(env!("CARGO_PKG_VERSION"))
@@ -114,6 +119,12 @@ pub fn parse_args() -> DatabaseConfig {
         None => TableMode::default(),
     };
 
+    let output_path = matches.get_one::<String>("output").map(PathBuf::from);
+    if output_path.is_some() && matches!(query_source, QuerySource::Interactive) {
+        eprintln!("Error: --output cannot be used in interactive mode");
+        exit(1);
+    }
+
     DatabaseConfig {
         driver_name,
         uri,
@@ -122,6 +133,7 @@ pub fn parse_args() -> DatabaseConfig {
         options,
         query_source,
         table_mode,
+        output_path,
     }
 }
 
@@ -185,6 +197,7 @@ mod tests {
             options: vec![("key1".to_string(), "val1".to_string())],
             query_source: QuerySource::Interactive,
             table_mode: TableMode::default(),
+            output_path: None,
         };
 
         assert_eq!(config.driver_name, "test_driver");
@@ -205,6 +218,7 @@ mod tests {
             options: vec![],
             query_source: QuerySource::Interactive,
             table_mode: TableMode::AsciiMarkdown,
+            output_path: None,
         };
 
         assert_eq!(config.driver_name, "test_driver");
@@ -214,4 +228,20 @@ mod tests {
         assert!(config.options.is_empty());
         assert_eq!(config.table_mode, TableMode::AsciiMarkdown);
     }
+
+    #[test]
+    fn test_database_config_with_output_path() {
+        let config = DatabaseConfig {
+            driver_name: "test_driver".to_string(),
+            uri: None,
+            username: None,
+            password: None,
+            options: vec![],
+            query_source: QuerySource::Query("SELECT 1".to_string()),
+            table_mode: TableMode::default(),
+            output_path: Some(PathBuf::from("output.json")),
+        };
+
+        assert_eq!(config.output_path, Some(PathBuf::from("output.json")));
+    }
 }

src/database.rs

@@ -2,10 +2,10 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::cli::DatabaseConfig;
-use crate::table::{self, TableMode};
 use adbc_core::options::{AdbcVersion, OptionDatabase, OptionValue};
 use adbc_core::{Connection, Database, Driver, LOAD_FLAG_DEFAULT, Statement};
 use adbc_driver_manager::ManagedDriver;
+use arrow_array::RecordBatch;
 use std::process::exit;
 
 pub fn initialize_connection(config: DatabaseConfig) -> impl Connection {
@@ -61,10 +61,9 @@ pub fn initialize_connection(config: DatabaseConfig) -> impl Connection {
 pub fn execute_query(
     connection: &mut impl adbc_core::Connection,
     sql: &str,
-    table_mode: TableMode,
-) -> Result<(), String> {
+) -> Result<Vec<RecordBatch>, String> {
     if sql.trim().is_empty() {
-        return Ok(());
+        return Ok(vec![]);
     }
 
     let mut statement = connection
@@ -79,9 +78,9 @@ pub fn execute_query(
         .execute()
         .map_err(|e| format!("Failed to execute statement: {e}"))?;
 
-    let batches: Vec<arrow_array::RecordBatch> = reader
+    let batches: Vec<RecordBatch> = reader
         .collect::<Result<_, _>>()
         .map_err(|e| format!("Failed to collect batches: {e}"))?;
 
-    table::print_batches(&batches, table_mode).map_err(|e| format!("Failed to print batches: {e}"))
+    Ok(batches)
 }

src/main.rs

@@ -4,6 +4,7 @@
 mod cli;
 mod database;
 mod highlighter;
+mod output;
 mod repl;
 mod table;
 
@@ -15,44 +16,54 @@ fn main() {
     let config = parse_args();
     let query_source = config.query_source.clone();
     let table_mode = config.table_mode;
+    let output_path = config.output_path.clone();
 
-    match query_source {
-        QuerySource::Query(sql) => {
-            let mut connection = database::initialize_connection(config);
-            if let Err(e) = database::execute_query(&mut connection, &sql, table_mode) {
-                eprintln!("{e}");
-                exit(1);
-            }
-        }
-        QuerySource::File(path) => {
-            let sql = match std::fs::read_to_string(&path) {
-                Ok(content) => content,
-                Err(e) => {
-                    eprintln!("Failed to read file {}: {e}", path.display());
-                    exit(1);
-                }
-            };
-            let mut connection = database::initialize_connection(config);
-            if let Err(e) = database::execute_query(&mut connection, &sql, table_mode) {
-                eprintln!("{e}");
-                exit(1);
-            }
-        }
+    if matches!(query_source, QuerySource::Interactive) {
+        let connection = database::initialize_connection(config);
+        repl::run_repl(connection, table_mode);
+        return;
+    }
+
+    let sql = match query_source {
+        QuerySource::Query(sql) => sql,
+        QuerySource::File(path) => std::fs::read_to_string(&path).unwrap_or_else(|e| {
+            eprintln!("Failed to read file {}: {e}", path.display());
+            exit(1);
+        }),
         QuerySource::Stdin => {
             let mut sql = String::new();
-            if let Err(e) = std::io::stdin().read_to_string(&mut sql) {
-                eprintln!("Failed to read from stdin: {e}");
-                exit(1);
-            }
-            let mut connection = database::initialize_connection(config);
-            if let Err(e) = database::execute_query(&mut connection, &sql, table_mode) {
-                eprintln!("{e}");
-                exit(1);
-            }
-        }
-        QuerySource::Interactive => {
-            let connection = database::initialize_connection(config);
-            repl::run_repl(connection, table_mode);
+            std::io::stdin()
+                .read_to_string(&mut sql)
+                .unwrap_or_else(|e| {
+                    eprintln!("Failed to read from stdin: {e}");
+                    exit(1);
+                });
+            sql
         }
+        QuerySource::Interactive => unreachable!(),
+    };
+
+    let mut connection = database::initialize_connection(config);
+    let batches = database::execute_query(&mut connection, &sql).unwrap_or_else(|e| {
+        eprintln!("{e}");
+        exit(1);
+    });
+    if let Err(e) = output_results(&batches, table_mode, output_path.as_deref()) {
+        eprintln!("{e}");
+        exit(1);
+    }
+}
+
+fn output_results(
+    batches: &[arrow_array::RecordBatch],
+    table_mode: table::TableMode,
+    output_path: Option<&std::path::Path>,
+) -> Result<(), String> {
+    if let Some(path) = output_path {
+        output::write_batches_to_file(batches, path)
+            .map_err(|e| format!("Failed to write output file: {e}"))
+    } else {
+        table::print_batches(batches, table_mode)
+            .map_err(|e| format!("Failed to print batches: {e}"))
     }
 }