duckdb/
config.rs

1use super::{Result, ffi};
2use crate::error::Error;
3use std::{default::Default, ffi::CString, os::raw::c_char, ptr};
4
5use strum::{AsRefStr, Display, EnumString};
6
7/// duckdb access mode, default is Automatic
8#[derive(Debug, Eq, PartialEq, Clone, EnumString, Display, AsRefStr)]
9pub enum AccessMode {
10    /// Access mode of the database AUTOMATIC
11    #[strum(to_string = "AUTOMATIC")]
12    Automatic,
13    /// Access mode of the database READ_ONLY
14    #[strum(to_string = "READ_ONLY")]
15    ReadOnly,
16    /// Access mode of the database READ_WRITE
17    #[strum(to_string = "READ_WRITE")]
18    ReadWrite,
19}
20
21/// duckdb default order, default is Asc
22#[derive(Debug, Eq, PartialEq, Clone, EnumString, Display, AsRefStr)]
23pub enum DefaultOrder {
24    /// The order type, ASC
25    #[strum(to_string = "ASC")]
26    Asc,
27    /// The order type, DESC
28    #[strum(to_string = "DESC")]
29    Desc,
30}
31
32/// duckdb default null order, default is nulls first
33#[derive(Debug, Eq, PartialEq, Clone, EnumString, Display, AsRefStr)]
34pub enum DefaultNullOrder {
35    /// Null ordering, NullsFirst
36    #[strum(to_string = "NULLS_FIRST")]
37    NullsFirst,
38    /// Null ordering, NullsLast
39    #[strum(to_string = "NULLS_LAST")]
40    NullsLast,
41}
42
43/// duckdb configuration
44/// Refer to <https://github.com/duckdb/duckdb/blob/master/src/main/config.cpp>
45#[derive(Default)]
46pub struct Config {
47    config: Option<ffi::duckdb_config>,
48}
49
50impl Config {
51    pub(crate) fn duckdb_config(&self) -> ffi::duckdb_config {
52        self.config.unwrap_or(std::ptr::null_mut() as ffi::duckdb_config)
53    }
54
55    /// enable autoload extensions
56    pub fn enable_autoload_extension(mut self, enabled: bool) -> Result<Self> {
57        let value = (enabled as u8).to_string();
58        self.set("autoinstall_known_extensions", &value)?;
59        self.set("autoload_known_extensions", &value)?;
60        Ok(self)
61    }
62
63    /// Access mode of the database (`AUTOMATIC`, `READ_ONLY` or `READ_WRITE`)
64    pub fn access_mode(mut self, mode: AccessMode) -> Result<Self> {
65        self.set("access_mode", mode)?;
66        Ok(self)
67    }
68
69    /// Metadata from DuckDB callers
70    pub fn custom_user_agent(mut self, custom_user_agent: &str) -> Result<Self> {
71        self.set("custom_user_agent", custom_user_agent)?;
72        Ok(self)
73    }
74
75    /// The order type used when none is specified (`ASC` or `DESC`)
76    pub fn default_order(mut self, order: DefaultOrder) -> Result<Self> {
77        self.set("default_order", order)?;
78        Ok(self)
79    }
80
81    /// Null ordering used when none is specified (`NULLS_FIRST` or `NULLS_LAST`)
82    pub fn default_null_order(mut self, null_order: DefaultNullOrder) -> Result<Self> {
83        self.set("default_null_order", null_order)?;
84        Ok(self)
85    }
86
87    /// Allow the database to access external state (through e.g. COPY TO/FROM, CSV readers, pandas replacement scans, etc)
88    pub fn enable_external_access(mut self, enabled: bool) -> Result<Self> {
89        self.set("enable_external_access", enabled.to_string())?;
90        Ok(self)
91    }
92
93    /// Whether or not object cache is used to cache e.g. Parquet metadata
94    pub fn enable_object_cache(mut self, enabled: bool) -> Result<Self> {
95        self.set("enable_object_cache", enabled.to_string())?;
96        Ok(self)
97    }
98
99    /// Allow to load third-party duckdb extensions.
100    pub fn allow_unsigned_extensions(mut self) -> Result<Self> {
101        self.set("allow_unsigned_extensions", "true")?;
102        Ok(self)
103    }
104
105    /// The maximum memory of the system (e.g. 1GB)
106    pub fn max_memory(mut self, memory: &str) -> Result<Self> {
107        self.set("max_memory", memory)?;
108        Ok(self)
109    }
110
111    /// The number of total threads used by the system
112    pub fn threads(mut self, thread_num: i64) -> Result<Self> {
113        self.set("threads", thread_num.to_string())?;
114        Ok(self)
115    }
116
117    /// Add any setting to the config. DuckDB will return an error if the setting is unknown or otherwise invalid.
118    pub fn with(mut self, key: impl AsRef<str>, value: impl AsRef<str>) -> Result<Self> {
119        self.set(key, value)?;
120        Ok(self)
121    }
122
123    fn set(&mut self, key: impl AsRef<str>, value: impl AsRef<str>) -> Result<()> {
124        let key = key.as_ref();
125        let value = value.as_ref();
126        if self.config.is_none() {
127            let mut config: ffi::duckdb_config = ptr::null_mut();
128            let state = unsafe { ffi::duckdb_create_config(&mut config) };
129            if state != ffi::DuckDBSuccess {
130                return Err(Error::DuckDBFailure(
131                    ffi::Error::new(state),
132                    Some("failed to create config".to_string()),
133                ));
134            }
135            self.config = Some(config);
136        }
137        let c_key = CString::new(key)?;
138        let c_value = CString::new(value)?;
139        let state = unsafe {
140            ffi::duckdb_set_config(
141                self.config.unwrap(),
142                c_key.as_ptr() as *const c_char,
143                c_value.as_ptr() as *const c_char,
144            )
145        };
146        if state != ffi::DuckDBSuccess {
147            return Err(Error::DuckDBFailure(
148                ffi::Error::new(state),
149                Some(format!("set {key}:{value} error")),
150            ));
151        }
152        Ok(())
153    }
154}
155
156impl Drop for Config {
157    fn drop(&mut self) {
158        if let Some(ref mut config) = self.config {
159            unsafe { ffi::duckdb_destroy_config(config) };
160        }
161    }
162}
163
164#[cfg(test)]
165mod test {
166    use crate::{Config, Connection, Result, types::Value};
167
168    #[test]
169    fn test_default_config() -> Result<()> {
170        let config = Config::default();
171        let db = Connection::open_in_memory_with_flags(config)?;
172        db.execute_batch("CREATE TABLE foo(x Text)")?;
173
174        let mut stmt = db.prepare("INSERT INTO foo(x) VALUES (?)")?;
175        stmt.execute([&"a"])?;
176        stmt.execute([&"b"])?;
177        stmt.execute([&"c"])?;
178        stmt.execute([Value::Null])?;
179
180        let val: Result<Vec<Option<String>>> = db
181            .prepare("SELECT x FROM foo ORDER BY x")?
182            .query_and_then([], |row| row.get(0))?
183            .collect();
184        let val = val?;
185        let mut iter = val.iter();
186        assert_eq!(val.len(), 4);
187        assert_eq!(iter.next().unwrap().as_ref().unwrap(), "a");
188        assert_eq!(iter.next().unwrap().as_ref().unwrap(), "b");
189        assert_eq!(iter.next().unwrap().as_ref().unwrap(), "c");
190        assert!(iter.next().unwrap().is_none());
191        assert_eq!(iter.next(), None);
192
193        Ok(())
194    }
195
196    #[test]
197    fn test_all_config() -> Result<()> {
198        let config = Config::default()
199            .access_mode(crate::AccessMode::ReadWrite)?
200            .default_null_order(crate::DefaultNullOrder::NullsLast)?
201            .default_order(crate::DefaultOrder::Desc)?
202            .enable_external_access(true)?
203            .enable_object_cache(false)?
204            .enable_autoload_extension(true)?
205            .allow_unsigned_extensions()?
206            .custom_user_agent("test_user_agent")?
207            .max_memory("2GB")?
208            .threads(4)?
209            .with("preserve_insertion_order", "true")?;
210
211        let db = Connection::open_in_memory_with_flags(config)?;
212        db.execute_batch("CREATE TABLE foo(x Text)")?;
213
214        let mut stmt = db.prepare("INSERT INTO foo(x) VALUES (?)")?;
215        stmt.execute([&"a"])?;
216        stmt.execute([&"b"])?;
217        stmt.execute([&"c"])?;
218        stmt.execute([Value::Null])?;
219
220        let val: Result<Vec<Option<String>>> = db
221            .prepare("SELECT x FROM foo ORDER BY x")?
222            .query_and_then([], |row| row.get(0))?
223            .collect();
224        let val = val?;
225        let mut iter = val.iter();
226        assert_eq!(iter.next().unwrap().as_ref().unwrap(), "c");
227        assert_eq!(iter.next().unwrap().as_ref().unwrap(), "b");
228        assert_eq!(iter.next().unwrap().as_ref().unwrap(), "a");
229        assert!(iter.next().unwrap().is_none());
230        assert_eq!(iter.next(), None);
231
232        let user_agent: Result<String> = db.query_row("PRAGMA USER_AGENT", [], |row| row.get(0));
233        let user_agent = user_agent.unwrap();
234        assert!(&user_agent.ends_with("rust test_user_agent"));
235
236        Ok(())
237    }
238
239    #[test]
240    fn test_null_byte_in_key_returns_error() {
241        let err = Config::default().with("key\0name", "value");
242        assert!(err.is_err(), "expected Config::set to return Err for null byte in key");
243    }
244
245    #[test]
246    fn test_invalid_setting() -> Result<()> {
247        let config = Config::default().with("some-invalid-setting", "true")?;
248        let res = Connection::open_in_memory_with_flags(config);
249        assert_eq!(
250            res.unwrap_err().to_string(),
251            "Invalid Input Error: The following options were not recognized: some-invalid-setting"
252        );
253        Ok(())
254    }
255}
duckdb/config.rs

duckdb/
config.rs
