diff --git a/Cargo.toml b/Cargo.toml index 5e27fc1..92affd8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -25,4 +25,4 @@ tempdir = "0.3.7" [package.metadata.docs.rs] # docs.rs cannot download libchdb or link native libraries. # Setting DOCS_RS env var tells build.rs to skip native build steps. -rustc-args = ["--cfg", "docsrs"] \ No newline at end of file +rustc-args = ["--cfg", "docsrs"] diff --git a/src/arg.rs b/src/arg.rs index 1b7bda3..be99819 100644 --- a/src/arg.rs +++ b/src/arg.rs @@ -4,9 +4,8 @@ //! log level, and custom command-line arguments. use std::borrow::Cow; -use std::ffi::CString; +use std::fmt::Display; -use crate::error::Error; use crate::format::OutputFormat; use crate::log_level::LogLevel; @@ -35,12 +34,31 @@ use crate::log_level::LogLevel; #[derive(Debug)] pub enum Arg<'a> { /// `--config-file=` + /// + /// Can be used to specify a custom configuration + /// file for the session, allowing one to configure various aspects of the + /// session's behavior (e.g., SSL settings to use with the `remoteSecure` + /// function). + /// + /// An example SSL configuration file (`chdb_ssl.xml`) might look like this: + /// ```xml + /// + /// + /// + /// path_to_server_ca_cert.pem + /// + /// + /// + /// ``` + /// where `path_to_server_ca_cert.pem` is the path to the CA certificate + /// of the remote server. ConfigFilePath(Cow<'a, str>), /// `--log-level=` LogLevel(LogLevel), /// `--output-format=` OutputFormat(OutputFormat), /// --multiquery + /// Emitted as `-n` (short for `--multiquery`). MultiQuery, /// Custom argument. /// @@ -52,26 +70,12 @@ pub enum Arg<'a> { /// 1. Arg::Custom("multiline".to_string().into(), None). /// 2. Arg::Custom("multiline".into(), None). /// - /// We should tell user where to look for officially supported arguments. - /// Here is some hint for now: . + /// Officially supported arguments can be found by running + /// `clickhouse-client --help` in the terminal. Custom(Cow<'a, str>, Option>), } impl<'a> Arg<'a> { - #[allow(dead_code)] - pub(crate) fn to_cstring(&self) -> Result { - Ok(match self { - Self::ConfigFilePath(v) => CString::new(format!("--config-file={v}")), - Self::LogLevel(v) => CString::new(format!("--log-level={}", v.as_str())), - Self::OutputFormat(v) => CString::new(format!("--output-format={}", v.as_str())), - Self::MultiQuery => CString::new("-n"), - Self::Custom(k, v) => match v { - None => CString::new(k.as_ref()), - Some(v) => CString::new(format!("--{k}={v}")), - }, - }?) - } - /// Extract `OutputFormat` from an `Arg` if it is an `OutputFormat` variant. /// /// This is a helper method used internally to extract output format information @@ -84,6 +88,21 @@ impl<'a> Arg<'a> { } } +impl Display for Arg<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::ConfigFilePath(v) => write!(f, "--config-file={v}"), + Self::LogLevel(v) => write!(f, "--log-level={}", v.as_str()), + Self::OutputFormat(v) => write!(f, "--output-format={}", v.as_str()), + Self::MultiQuery => write!(f, "-n"), + Self::Custom(k, v) => match v { + None => write!(f, "--{}", k.as_ref()), + Some(v) => write!(f, "--{k}={v}"), + }, + } + } +} + /// Extract `OutputFormat` from a slice of `Arg`s. /// /// This function searches through the provided arguments and returns the first @@ -96,7 +115,73 @@ impl<'a> Arg<'a> { /// # Returns /// /// Returns the first `OutputFormat` found, or `OutputFormat::TabSeparated` as default. -pub(crate) fn extract_output_format(args: Option<&[Arg]>) -> OutputFormat { +pub(crate) fn extract_output_format(args: Option<&[Arg]>, default: OutputFormat) -> OutputFormat { args.and_then(|args| args.iter().find_map(|a| a.as_output_format())) - .unwrap_or(OutputFormat::TabSeparated) + .unwrap_or(default) +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_arg_display_config_file_path() { + assert_eq!( + Arg::ConfigFilePath(Cow::from("my.xml")).to_string(), + "--config-file=my.xml" + ); + } + + #[test] + fn test_arg_display_log_level() { + assert_eq!( + Arg::LogLevel(LogLevel::Trace).to_string(), + "--log-level=trace" + ); + assert_eq!( + Arg::LogLevel(LogLevel::Debug).to_string(), + "--log-level=debug" + ); + assert_eq!( + Arg::LogLevel(LogLevel::Info).to_string(), + "--log-level=information" + ); + assert_eq!( + Arg::LogLevel(LogLevel::Warn).to_string(), + "--log-level=warning" + ); + assert_eq!( + Arg::LogLevel(LogLevel::Error).to_string(), + "--log-level=error" + ); + } + + #[test] + fn test_arg_display_output_format() { + assert_eq!( + Arg::OutputFormat(OutputFormat::JSONEachRow).to_string(), + "--output-format=JSONEachRow" + ); + } + + #[test] + fn test_arg_display_multi_query() { + assert_eq!(Arg::MultiQuery.to_string(), "-n"); + } + + #[test] + fn test_arg_display_custom_key_only() { + assert_eq!( + Arg::Custom("multiline".into(), None).to_string(), + "--multiline" + ); + } + + #[test] + fn test_arg_display_custom_key_value() { + assert_eq!( + Arg::Custom("priority".into(), Some("1".into())).to_string(), + "--priority=1" + ); + } } diff --git a/src/connection.rs b/src/connection.rs index b6534f0..99f0a54 100644 --- a/src/connection.rs +++ b/src/connection.rs @@ -5,10 +5,10 @@ use std::ffi::{c_char, CString}; use crate::arrow_stream::{ArrowArray, ArrowSchema, ArrowStream}; -use crate::bindings; use crate::error::{Error, Result}; use crate::format::OutputFormat; use crate::query_result::QueryResult; +use crate::{bindings, CHDB_PROGRAM_NAME}; /// A connection to a chDB database. /// @@ -49,13 +49,12 @@ unsafe impl Send for Connection {} impl Connection { /// Connect to chDB with the given command-line arguments. /// - /// This is a low-level function that allows you to pass arbitrary arguments - /// to the chDB connection. For most use cases, prefer [`open_in_memory`](Self::open_in_memory) - /// or [`open_with_path`](Self::open_with_path). + /// Use [crate::session::SessionBuilder] for a higher-level API that supports + /// sessions and persistent storage. /// /// # Arguments /// - /// * `args` - Array of command-line arguments (e.g., `["clickhouse", "--path=/tmp/db"]`) + /// * `args` - Array of command-line arguments (e.g., `["--path=/tmp/db"]`) /// /// # Examples /// @@ -63,7 +62,7 @@ impl Connection { /// use chdb_rust::connection::Connection; /// /// // Connect with custom arguments - /// let conn = Connection::open(&["clickhouse", "--path=/tmp/mydb"])?; + /// let conn = Connection::open(&["--path=/tmp/mydb"])?; /// # Ok::<(), chdb_rust::error::Error>(()) /// ``` /// @@ -72,14 +71,14 @@ impl Connection { /// Returns [`Error::ConnectionFailed`] if the /// connection cannot be established. pub fn open(args: &[&str]) -> Result { - let c_args: Vec = args - .iter() - .map(|s| CString::new(*s)) - .collect::, _>>()?; + let c_args: Vec = std::iter::once(CHDB_PROGRAM_NAME) + .chain(args.iter().copied()) + .map(CString::new) + .collect::>()?; - let mut argv: Vec<*mut c_char> = c_args.iter().map(|s| s.as_ptr() as *mut c_char).collect(); - - let conn_ptr = unsafe { bindings::chdb_connect(argv.len() as i32, argv.as_mut_ptr()) }; + let argv: Vec<*const c_char> = c_args.iter().map(|s| s.as_ptr()).collect(); + let conn_ptr = + unsafe { bindings::chdb_connect(argv.len() as i32, argv.as_ptr() as *mut *mut c_char) }; if conn_ptr.is_null() { return Err(Error::ConnectionFailed); @@ -113,7 +112,7 @@ impl Connection { /// Returns [`Error::ConnectionFailed`] if the /// connection cannot be established. pub fn open_in_memory() -> Result { - Self::open(&["clickhouse"]) + Self::open(&[]) } /// Connect to a database at the given path. @@ -138,9 +137,10 @@ impl Connection { /// /// Returns [`Error::ConnectionFailed`] if the /// connection cannot be established. + #[deprecated(note = "Use `SessionBuilder` instead")] pub fn open_with_path(path: &str) -> Result { let path_arg = format!("--path={path}"); - Self::open(&["clickhouse", &path_arg]) + Self::open(&[&path_arg]) } /// Execute a query and return the result. @@ -383,3 +383,26 @@ impl Drop for Connection { } } } + +#[cfg(test)] +mod test { + use super::*; + use crate::{error::Result, test_utils::tempdir}; + + #[test] + fn test_connection_open_with_explicit_data_path() -> Result<()> { + let tmp = tempdir(); + let path_arg = format!( + "--path={}", + tmp.path().to_str().expect("temp path is not valid UTF-8") + ); + Connection::open(&[&path_arg])?; + + assert!( + tmp.path().read_dir()?.next().is_some(), + "expected chDB to create files in the data dir" + ); + + Ok(()) + } +} diff --git a/src/lib.rs b/src/lib.rs index c96ce38..a08cf58 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -53,11 +53,17 @@ pub mod log_level; pub mod query_result; pub mod session; +#[cfg(test)] +mod test_utils; + use crate::arg::{extract_output_format, Arg}; use crate::connection::Connection; use crate::error::Result; +use crate::format::OutputFormat; use crate::query_result::QueryResult; +pub(crate) const CHDB_PROGRAM_NAME: &str = "clickhouse"; + /// Execute a one-off query using an in-memory connection. /// /// This function creates a temporary in-memory database connection, executes the query, @@ -101,6 +107,6 @@ use crate::query_result::QueryResult; /// - The query execution fails pub fn execute(query: &str, query_args: Option<&[Arg]>) -> Result { let conn = Connection::open_in_memory()?; - let fmt = extract_output_format(query_args); + let fmt = extract_output_format(query_args, OutputFormat::TabSeparated); conn.query(query, fmt) } diff --git a/src/session.rs b/src/session.rs index 673e03b..49fb74c 100644 --- a/src/session.rs +++ b/src/session.rs @@ -3,12 +3,12 @@ //! This module provides the [`Session`] and [`SessionBuilder`] types for managing //! persistent database connections with automatic cleanup. -use std::fs; use std::path::PathBuf; +use std::{fs, io}; -use crate::arg::Arg; +use crate::arg::{extract_output_format, Arg}; use crate::connection::Connection; -use crate::error::Error; +use crate::error::{Error, Result}; use crate::format::OutputFormat; use crate::query_result::QueryResult; @@ -33,7 +33,7 @@ use crate::query_result::QueryResult; pub struct SessionBuilder<'a> { data_path: PathBuf, default_format: OutputFormat, - _marker: std::marker::PhantomData<&'a ()>, + arguments: Vec>, auto_cleanup: bool, } @@ -81,7 +81,7 @@ pub struct SessionBuilder<'a> { #[derive(Debug)] pub struct Session { conn: Connection, - data_path: String, + data_path: PathBuf, default_format: OutputFormat, auto_cleanup: bool, } @@ -90,7 +90,8 @@ impl<'a> SessionBuilder<'a> { /// Create a new `SessionBuilder` with default settings. /// /// The default settings are: - /// - Data path: `./chdb` in the current working directory + /// - Data path: current working directory. The `build` method will create a `chdb` + /// subdirectory if the path is not set. /// - Output format: `TabSeparated` /// - Auto cleanup: `false` /// @@ -103,13 +104,10 @@ impl<'a> SessionBuilder<'a> { /// # Ok::<(), chdb_rust::error::Error>(()) /// ``` pub fn new() -> Self { - let mut data_path = std::env::current_dir().unwrap(); - data_path.push("chdb"); - Self { - data_path, + data_path: PathBuf::new(), default_format: OutputFormat::TabSeparated, - _marker: std::marker::PhantomData, + arguments: Vec::new(), auto_cleanup: false, } } @@ -139,12 +137,15 @@ impl<'a> SessionBuilder<'a> { /// Add a query argument to the session builder. /// - /// Currently, only `OutputFormat` arguments are supported and will be used + /// If `OutputFormat` argument is provided it will be used /// as the default output format for queries executed on this session. /// + /// Moreover, `OutputFormat` is consumed by the session layer and is not + /// forwarded as a command-line argument to the connection. + /// /// # Arguments /// - /// * `arg` - The argument to add (currently only `OutputFormat` is supported) + /// * `arg` - The argument to add /// /// # Examples /// @@ -152,15 +153,22 @@ impl<'a> SessionBuilder<'a> { /// use chdb_rust::session::SessionBuilder; /// use chdb_rust::arg::Arg; /// use chdb_rust::format::OutputFormat; + /// use std::borrow::Cow; /// /// let builder = SessionBuilder::new() - /// .with_arg(Arg::OutputFormat(OutputFormat::JSONEachRow)); + /// .with_arg(Arg::OutputFormat(OutputFormat::JSONEachRow)) + /// .with_arg(Arg::ConfigFilePath(Cow::from("chdb_ssl.xml"))) + /// .with_arg(Arg::Custom(Cow::from("progress"), Some(Cow::from("err")))); /// # Ok::<(), chdb_rust::error::Error>(()) /// ``` + /// Custom arguments can be used to specify any additional command-line arguments + /// that can be found by running `clickhouse-client --help` in + /// the terminal. pub fn with_arg(mut self, arg: Arg<'a>) -> Self { - // Only OutputFormat is supported with the new API if let Some(fmt) = arg.as_output_format() { self.default_format = fmt; + } else { + self.arguments.push(arg); } self } @@ -194,7 +202,8 @@ impl<'a> SessionBuilder<'a> { /// Build the session with the configured settings. /// /// This creates the data directory if it doesn't exist and establishes - /// a connection to the database. + /// a connection to the database. All configured arguments are passed to + /// the connection. /// /// # Returns /// @@ -218,19 +227,44 @@ impl<'a> SessionBuilder<'a> { /// .build()?; /// # Ok::<(), chdb_rust::error::Error>(()) /// ``` - pub fn build(self) -> Result { - let data_path = self.data_path.to_str().ok_or(Error::PathError)?.to_string(); + pub fn build(mut self) -> Result { + if self.data_path.as_os_str().is_empty() { + let mut default_path = std::env::current_dir()?; + default_path.push("chdb"); + self.data_path = default_path; + } + let dir_already_existed = self.data_path.exists(); fs::create_dir_all(&self.data_path)?; - if fs::metadata(&self.data_path)?.permissions().readonly() { - return Err(Error::InsufficientPermissions); + let probe_path = self.data_path.join(".write_probe"); + if let Err(e) = fs::OpenOptions::new() + .write(true) + .create(true) + .truncate(true) + .open(&probe_path) + { + if !dir_already_existed { + fs::remove_dir(&self.data_path).ok(); + } + return Err(match e.kind() { + io::ErrorKind::PermissionDenied => Error::InsufficientPermissions, + _ => Error::Io(e), + }); } + fs::remove_file(&probe_path).ok(); // best-effort cleanup of write probe + + let data_path = self.data_path.to_str().ok_or(Error::PathError)?; + let path_arg = format!("--path={data_path}"); + let owned: Vec = self.arguments.iter().map(|a| a.to_string()).collect(); + let mut args: Vec<&str> = Vec::with_capacity(1 + owned.len()); + args.push(&path_arg); + args.extend(owned.iter().map(String::as_str)); - let conn = Connection::open_with_path(&data_path)?; + let conn = Connection::open(&args)?; Ok(Session { conn, - data_path, + data_path: self.data_path, default_format: self.default_format, auto_cleanup: self.auto_cleanup, }) @@ -252,7 +286,10 @@ impl Session { /// # Arguments /// /// * `query` - The SQL query string to execute - /// * `query_args` - Optional array of query arguments (e.g., output format) + /// * `query_args` - Optional array of query arguments (e.g., output format). + /// + /// Only `OutputFormat` is currently supported and will override the + /// session's default output format for this query. /// /// # Returns /// @@ -291,10 +328,8 @@ impl Session { /// - The query syntax is invalid /// - The query references non-existent tables or columns /// - The query execution fails for any other reason - pub fn execute(&self, query: &str, query_args: Option<&[Arg]>) -> Result { - let fmt = query_args - .and_then(|args| args.iter().find_map(|a| a.as_output_format())) - .unwrap_or(self.default_format); + pub fn execute(&self, query: &str, query_args: Option<&[Arg]>) -> Result { + let fmt = extract_output_format(query_args, self.default_format); self.conn.query(query, fmt) } } @@ -306,3 +341,47 @@ impl Drop for Session { } } } + +#[cfg(test)] +mod test { + use crate::test_utils::tempdir; + + use super::*; + + #[test] + fn test_session_builder_no_args_still_builds() -> Result<()> { + let tmp = tempdir(); + SessionBuilder::new().with_data_path(tmp.path()).build()?; + Ok(()) + } + + #[test] + fn test_query_uses_default_output_format() -> Result<()> { + let tmp = tempdir(); + let session = SessionBuilder::new().with_data_path(tmp.path()).build()?; + + let resp = session.execute("SELECT 'foo' AS name, 1 AS count", None)?; + + assert_eq!(resp.data_utf8_lossy(), "foo\t1\n"); + + Ok(()) + } + + #[test] + fn test_query_output_format_overrides_session_builder_output_format() -> Result<()> { + let tmp = tempdir(); + let session = SessionBuilder::new() + .with_data_path(tmp.path()) + .with_arg(Arg::OutputFormat(OutputFormat::Parquet)) + .build()?; + + let resp = session.execute( + "SELECT 1 AS count", + Some(&[Arg::OutputFormat(OutputFormat::JSONEachRow)]), + )?; + + assert_eq!(resp.data_utf8_lossy(), "{\"count\":1}\n"); + + Ok(()) + } +} diff --git a/src/test_utils.rs b/src/test_utils.rs new file mode 100644 index 0000000..2263d9a --- /dev/null +++ b/src/test_utils.rs @@ -0,0 +1,4 @@ +#[cfg(test)] +pub(crate) fn tempdir() -> tempdir::TempDir { + tempdir::TempDir::new("chdb-rust").expect("failed to create temp dir") +} diff --git a/tests/common/mod.rs b/tests/common/mod.rs new file mode 100644 index 0000000..4ed0f76 --- /dev/null +++ b/tests/common/mod.rs @@ -0,0 +1,3 @@ +pub fn tempdir() -> tempdir::TempDir { + tempdir::TempDir::new("chdb-rust").expect("failed to create temp dir") +} diff --git a/tests/examples.rs b/tests/examples.rs index ec3df0b..bd23463 100644 --- a/tests/examples.rs +++ b/tests/examples.rs @@ -4,6 +4,7 @@ //! in the chDB library. Run with `RUST_TEST_THREADS=1 cargo test --test examples` //! to ensure reliable execution. +mod common; use chdb_rust::arg::Arg; use chdb_rust::error::Result; use chdb_rust::execute; @@ -18,7 +19,7 @@ fn test_stateful() -> Result<()> { // // Create session. // - let tmp = tempdir::TempDir::new("chdb-rust")?; + let tmp = common::tempdir(); let session = SessionBuilder::new() .with_data_path(tmp.path()) .with_arg(Arg::LogLevel(LogLevel::Debug)) @@ -181,7 +182,7 @@ fn test_query_result_data_methods() -> Result<()> { #[test] fn test_multiple_inserts_and_aggregation() -> Result<()> { - let tmp = tempdir::TempDir::new("chdb-rust")?; + let tmp = common::tempdir(); let session = SessionBuilder::new() .with_data_path(tmp.path()) .with_auto_cleanup(true) @@ -233,7 +234,7 @@ fn test_multiple_inserts_and_aggregation() -> Result<()> { #[test] fn test_different_data_types() -> Result<()> { - let tmp = tempdir::TempDir::new("chdb-rust")?; + let tmp = common::tempdir(); let session = SessionBuilder::new() .with_data_path(tmp.path()) .with_auto_cleanup(true) @@ -307,7 +308,7 @@ fn test_error_handling_invalid_syntax() { #[test] fn test_session_auto_cleanup() -> Result<()> { - let tmp = tempdir::TempDir::new("chdb-rust")?; + let tmp = common::tempdir(); let data_path = tmp.path().to_path_buf(); { @@ -342,7 +343,7 @@ fn test_session_auto_cleanup() -> Result<()> { #[test] fn test_complex_query_with_joins() -> Result<()> { - let tmp = tempdir::TempDir::new("chdb-rust")?; + let tmp = common::tempdir(); let session = SessionBuilder::new() .with_data_path(tmp.path()) .with_auto_cleanup(true) @@ -427,7 +428,7 @@ fn test_default_output_format() -> Result<()> { #[test] fn test_session_without_auto_cleanup() -> Result<()> { - let tmp = tempdir::TempDir::new("chdb-rust")?; + let tmp = common::tempdir(); let session = SessionBuilder::new() .with_data_path(tmp.path()) .with_auto_cleanup(false) // Explicitly disable cleanup