Merge pull request #20 from wangfenjin/config

feat: impl config

Author: wangfenjin

Cargo.toml

@@ -32,7 +32,6 @@ modern-full = [
 
 [dependencies]
 # time = { version = "0.3.2", features = ["formatting", "parsing"], optional = true }
-bitflags = "1.2"
 hashlink = "0.7"
 chrono = { version = "0.4", optional = true }
 serde_json = { version = "1.0", optional = true }
@@ -48,6 +47,7 @@ smallvec = "1.6.1"
 cast = { version = "0.2", features = ["std", "x128"] }
 arrow = { version = "5.0", default-features = false, features = ["prettyprint"] }
 rust_decimal = "1.14"
+strum = { version = "0.21", features = ["derive"] }
 
 [dev-dependencies]
 doc-comment = "0.3"

src/config.rs

@@ -0,0 +1,205 @@
+use super::ffi;
+use super::Result;
+use crate::error::Error;
+use std::default::Default;
+use std::ffi::CString;
+use std::os::raw::c_char;
+use std::ptr;
+
+use strum::{EnumString, ToString};
+
+/// duckdb access mode, default is Automatic
+#[derive(Debug, Eq, PartialEq, EnumString, ToString)]
+pub enum AccessMode {
+    /// Access mode of the database AUTOMATIC
+    #[strum(to_string = "AUTOMATIC")]
+    Automatic,
+    /// Access mode of the database READ_ONLY
+    #[strum(to_string = "READ_ONLY")]
+    ReadOnly,
+    /// Access mode of the database READ_WRITE
+    #[strum(to_string = "READ_WRITE")]
+    ReadWrite,
+}
+
+/// duckdb default order, default is Asc
+#[derive(Debug, Eq, PartialEq, EnumString, ToString)]
+pub enum DefaultOrder {
+    /// The order type, ASC
+    #[strum(to_string = "ASC")]
+    Asc,
+    /// The order type, DESC
+    #[strum(to_string = "DESC")]
+    Desc,
+}
+
+/// duckdb default null order, default is nulls first
+#[derive(Debug, Eq, PartialEq, EnumString, ToString)]
+pub enum DefaultNullOrder {
+    /// Null ordering, NullsFirst
+    #[strum(to_string = "NULLS_FIRST")]
+    NullsFirst,
+    /// Null ordering, NullsLast
+    #[strum(to_string = "NULLS_LAST")]
+    NullsLast,
+}
+
+/// duckdb configuration
+/// Refer to https://github.com/duckdb/duckdb/blob/master/src/main/config.cpp
+pub struct Config {
+    config: Option<ffi::duckdb_config>,
+}
+
+impl Config {
+    pub(crate) fn duckdb_config(&self) -> ffi::duckdb_config {
+        self.config.unwrap_or(std::ptr::null_mut() as ffi::duckdb_config)
+    }
+
+    /// Access mode of the database ([AUTOMATIC], READ_ONLY or READ_WRITE)
+    pub fn access_mode(mut self, mode: AccessMode) -> Result<Config> {
+        self.set("access_mode", &mode.to_string())?;
+        Ok(self)
+    }
+
+    /// The order type used when none is specified ([ASC] or DESC)
+    pub fn default_order(mut self, order: DefaultOrder) -> Result<Config> {
+        self.set("default_order", &order.to_string())?;
+        Ok(self)
+    }
+
+    /// Null ordering used when none is specified ([NULLS_FIRST] or NULLS_LAST)
+    pub fn default_null_order(mut self, null_order: DefaultNullOrder) -> Result<Config> {
+        self.set("default_null_order", &null_order.to_string())?;
+        Ok(self)
+    }
+
+    /// Allow the database to access external state (through e.g. COPY TO/FROM, CSV readers, pandas replacement scans, etc)
+    pub fn enable_external_access(mut self, enabled: bool) -> Result<Config> {
+        self.set("enable_external_access", &enabled.to_string())?;
+        Ok(self)
+    }
+
+    /// Whether or not object cache is used to cache e.g. Parquet metadata
+    pub fn enable_object_cache(mut self, enabled: bool) -> Result<Config> {
+        self.set("enable_object_cache", &enabled.to_string())?;
+        Ok(self)
+    }
+
+    /// The maximum memory of the system (e.g. 1GB)
+    pub fn max_memory(mut self, memory: &str) -> Result<Config> {
+        self.set("max_memory", memory)?;
+        Ok(self)
+    }
+
+    /// The number of total threads used by the system
+    pub fn threads(mut self, thread_num: i64) -> Result<Config> {
+        self.set("threads", &thread_num.to_string())?;
+        Ok(self)
+    }
+
+    fn set(&mut self, key: &str, value: &str) -> Result<()> {
+        if self.config.is_none() {
+            let mut config: ffi::duckdb_config = ptr::null_mut();
+            let state = unsafe { ffi::duckdb_create_config(&mut config) };
+            assert_eq!(state, ffi::DuckDBSuccess);
+            self.config = Some(config);
+        }
+        let c_key = CString::new(key).unwrap();
+        let c_value = CString::new(value).unwrap();
+        let state = unsafe {
+            ffi::duckdb_set_config(
+                self.config.unwrap(),
+                c_key.as_ptr() as *const c_char,
+                c_value.as_ptr() as *const c_char,
+            )
+        };
+        if state != ffi::DuckDBSuccess {
+            return Err(Error::DuckDBFailure(
+                ffi::Error::new(state as u32),
+                Some(format!("set {}:{} error", key, value)),
+            ));
+        }
+        Ok(())
+    }
+}
+
+impl Default for Config {
+    fn default() -> Config {
+        Config { config: None }
+    }
+}
+
+impl Drop for Config {
+    fn drop(&mut self) {
+        if self.config.is_some() {
+            unsafe { ffi::duckdb_destroy_config(&mut self.config.unwrap()) };
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use crate::types::Value;
+    use crate::{Config, Connection, Result};
+
+    #[test]
+    fn test_default_config() -> Result<()> {
+        let config = Config::default();
+        let db = Connection::open_in_memory_with_flags(config)?;
+        db.execute_batch("CREATE TABLE foo(x Text)")?;
+
+        let mut stmt = db.prepare("INSERT INTO foo(x) VALUES (?)")?;
+        stmt.execute(&[&"a"])?;
+        stmt.execute(&[&"b"])?;
+        stmt.execute(&[&"c"])?;
+        stmt.execute([Value::Null])?;
+
+        let val: Result<Vec<Option<String>>> = db
+            .prepare("SELECT x FROM foo ORDER BY x")?
+            .query_and_then([], |row| row.get(0))?
+            .collect();
+        let val = val?;
+        let mut iter = val.iter();
+        assert!(iter.next().unwrap().is_none());
+        assert_eq!(iter.next().unwrap().as_ref().unwrap(), "a");
+        assert_eq!(iter.next().unwrap().as_ref().unwrap(), "b");
+        assert_eq!(iter.next().unwrap().as_ref().unwrap(), "c");
+        assert_eq!(iter.next(), None);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_all_config() -> Result<()> {
+        let config = Config::default()
+            .access_mode(crate::AccessMode::ReadWrite)?
+            .default_null_order(crate::DefaultNullOrder::NullsLast)?
+            .default_order(crate::DefaultOrder::Desc)?
+            .enable_external_access(true)?
+            .enable_object_cache(false)?
+            .max_memory("2GB")?
+            .threads(4)?;
+        let db = Connection::open_in_memory_with_flags(config)?;
+        db.execute_batch("CREATE TABLE foo(x Text)")?;
+
+        let mut stmt = db.prepare("INSERT INTO foo(x) VALUES (?)")?;
+        stmt.execute(&[&"a"])?;
+        stmt.execute(&[&"b"])?;
+        stmt.execute(&[&"c"])?;
+        stmt.execute([Value::Null])?;
+
+        let val: Result<Vec<Option<String>>> = db
+            .prepare("SELECT x FROM foo ORDER BY x")?
+            .query_and_then([], |row| row.get(0))?
+            .collect();
+        let val = val?;
+        let mut iter = val.iter();
+        assert_eq!(iter.next().unwrap().as_ref().unwrap(), "c");
+        assert_eq!(iter.next().unwrap().as_ref().unwrap(), "b");
+        assert_eq!(iter.next().unwrap().as_ref().unwrap(), "a");
+        assert!(iter.next().unwrap().is_none());
+        assert_eq!(iter.next(), None);
+
+        Ok(())
+    }
+}

src/inner_connection.rs

@@ -1,19 +1,18 @@
-use std::ffi::{CStr, CString};
+use std::ffi::{c_void, CStr, CString};
 use std::mem;
 use std::os::raw::c_char;
 use std::ptr;
 use std::str;
 
 use super::ffi;
-use super::{Appender, Connection, OpenFlags, Result};
+use super::{Appender, Config, Connection, Result};
 use crate::error::{result_from_duckdb_appender, result_from_duckdb_arrow, result_from_duckdb_prepare, Error};
 use crate::raw_statement::RawStatement;
 use crate::statement::Statement;
 
 pub struct InnerConnection {
     pub db: ffi::duckdb_database,
     pub con: ffi::duckdb_connection,
-    // pub result: ffi::duckdb_result,
     owned: bool,
 }
 
@@ -35,25 +34,18 @@ impl InnerConnection {
             // TODO: fix this
             panic!("error {:?}", e);
         }
-        InnerConnection {
-            db,
-            con,
-            // result: mem::zeroed(),
-            owned,
-        }
+        InnerConnection { db, con, owned }
     }
 
-    pub fn open_with_flags(c_path: &CStr, _: OpenFlags, _: Option<&CStr>) -> Result<InnerConnection> {
+    pub fn open_with_flags(c_path: &CStr, config: Config) -> Result<InnerConnection> {
         unsafe {
             let mut db: ffi::duckdb_database = ptr::null_mut();
-            let r = ffi::duckdb_open(c_path.as_ptr(), &mut db);
+            let mut c_err = std::ptr::null_mut();
+            let r = ffi::duckdb_open_ext(c_path.as_ptr(), &mut db, config.duckdb_config(), &mut c_err);
             if r != ffi::DuckDBSuccess {
-                ffi::duckdb_close(&mut db);
-                let e = Error::DuckDBFailure(
-                    ffi::Error::new(r),
-                    Some(format!("{}: {}", "duckdb_open error", c_path.to_string_lossy())),
-                );
-                return Err(e);
+                let msg = Some(CStr::from_ptr(c_err).to_string_lossy().to_string());
+                ffi::duckdb_free(c_err as *mut c_void);
+                return Err(Error::DuckDBFailure(ffi::Error::new(r as u32), msg));
             }
             Ok(InnerConnection::new(db, true))
         }