diff --git a/Cargo.lock b/Cargo.lock index d2f6993..2ffbe2b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -609,6 +609,7 @@ dependencies = [ "clap", "comfy-table", "nu-ansi-term", + "parquet", "reedline", "syntect", "tempfile", @@ -1096,6 +1097,32 @@ dependencies = [ "windows-link 0.2.1", ] +[[package]] +name = "parquet" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff322f54b1a0f9288e614ed1f2d329b380af5476420db19f46ffb865e1163d73" +dependencies = [ + "ahash", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-ipc", + "arrow-schema", + "arrow-select", + "base64", + "bytes", + "chrono", + "half", + "hashbrown", + "num-bigint", + "num-integer", + "num-traits", + "seq-macro", + "snap", + "twox-hash", +] + [[package]] name = "path-slash" version = "0.2.1" @@ -1299,6 +1326,12 @@ version = "1.0.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" +[[package]] +name = "seq-macro" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" + [[package]] name = "serde" version = "1.0.229" @@ -1418,6 +1451,12 @@ version = "1.16.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f9395f0f0eee849a9b707b2f06bb92a6a422090e2123bb2ef8e87a0e61892a8e" +[[package]] +name = "snap" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "199905e6153d6405f9728fe44daace35f8f837bbf830bb6e85fbd5828709a886" + [[package]] name = "strip-ansi-escapes" version = "0.2.1" @@ -1632,6 +1671,12 @@ version = "1.1.2+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7d56353a2a665ad0f41a421187180aab746c8c325620617ad883a99a1cbe66d2" +[[package]] +name = "twox-hash" +version = "2.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5283634e518fe9e82c7b20520bb4bc209009fd16c82077c802f8111ecbb0117a" + [[package]] name = "unicode-ident" version = "1.0.26" diff --git a/Cargo.toml b/Cargo.toml index 303b6d3..a9ea5e7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,10 +20,11 @@ arrow-schema = "59.3.0" clap = { version = "4.6.7", features = ["derive"] } comfy-table = "8.0.1" nu-ansi-term = "0.50.3" +parquet = { version = "59.3.0", default-features = false, features = ["arrow", "snap"] } reedline = "0.52.0" syntect = "5.3.0" +tempfile = "3.27.0" terminal-colorsaurus = "1.0.3" [dev-dependencies] arrow-buffer = "59.3.0" -tempfile = "3.27.0" diff --git a/README.md b/README.md index 961fb4d..1a0cdc3 100644 --- a/README.md +++ b/README.md @@ -21,7 +21,7 @@ A command-line tool for querying databases via [ADBC](https://arrow.apache.org/a - **Interactive SQL shell** - Execute SQL queries with command history and intuitive navigation - **Syntax highlighting** - SQL queries highlighted for improved readability - **Formatted output** - Results displayed in clean, aligned tables with dynamic column width -- **File export** - Export query results to JSON, CSV, or Arrow IPC files +- **File export** - Export query results to JSON, JSON lines, CSV, Arrow IPC, or Parquet files - **Fast and lightweight** - Built in Rust for high performance and minimal resource usage ## Installation @@ -107,8 +107,11 @@ Execute a query and output the result to a file: ```sh databow --driver duckdb --query "SELECT 42 AS the_answer" --output result.json +databow --driver duckdb --query "SELECT 42 AS the_answer" --output result.jsonl databow --driver duckdb --query "SELECT 42 AS the_answer" --output result.csv databow --driver duckdb --query "SELECT 42 AS the_answer" --output result.arrow +databow --driver duckdb --query "SELECT 42 AS the_answer" --output result.arrows +databow --driver duckdb --query "SELECT 42 AS the_answer" --output result.parquet ``` ## Reference diff --git a/docs/index.md b/docs/index.md index 06e5a75..2954f7e 100644 --- a/docs/index.md +++ b/docs/index.md @@ -19,5 +19,5 @@ databow is a command-line tool for querying databases. - **Interactive SQL shell** - Execute SQL queries with command history and intuitive navigation - **Syntax highlighting** - SQL queries highlighted for improved readability - **Formatted output** - Results displayed in clean, aligned tables with dynamic column width -- **File export** - Export query results to JSON, CSV, or Arrow IPC files +- **File export** - Export query results to JSON, JSON lines, CSV, Arrow IPC, or Parquet files - **Fast and lightweight** - Built in Rust for high performance and minimal resource usage diff --git a/docs/reference.md b/docs/reference.md index 25c203a..14bcd52 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -108,11 +108,14 @@ databow --driver duckdb --query "SELECT 42 AS the_answer" --output result.json The output format is inferred from the file extension: -| Extension | Format | -|-----------------|-----------| -| `.json` | JSON | -| `.csv` | CSV | -| `.arrow`, `.ipc`| Arrow IPC | +| Extension | Format | +|-----------------|------------------| +| `.json` | JSON | +| `.jsonl` | JSON lines | +| `.csv` | CSV | +| `.arrow`, `.ipc`| Arrow IPC file | +| `.arrows` | Arrow IPC stream | +| `.parquet` | Parquet | ## --help diff --git a/docs/tutorial.md b/docs/tutorial.md index 59ec6ae..e38cb24 100644 --- a/docs/tutorial.md +++ b/docs/tutorial.md @@ -143,7 +143,7 @@ $ databow --profile warehouse --file query.sql └───────────────────┘ ``` -Instead of printing query results to stdout, the [`--output` argument](/reference/#-output) can be used to write results to JSON, CSV, or Arrow IPC files: +Instead of printing query results to stdout, the [`--output` argument](/reference/#-output) can be used to write results to JSON, JSON lines, CSV, Arrow IPC file, Arrow IPC stream, or Parquet files: ```console $ databow --profile warehouse --query "SELECT * FROM penguins" --output penguins.csv diff --git a/src/cli.rs b/src/cli.rs index 8dac594..9dd9061 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -1,6 +1,7 @@ // Copyright 2026 Columnar Technologies Inc. // SPDX-License-Identifier: Apache-2.0 +use crate::output::OutputFormat; use crate::table::TableMode; use clap::{Arg, ArgAction, Command, value_parser}; use std::path::PathBuf; @@ -84,7 +85,8 @@ pub fn parse_args() -> AppConfig { Arg::new("output") .long("output") .help("Write result to file") - .value_name("file"), + .value_name("file") + .value_parser(parse_output_path), ]; let command = Command::new("databow") .version(env!("CARGO_PKG_VERSION")) @@ -163,7 +165,7 @@ pub fn parse_args() -> AppConfig { .copied() .unwrap_or_default(); - let output_path = matches.get_one::("output").map(PathBuf::from); + let output_path = matches.get_one::("output").cloned(); if output_path.is_some() && matches!(query_source, QuerySource::Interactive) { eprintln!("Error: --output cannot be used in interactive mode"); exit(1); @@ -177,6 +179,12 @@ pub fn parse_args() -> AppConfig { } } +fn parse_output_path(s: &str) -> Result { + let path = PathBuf::from(s); + OutputFormat::from_path(&path)?; + Ok(path) +} + fn uri_has_driver_scheme(uri: &str) -> bool { let Some(idx) = uri.find(':') else { return false; @@ -202,6 +210,26 @@ fn parse_option(option: &str) -> Result<(String, String), String> { mod tests { use super::*; + #[test] + fn test_parse_output_path_valid() { + assert_eq!( + parse_output_path("out.parquet").unwrap(), + PathBuf::from("out.parquet") + ); + } + + #[test] + fn test_parse_output_path_unsupported_extension() { + let err = parse_output_path("out.xyz").unwrap_err(); + assert!(err.contains("Unsupported file extension")); + } + + #[test] + fn test_parse_output_path_no_extension() { + let err = parse_output_path("out").unwrap_err(); + assert!(err.contains("no file extension")); + } + #[test] fn test_connection_source_direct() { let source = ConnectionSource::Direct { diff --git a/src/database.rs b/src/database.rs index 9dd5c4f..113273e 100644 --- a/src/database.rs +++ b/src/database.rs @@ -9,7 +9,8 @@ use adbc_driver_manager::profile::{ }; use adbc_driver_manager::{ManagedConnection, ManagedDatabase, ManagedDriver}; use arrow_array::cast::AsArray; -use arrow_array::{Array, RecordBatch, UnionArray}; +use arrow_array::{Array, RecordBatch, RecordBatchReader, UnionArray}; +use arrow_schema::{Schema, SchemaRef}; use std::collections::HashSet; /// Vendor (database product) metadata reported by an ADBC driver. @@ -304,6 +305,44 @@ fn merge_options( merged } +/// The result of executing a query: the result schema reported by the driver +/// and the collected record batches (which may be empty). +pub struct QueryResult { + pub schema: SchemaRef, + pub batches: Vec, +} + +pub fn execute_query( + connection: &mut impl adbc_core::Connection, + sql: &str, +) -> Result { + if sql.trim().is_empty() { + return Ok(QueryResult { + schema: std::sync::Arc::new(Schema::empty()), + batches: vec![], + }); + } + + let mut statement = connection + .new_statement() + .map_err(|e| format!("Failed to create statement: {e}"))?; + + statement + .set_sql_query(sql) + .map_err(|e| format!("Failed to set SQL query: {e}"))?; + + let reader = statement + .execute() + .map_err(|e| format!("Failed to execute statement: {e}"))?; + + let schema = reader.schema(); + let batches: Vec = reader + .collect::>() + .map_err(|e| format!("Failed to collect batches: {e}"))?; + + Ok(QueryResult { schema, batches }) +} + #[cfg(test)] mod tests { use super::*; @@ -487,30 +526,3 @@ mod tests { assert_eq!(merged.len(), 2); } } - -pub fn execute_query( - connection: &mut impl adbc_core::Connection, - sql: &str, -) -> Result, String> { - if sql.trim().is_empty() { - return Ok(vec![]); - } - - let mut statement = connection - .new_statement() - .map_err(|e| format!("Failed to create statement: {e}"))?; - - statement - .set_sql_query(sql) - .map_err(|e| format!("Failed to set SQL query: {e}"))?; - - let reader = statement - .execute() - .map_err(|e| format!("Failed to execute statement: {e}"))?; - - let batches: Vec = reader - .collect::>() - .map_err(|e| format!("Failed to collect batches: {e}"))?; - - Ok(batches) -} diff --git a/src/main.rs b/src/main.rs index 90a1aed..85e9c41 100644 --- a/src/main.rs +++ b/src/main.rs @@ -48,30 +48,37 @@ fn main() { QuerySource::Interactive => unreachable!(), }; + // A blank query has no result schema, so there is nothing meaningful to + // write; fail rather than silently producing an empty or garbage file. + if output_path.is_some() && sql.trim().is_empty() { + eprintln!("No query to execute; refusing to write an empty output file"); + exit(1); + } + let mut connection = database::initialize_connection(connection).unwrap_or_else(|e| { eprintln!("{e}"); exit(1); }); - let batches = database::execute_query(&mut connection, &sql).unwrap_or_else(|e| { + let result = 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()) { + if let Err(e) = output_results(&result, table_mode, output_path.as_deref()) { eprintln!("{e}"); exit(1); } } fn output_results( - batches: &[arrow_array::RecordBatch], + result: &database::QueryResult, 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) + output::write_batches_to_file(&result.schema, &result.batches, path) .map_err(|e| format!("Failed to write output file: {e}")) } else { - table::print_batches(batches, table_mode) + table::print_batches(&result.batches, table_mode) .map_err(|e| format!("Failed to print batches: {e}")) } } diff --git a/src/output.rs b/src/output.rs index 690335a..b8d3c9b 100644 --- a/src/output.rs +++ b/src/output.rs @@ -2,49 +2,172 @@ // SPDX-License-Identifier: Apache-2.0 use arrow::csv::writer::Writer as CsvWriter; -use arrow::ipc::writer::FileWriter as IpcWriter; -use arrow::json::writer::{JsonArray, Writer as JsonWriter}; +use arrow::ipc::writer::{FileWriter as IpcWriter, StreamWriter as IpcStreamWriter}; +use arrow::json::writer::{JsonArray, LineDelimited, Writer as JsonWriter}; use arrow_array::RecordBatch; -use arrow_schema::ArrowError; -use std::fs::File; -use std::path::Path; +use arrow_schema::{ArrowError, SchemaRef}; +use parquet::arrow::ArrowWriter; +use parquet::basic::Compression; +use parquet::file::properties::WriterProperties; +use std::io::Write; +use std::path::{Path, PathBuf}; +use tempfile::NamedTempFile; #[derive(Debug, Clone, Copy, PartialEq)] pub enum OutputFormat { Json, + Jsonl, Csv, Arrow, + ArrowStream, + Parquet, } impl OutputFormat { pub fn from_path(path: &Path) -> Result { match path.extension().and_then(|ext| ext.to_str()) { Some("json") => Ok(OutputFormat::Json), + Some("jsonl") => Ok(OutputFormat::Jsonl), Some("csv") => Ok(OutputFormat::Csv), Some("arrow" | "ipc") => Ok(OutputFormat::Arrow), + Some("arrows") => Ok(OutputFormat::ArrowStream), + Some("parquet") => Ok(OutputFormat::Parquet), Some(ext) => Err(format!("Unsupported file extension: '.{ext}'")), None => Err("Cannot infer format: no file extension".to_string()), } } } -pub fn write_batches_to_file(batches: &[RecordBatch], path: &Path) -> Result<(), ArrowError> { - if batches.is_empty() { - return Ok(()); - } - +/// Write query results to `path`, inferring the format from its extension. +/// +/// Output is written to a temporary file in the same directory and atomically +/// renamed into place only after the write succeeds, so a failed write never +/// leaves a partial file behind or clobbers an existing one. If `path` is a +/// symlink, the file it points to is replaced and the link is left intact. +/// +/// `schema` is used when there are no batches so that formats which can carry +/// a schema (Arrow IPC, Parquet, CSV header) still produce a valid, empty file. +pub fn write_batches_to_file( + schema: &SchemaRef, + batches: &[RecordBatch], + path: &Path, +) -> Result<(), ArrowError> { let format = OutputFormat::from_path(path).map_err(ArrowError::InvalidArgumentError)?; - let file = File::create(path).map_err(|e| ArrowError::IoError(e.to_string(), e))?; + + // Prefer the schema of the actual batches: some writers (e.g. Parquet) + // reject batches whose schema differs from the writer's schema, and the + // reader-reported schema may differ in metadata. + let schema = batches + .first() + .map(|b| b.schema()) + .unwrap_or_else(|| schema.clone()); + + let target = resolve_symlinks(path).map_err(io_error)?; + check_writable(&target).map_err(io_error)?; + + let mut tmp = create_temp_file(&target).map_err(io_error)?; + let file = tmp.as_file_mut(); match format { - OutputFormat::Json => write_json(batches, file), - OutputFormat::Csv => write_csv(batches, file), - OutputFormat::Arrow => write_arrow_ipc(batches, file), + OutputFormat::Json => write_json(batches, &mut *file), + OutputFormat::Jsonl => write_jsonl(batches, &mut *file), + OutputFormat::Csv => write_csv(&schema, batches, &mut *file), + OutputFormat::Arrow => write_arrow_ipc(&schema, batches, &mut *file), + OutputFormat::ArrowStream => write_arrow_stream(&schema, batches, &mut *file), + OutputFormat::Parquet => write_parquet(&schema, batches, &mut *file), + }?; + + // Make sure the data is on disk before the rename makes it visible, so a + // crash can't leave an empty or truncated file in place of the old one. + file.sync_all().map_err(io_error)?; + + // On any error above, `tmp` is dropped and the temp file is removed. + tmp.persist(&target).map_err(|e| io_error(e.error))?; + + Ok(()) +} + +fn io_error(e: std::io::Error) -> ArrowError { + ArrowError::IoError(e.to_string(), e) +} + +/// Follow `path` through any symlinks to the file that should actually be +/// replaced. The final target need not exist (a dangling symlink resolves to +/// the path it points at). +fn resolve_symlinks(path: &Path) -> std::io::Result { + // Same limit Linux uses (MAXSYMLINKS) to detect loops. + const MAX_SYMLINKS: usize = 40; + + let mut current = path.to_path_buf(); + for _ in 0..MAX_SYMLINKS { + match std::fs::symlink_metadata(¤t) { + Ok(meta) if meta.file_type().is_symlink() => { + let link = std::fs::read_link(¤t)?; + current = match current.parent() { + Some(parent) if link.is_relative() => parent.join(link), + _ => link, + }; + } + _ => return Ok(current), + } } + Err(std::io::Error::other(format!( + "Too many levels of symbolic links: {}", + path.display() + ))) } -fn write_json(batches: &[RecordBatch], file: File) -> Result<(), ArrowError> { - let mut writer = JsonWriter::<_, JsonArray>::new(file); +/// Fail early if `path` exists but can't be written, matching what +/// `File::create` would do. Without this, renaming over a read-only file +/// would succeed on Unix, since rename only needs directory permissions. +fn check_writable(path: &Path) -> std::io::Result<()> { + match std::fs::OpenOptions::new().write(true).open(path) { + Ok(_) => Ok(()), + Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(e) => Err(e), + } +} + +/// Directory in which to create the temp file for `path`: its parent, or the +/// current directory for a bare file name. +fn temp_dir_for(path: &Path) -> &Path { + match path.parent() { + Some(p) if !p.as_os_str().is_empty() => p, + _ => Path::new("."), + } +} + +/// Create a temporary file next to `path` so it can be atomically renamed +/// over the destination. +fn create_temp_file(path: &Path) -> std::io::Result { + let file_name = path.file_name().unwrap_or_default(); + let mut prefix = std::ffi::OsString::from("."); + prefix.push(file_name); + prefix.push("."); + + let mut builder = tempfile::Builder::new(); + builder.prefix(&prefix).suffix(".tmp"); + + // tempfile defaults to 0o600; match what `File::create` would produce + // (0o666 minus umask), or keep the existing file's permissions. + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + builder.permissions(std::fs::Permissions::from_mode(0o666)); + } + + let tmp = builder.tempfile_in(temp_dir_for(path))?; + + #[cfg(unix)] + if let Ok(metadata) = std::fs::metadata(path) { + std::fs::set_permissions(tmp.path(), metadata.permissions())?; + } + + Ok(tmp) +} + +fn write_json(batches: &[RecordBatch], out: W) -> Result<(), ArrowError> { + let mut writer = JsonWriter::<_, JsonArray>::new(out); for batch in batches { writer.write(batch)?; } @@ -53,21 +176,39 @@ fn write_json(batches: &[RecordBatch], file: File) -> Result<(), ArrowError> { Ok(()) } -fn write_csv(batches: &[RecordBatch], file: File) -> Result<(), ArrowError> { - let mut writer = CsvWriter::new(file); +fn write_jsonl(batches: &[RecordBatch], out: W) -> Result<(), ArrowError> { + let mut writer = JsonWriter::<_, LineDelimited>::new(out); for batch in batches { writer.write(batch)?; } + writer.finish()?; Ok(()) } -fn write_arrow_ipc(batches: &[RecordBatch], file: File) -> Result<(), ArrowError> { +fn write_csv( + schema: &SchemaRef, + batches: &[RecordBatch], + out: W, +) -> Result<(), ArrowError> { + let mut writer = CsvWriter::new(out); if batches.is_empty() { - return Ok(()); + // Writing an empty batch emits just the header row. + writer.write(&RecordBatch::new_empty(schema.clone()))?; } - let schema = batches[0].schema(); - let mut writer = IpcWriter::try_new(file, &schema)?; + for batch in batches { + writer.write(batch)?; + } + + Ok(()) +} + +fn write_arrow_ipc( + schema: &SchemaRef, + batches: &[RecordBatch], + out: W, +) -> Result<(), ArrowError> { + let mut writer = IpcWriter::try_new(out, schema)?; for batch in batches { writer.write(batch)?; } @@ -76,11 +217,43 @@ fn write_arrow_ipc(batches: &[RecordBatch], file: File) -> Result<(), ArrowError Ok(()) } +fn write_arrow_stream( + schema: &SchemaRef, + batches: &[RecordBatch], + out: W, +) -> Result<(), ArrowError> { + let mut writer = IpcStreamWriter::try_new(out, schema)?; + for batch in batches { + writer.write(batch)?; + } + writer.finish()?; + + Ok(()) +} + +fn write_parquet( + schema: &SchemaRef, + batches: &[RecordBatch], + out: W, +) -> Result<(), ArrowError> { + let props = WriterProperties::builder() + .set_compression(Compression::SNAPPY) + .build(); + let mut writer = ArrowWriter::try_new(out, schema.clone(), Some(props))?; + for batch in batches { + writer.write(batch)?; + } + writer.close()?; + + Ok(()) +} + #[cfg(test)] mod tests { use super::*; use arrow_array::{ArrayRef, Int32Array, StringArray}; use arrow_schema::{DataType, Field, Schema}; + use std::fs::File; use std::io::Read; use std::sync::Arc; use tempfile::tempdir; @@ -142,7 +315,7 @@ mod tests { let path = dir.path().join("output.json"); let batch = create_test_batch(); - write_batches_to_file(&[batch], &path).unwrap(); + write_batches_to_file(&batch.schema(), &[batch], &path).unwrap(); let mut file = File::open(&path).unwrap(); let mut contents = String::new(); @@ -161,7 +334,7 @@ mod tests { let path = dir.path().join("output.csv"); let batch = create_test_batch(); - write_batches_to_file(&[batch], &path).unwrap(); + write_batches_to_file(&batch.schema(), &[batch], &path).unwrap(); let mut file = File::open(&path).unwrap(); let mut contents = String::new(); @@ -179,7 +352,7 @@ mod tests { let path = dir.path().join("output.arrow"); let batch = create_test_batch(); - write_batches_to_file(std::slice::from_ref(&batch), &path).unwrap(); + write_batches_to_file(&batch.schema(), std::slice::from_ref(&batch), &path).unwrap(); // Verify by reading it back let file = File::open(&path).unwrap(); @@ -192,55 +365,94 @@ mod tests { } #[test] - fn test_write_arrow_ipc_empty_batches_direct() { - // Test that write_arrow_ipc handles empty batches directly (defensive coding) - // This is a regression test for the panic that would occur at src/output.rs:66 - let dir = tempdir().unwrap(); - let path = dir.path().join("output.arrow"); - let file = File::create(&path).unwrap(); + fn test_output_format_from_path_arrows() { + let path = Path::new("output.arrows"); + assert_eq!( + OutputFormat::from_path(path).unwrap(), + OutputFormat::ArrowStream + ); + } - // Direct call to write_arrow_ipc with empty batches - should not panic - let result = write_arrow_ipc(&[], file); + #[test] + fn test_output_format_from_path_parquet() { + let path = Path::new("output.parquet"); + assert_eq!( + OutputFormat::from_path(path).unwrap(), + OutputFormat::Parquet + ); + } - assert!(result.is_ok()); - // File should be created but empty since we opened it before the empty check - // Actually, let's just verify it doesn't panic and returns Ok + #[test] + fn test_output_format_from_path_jsonl() { + let path = Path::new("output.jsonl"); + assert_eq!(OutputFormat::from_path(path).unwrap(), OutputFormat::Jsonl); } #[test] - fn test_write_empty_batches_json() { + fn test_write_arrow_stream() { let dir = tempdir().unwrap(); - let path = dir.path().join("output.json"); + let path = dir.path().join("output.arrows"); + let batch = create_test_batch(); + + write_batches_to_file(&batch.schema(), std::slice::from_ref(&batch), &path).unwrap(); - // Empty batches should not create a file - write_batches_to_file(&[], &path).unwrap(); + // Verify by reading it back as an IPC stream + let file = File::open(&path).unwrap(); + let reader = arrow::ipc::reader::StreamReader::try_new(file, None).unwrap(); + let read_batches: Vec = reader.map(|r| r.unwrap()).collect(); - // Verify file was not created - assert!(!path.exists()); + assert_eq!(read_batches.len(), 1); + assert_eq!(read_batches[0].num_rows(), batch.num_rows()); + assert_eq!(read_batches[0].num_columns(), batch.num_columns()); } #[test] - fn test_write_empty_batches_arrow() { + fn test_write_parquet() { + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + let dir = tempdir().unwrap(); - let path = dir.path().join("output.arrow"); + let path = dir.path().join("output.parquet"); + let batch = create_test_batch(); - // Empty batches should not create a file - write_batches_to_file(&[], &path).unwrap(); + write_batches_to_file(&batch.schema(), std::slice::from_ref(&batch), &path).unwrap(); - // Verify file was not created - assert!(!path.exists()); + // Verify by reading it back + let file = File::open(&path).unwrap(); + let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + for column in builder.metadata().row_group(0).columns() { + assert_eq!(column.compression(), Compression::SNAPPY); + } + let reader = builder.build().unwrap(); + let read_batches: Vec = reader.map(|r| r.unwrap()).collect(); + + let total_rows: usize = read_batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, batch.num_rows()); + assert_eq!(read_batches[0].num_columns(), batch.num_columns()); } #[test] - fn test_write_empty_batches_csv() { + fn test_write_jsonl() { let dir = tempdir().unwrap(); - let path = dir.path().join("output.csv"); + let path = dir.path().join("output.jsonl"); + let batch = create_test_batch(); - // Empty batches should not create a file - write_batches_to_file(&[], &path).unwrap(); + write_batches_to_file(&batch.schema(), std::slice::from_ref(&batch), &path).unwrap(); - // Verify file was not created - assert!(!path.exists()); + let mut file = File::open(&path).unwrap(); + let mut contents = String::new(); + file.read_to_string(&mut contents).unwrap(); + + // JSON lines: one object per line, not a wrapping array + assert!(!contents.trim_start().starts_with('[')); + let lines: Vec<&str> = contents.lines().filter(|l| !l.is_empty()).collect(); + assert_eq!(lines.len(), 3); + for line in &lines { + assert!(line.trim_start().starts_with('{')); + assert!(line.trim_end().ends_with('}')); + } + assert!(contents.contains("Alice")); + assert!(contents.contains("Bob")); + assert!(contents.contains("Charlie")); } #[test] @@ -250,7 +462,7 @@ mod tests { let batch1 = create_test_batch(); let batch2 = create_test_batch(); - write_batches_to_file(&[batch1, batch2], &path).unwrap(); + write_batches_to_file(&batch1.schema(), &[batch1.clone(), batch2], &path).unwrap(); let mut file = File::open(&path).unwrap(); let mut contents = String::new(); @@ -260,4 +472,284 @@ mod tests { let count = contents.matches("Alice").count(); assert_eq!(count, 2); } + + fn read_to_string(path: &Path) -> String { + let mut contents = String::new(); + File::open(path) + .unwrap() + .read_to_string(&mut contents) + .unwrap(); + contents + } + + fn dir_entries(dir: &Path) -> Vec { + let mut names: Vec = std::fs::read_dir(dir) + .unwrap() + .map(|e| e.unwrap().file_name().to_string_lossy().into_owned()) + .collect(); + names.sort(); + names + } + + #[test] + fn test_write_zero_rows_json() { + let dir = tempdir().unwrap(); + let path = dir.path().join("output.json"); + let schema = create_test_batch().schema(); + + write_batches_to_file(&schema, &[], &path).unwrap(); + + assert_eq!(read_to_string(&path).trim(), "[]"); + } + + #[test] + fn test_write_zero_rows_jsonl() { + let dir = tempdir().unwrap(); + let path = dir.path().join("output.jsonl"); + let schema = create_test_batch().schema(); + + write_batches_to_file(&schema, &[], &path).unwrap(); + + assert!(path.exists()); + assert_eq!(read_to_string(&path), ""); + } + + #[test] + fn test_write_zero_rows_csv() { + let dir = tempdir().unwrap(); + let path = dir.path().join("output.csv"); + let schema = create_test_batch().schema(); + + write_batches_to_file(&schema, &[], &path).unwrap(); + + assert_eq!(read_to_string(&path), "id,name\n"); + } + + #[test] + fn test_write_zero_rows_arrow_ipc() { + let dir = tempdir().unwrap(); + let path = dir.path().join("output.arrow"); + let schema = create_test_batch().schema(); + + write_batches_to_file(&schema, &[], &path).unwrap(); + + let reader = + arrow::ipc::reader::FileReader::try_new(File::open(&path).unwrap(), None).unwrap(); + assert_eq!(reader.schema(), schema); + assert_eq!(reader.count(), 0); + } + + #[test] + fn test_write_zero_rows_arrow_stream() { + let dir = tempdir().unwrap(); + let path = dir.path().join("output.arrows"); + let schema = create_test_batch().schema(); + + write_batches_to_file(&schema, &[], &path).unwrap(); + + let reader = + arrow::ipc::reader::StreamReader::try_new(File::open(&path).unwrap(), None).unwrap(); + assert_eq!(reader.schema(), schema); + assert_eq!(reader.count(), 0); + } + + #[test] + fn test_write_zero_rows_parquet() { + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + + let dir = tempdir().unwrap(); + let path = dir.path().join("output.parquet"); + let schema = create_test_batch().schema(); + + write_batches_to_file(&schema, &[], &path).unwrap(); + + let builder = ParquetRecordBatchReaderBuilder::try_new(File::open(&path).unwrap()).unwrap(); + assert_eq!(builder.schema().fields(), schema.fields()); + assert_eq!(builder.metadata().file_metadata().num_rows(), 0); + } + + #[test] + fn test_write_zero_rows_replaces_existing_file() { + let dir = tempdir().unwrap(); + let path = dir.path().join("output.csv"); + let batch = create_test_batch(); + + write_batches_to_file(&batch.schema(), std::slice::from_ref(&batch), &path).unwrap(); + assert!(read_to_string(&path).contains("Alice")); + + write_batches_to_file(&batch.schema(), &[], &path).unwrap(); + assert_eq!(read_to_string(&path), "id,name\n"); + } + + /// A batch the parquet crate cannot write (Interval(MonthDayNano) is NYI). + fn create_unwritable_parquet_batch() -> RecordBatch { + use arrow_array::IntervalMonthDayNanoArray; + use arrow_array::types::IntervalMonthDayNano; + + let schema = Arc::new(Schema::new(vec![Field::new( + "iv", + DataType::Interval(arrow_schema::IntervalUnit::MonthDayNano), + false, + )])); + let array: ArrayRef = Arc::new(IntervalMonthDayNanoArray::from(vec![ + IntervalMonthDayNano::new(0, 1, 0), + ])); + RecordBatch::try_new(schema, vec![array]).unwrap() + } + + #[test] + fn test_failed_write_preserves_existing_file() { + let dir = tempdir().unwrap(); + let path = dir.path().join("keep.parquet"); + let batch = create_test_batch(); + write_batches_to_file(&batch.schema(), std::slice::from_ref(&batch), &path).unwrap(); + let original = std::fs::read(&path).unwrap(); + + let bad = create_unwritable_parquet_batch(); + let result = write_batches_to_file(&bad.schema(), &[bad], &path); + + assert!(result.is_err()); + assert_eq!(std::fs::read(&path).unwrap(), original); + // The temporary file is cleaned up. + assert_eq!(dir_entries(dir.path()), vec!["keep.parquet"]); + } + + #[test] + fn test_failed_write_does_not_create_file() { + let dir = tempdir().unwrap(); + let path = dir.path().join("new.parquet"); + + let bad = create_unwritable_parquet_batch(); + assert!(write_batches_to_file(&bad.schema(), &[bad], &path).is_err()); + assert!(dir_entries(dir.path()).is_empty()); + } + + #[test] + fn test_write_read_only_file_fails_and_preserves_it() { + let dir = tempdir().unwrap(); + let path = dir.path().join("output.csv"); + let batch = create_test_batch(); + write_batches_to_file(&batch.schema(), std::slice::from_ref(&batch), &path).unwrap(); + let original = std::fs::read(&path).unwrap(); + let mut perms = std::fs::metadata(&path).unwrap().permissions(); + perms.set_readonly(true); + std::fs::set_permissions(&path, perms).unwrap(); + if std::fs::OpenOptions::new().write(true).open(&path).is_ok() { + // Running with privileges that bypass permissions (e.g. root). + return; + } + + let result = write_batches_to_file(&batch.schema(), &[], &path); + + assert!(result.is_err()); + assert_eq!(std::fs::read(&path).unwrap(), original); + assert_eq!(dir_entries(dir.path()), vec!["output.csv"]); + } + + #[cfg(unix)] + #[test] + fn test_write_preserves_existing_permissions() { + use std::os::unix::fs::PermissionsExt; + + let dir = tempdir().unwrap(); + let path = dir.path().join("output.csv"); + let batch = create_test_batch(); + write_batches_to_file(&batch.schema(), std::slice::from_ref(&batch), &path).unwrap(); + std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o640)).unwrap(); + + write_batches_to_file(&batch.schema(), std::slice::from_ref(&batch), &path).unwrap(); + + let mode = std::fs::metadata(&path).unwrap().permissions().mode() & 0o777; + assert_eq!(mode, 0o640); + } + + #[cfg(unix)] + #[test] + fn test_write_new_file_matches_file_create_permissions() { + use std::os::unix::fs::PermissionsExt; + + let dir = tempdir().unwrap(); + let path = dir.path().join("output.csv"); + let reference = dir.path().join("reference"); + File::create(&reference).unwrap(); + let batch = create_test_batch(); + + write_batches_to_file(&batch.schema(), std::slice::from_ref(&batch), &path).unwrap(); + + // Whatever the umask, the result should match `File::create` rather + // than tempfile's owner-only default. + let mode = |p: &Path| std::fs::metadata(p).unwrap().permissions().mode() & 0o777; + assert_eq!(mode(&path), mode(&reference)); + } + + #[cfg(unix)] + #[test] + fn test_write_through_symlink_replaces_target() { + let dir = tempdir().unwrap(); + std::fs::create_dir(dir.path().join("real")).unwrap(); + let target = dir.path().join("real").join("output.csv"); + let link = dir.path().join("link.csv"); + let batch = create_test_batch(); + write_batches_to_file(&batch.schema(), &[], &target).unwrap(); + std::os::unix::fs::symlink(Path::new("real").join("output.csv"), &link).unwrap(); + + write_batches_to_file(&batch.schema(), std::slice::from_ref(&batch), &link).unwrap(); + + assert!( + std::fs::symlink_metadata(&link) + .unwrap() + .file_type() + .is_symlink() + ); + assert!(read_to_string(&target).contains("Alice")); + // The temp file was created next to the target, not the link. + assert_eq!(dir_entries(dir.path()), vec!["link.csv", "real"]); + assert_eq!(dir_entries(&dir.path().join("real")), vec!["output.csv"]); + } + + #[cfg(unix)] + #[test] + fn test_write_through_dangling_symlink_creates_target() { + let dir = tempdir().unwrap(); + let target = dir.path().join("target.csv"); + let link = dir.path().join("link.csv"); + std::os::unix::fs::symlink("target.csv", &link).unwrap(); + let batch = create_test_batch(); + + write_batches_to_file(&batch.schema(), std::slice::from_ref(&batch), &link).unwrap(); + + assert!( + std::fs::symlink_metadata(&link) + .unwrap() + .file_type() + .is_symlink() + ); + assert!(read_to_string(&target).contains("Alice")); + } + + #[cfg(unix)] + #[test] + fn test_write_through_symlink_loop_fails() { + let dir = tempdir().unwrap(); + let a = dir.path().join("a.csv"); + let b = dir.path().join("b.csv"); + std::os::unix::fs::symlink("b.csv", &a).unwrap(); + std::os::unix::fs::symlink("a.csv", &b).unwrap(); + let batch = create_test_batch(); + + let result = write_batches_to_file(&batch.schema(), &[batch], &a); + + assert!(result.is_err()); + assert_eq!(dir_entries(dir.path()), vec!["a.csv", "b.csv"]); + } + + #[test] + fn test_temp_dir_for() { + assert_eq!(temp_dir_for(Path::new("output.json")), Path::new(".")); + assert_eq!(temp_dir_for(Path::new("out/output.json")), Path::new("out")); + assert_eq!( + temp_dir_for(Path::new("/tmp/output.json")), + Path::new("/tmp") + ); + } } diff --git a/src/repl.rs b/src/repl.rs index 2e83b03..3bc0f59 100644 --- a/src/repl.rs +++ b/src/repl.rs @@ -89,7 +89,7 @@ pub fn run_repl(mut connection: impl Connection, table_mode: TableMode) { let sql = buffer.trim_end().trim_end_matches(';'); let batches = match database::execute_query(&mut connection, sql) { - Ok(batches) => batches, + Ok(result) => result.batches, Err(err) => { eprintln!("{err}"); continue; diff --git a/tests/integration_test.rs b/tests/integration_test.rs index 302c35c..eb214c7 100644 --- a/tests/integration_test.rs +++ b/tests/integration_test.rs @@ -479,3 +479,52 @@ fn test_timestamp_with_time_zone() { stdout ); } + +#[test] +fn test_invalid_output_extension_fails_before_connect() { + let output = Command::new("cargo") + .args([ + "run", + "--", + "--driver", + "nonexistent_driver", + "--query", + "SELECT 1", + "--output", + "out.xyz", + ]) + .output() + .expect("Failed to execute command"); + + // Exit code 2 and "invalid value" come from clap argument parsing, + // which happens before any connection attempt. + assert_eq!(output.status.code(), Some(2)); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!(stderr.contains("invalid value 'out.xyz' for '--output '")); + assert!(stderr.contains("Unsupported file extension")); +} + +#[test] +fn test_blank_query_with_output_fails_before_connect() { + let dir = tempfile::tempdir().unwrap(); + let out = dir.path().join("out.csv"); + let output = Command::new("cargo") + .args([ + "run", + "--", + "--driver", + "nonexistent_driver", + "--query", + " ", + "--output", + ]) + .arg(&out) + .output() + .expect("Failed to execute command"); + + assert!(!output.status.success()); + let stderr = String::from_utf8_lossy(&output.stderr); + // Rejected before trying to load the (nonexistent) driver. + assert!(stderr.contains("No query to execute"), "stderr: {stderr}"); + assert!(!out.exists()); +}