better_duck_core/
pool.rs

1//! An [`r2d2`] connection pool backed by a shared [`Database`].
2//!
3//! Unlike opening one connection per pool slot, every connection this manager
4//! creates shares a single [`Database`] handle — so an in-memory pool observes
5//! one consistent database rather than one independent database per connection.
6
7use std::path::Path;
8
9use crate::{config::Config, connection::Connection, database::Database, error::Error};
10
11/// An [`r2d2::ManageConnection`] that creates connections from a shared [`Database`].
12#[derive(Clone, Debug)]
13pub struct DuckDbConnectionManager {
14    database: Database,
15}
16
17impl DuckDbConnectionManager {
18    /// Creates a manager over an already-open [`Database`].
19    pub fn new(database: Database) -> DuckDbConnectionManager {
20        DuckDbConnectionManager { database }
21    }
22
23    /// Creates a manager backed by a file-based database.
24    ///
25    /// # Errors
26    ///
27    /// Returns an error if the database cannot be opened.
28    pub fn file<P: AsRef<Path>>(path: P) -> crate::error::Result<DuckDbConnectionManager> {
29        Database::open(path).map(DuckDbConnectionManager::new)
30    }
31
32    /// Creates a manager backed by a file-based database with additional config.
33    ///
34    /// # Errors
35    ///
36    /// Returns an error if the database cannot be opened.
37    pub fn file_with_flags<P: AsRef<Path>>(
38        path: P,
39        config: Config,
40    ) -> crate::error::Result<DuckDbConnectionManager> {
41        Database::open_with_flags(path, config).map(DuckDbConnectionManager::new)
42    }
43
44    /// Creates a manager backed by a shared in-memory database.
45    ///
46    /// All connections checked out of a pool built from this manager observe the
47    /// same in-memory data.
48    ///
49    /// # Errors
50    ///
51    /// Returns an error if the database cannot be opened.
52    pub fn memory() -> crate::error::Result<DuckDbConnectionManager> {
53        Database::open_in_memory().map(DuckDbConnectionManager::new)
54    }
55
56    /// Creates a manager backed by a shared in-memory database with additional config.
57    ///
58    /// # Errors
59    ///
60    /// Returns an error if the database cannot be opened.
61    pub fn memory_with_flags(config: Config) -> crate::error::Result<DuckDbConnectionManager> {
62        Database::open_in_memory_with_flags(config).map(DuckDbConnectionManager::new)
63    }
64
65    /// Returns the shared [`Database`] backing this manager.
66    pub fn database(&self) -> &Database {
67        &self.database
68    }
69}
70
71impl r2d2::ManageConnection for DuckDbConnectionManager {
72    type Connection = Connection;
73    type Error = Error;
74
75    fn connect(&self) -> Result<Connection, Error> {
76        self.database.connect()
77    }
78
79    fn is_valid(
80        &self,
81        conn: &mut Connection,
82    ) -> Result<(), Error> {
83        conn.execute_batch("SELECT 1")
84    }
85
86    fn has_broken(
87        &self,
88        conn: &mut Connection,
89    ) -> bool {
90        !conn.is_open()
91    }
92}
93
94/// A DuckDB connection pool.
95pub type Pool = r2d2::Pool<DuckDbConnectionManager>;
96/// A connection checked out of a [`Pool`].
97pub type PooledConnection = r2d2::PooledConnection<DuckDbConnectionManager>;
98
99pub use r2d2::Builder as PoolBuilder;
100pub use r2d2::Error as PoolError;
101pub use r2d2::State as PoolState;
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106    use r2d2::ManageConnection;
107
108    #[test]
109    fn pool_shares_one_in_memory_database() {
110        let manager = DuckDbConnectionManager::memory().unwrap();
111        let pool = Pool::builder().max_size(4).build(manager).unwrap();
112
113        let mut a = pool.get().unwrap();
114        a.execute_batch("CREATE TABLE t (id INTEGER)").unwrap();
115        drop(a);
116
117        let mut b = pool.get().unwrap();
118        b.execute_batch("INSERT INTO t VALUES (1)").unwrap();
119        let mut result = b.execute("SELECT count(*) AS c FROM t").unwrap();
120        let row = result.next().unwrap().unwrap();
121        match row.get("c").unwrap() {
122            crate::types::value::DuckValue::BigInt(n) => assert_eq!(*n, 1),
123            other => panic!("expected BigInt, got {other:?}"),
124        }
125    }
126
127    #[test]
128    fn pool_max_size_is_enforced() {
129        let manager = DuckDbConnectionManager::memory().unwrap();
130        let pool = Pool::builder()
131            .max_size(1)
132            .connection_timeout(std::time::Duration::from_millis(100))
133            .build(manager)
134            .unwrap();
135
136        let _held = pool.get().unwrap();
137        assert!(pool.get().is_err());
138    }
139
140    #[test]
141    fn manager_is_send_sync() {
142        fn assert_send_sync<T: Send + Sync>() {}
143        assert_send_sync::<DuckDbConnectionManager>();
144    }
145
146    #[test]
147    fn manager_validates_open_connections() {
148        let manager = DuckDbConnectionManager::memory().unwrap();
149        let mut connection = manager.connect().unwrap();
150        assert!(manager.is_valid(&mut connection).is_ok());
151        assert!(!manager.has_broken(&mut connection));
152    }
153
154    #[test]
155    fn manager_file_connections_share_persisted_state() {
156        let directory = tempfile::tempdir().unwrap();
157        let path = directory.path().join("pool.duckdb");
158        let manager = DuckDbConnectionManager::file(&path).unwrap();
159        let mut first = manager.connect().unwrap();
160        first
161            .execute_batch("CREATE TABLE pooled (value INTEGER); INSERT INTO pooled VALUES (9)")
162            .unwrap();
163        drop(first);
164
165        let mut second = manager.connect().unwrap();
166        let mut rows = second.execute("SELECT value FROM pooled").unwrap();
167        let row = rows.next().unwrap().unwrap();
168        assert_eq!(row.get("value").unwrap(), &crate::types::value::DuckValue::Int(9));
169    }
170}
better_duck_core/pool.rs

better_duck_core/
pool.rs
