duckdb/
config.rs

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

duckdb/
config.rs
