better_duck_core/asynchronous/
connection.rs

1use std::{path::Path, sync::Arc};
2
3use parking_lot::Mutex;
4
5use crate::{
6    connection::Connection,
7    error::{Error, Result},
8    result_set::ResultSet,
9    types::{appendable::AppendAble, value::DuckValue},
10    Appender,
11};
12
13/// An async facade over a [`Connection`].
14///
15/// Every method dispatches to `tokio::task::spawn_blocking`, so the connection
16/// never blocks the async executor. The handle is cheap to clone; clones share
17/// one connection, serialized by an internal mutex.
18///
19/// # Panics
20///
21/// These methods panic if called outside a Tokio runtime context, matching
22/// `tokio::task::spawn_blocking`.
23#[derive(Clone)]
24pub struct AsyncConnection {
25    inner: Arc<Mutex<Connection>>,
26}
27
28impl AsyncConnection {
29    /// Wraps an existing [`Connection`] for async use.
30    pub fn new(conn: Connection) -> AsyncConnection {
31        AsyncConnection { inner: Arc::new(Mutex::new(conn)) }
32    }
33
34    /// Recovers the inner [`Connection`] if this is the last handle.
35    ///
36    /// Returns `self` unchanged (as `Err`) if other clones of this handle exist.
37    pub fn try_into_inner(self) -> std::result::Result<Connection, AsyncConnection> {
38        Arc::try_unwrap(self.inner)
39            .map(Mutex::into_inner)
40            .map_err(|inner| AsyncConnection { inner })
41    }
42
43    /// Opens a connection to a DuckDB database at the given file path.
44    ///
45    /// # Errors
46    ///
47    /// Returns an error if the database cannot be opened.
48    pub async fn open<P>(path: P) -> Result<AsyncConnection>
49    where
50        P: AsRef<Path> + Send + 'static,
51    {
52        let conn = tokio::task::spawn_blocking(move || Connection::open(path))
53            .await
54            .map_err(|e| Error::BackgroundTaskFailed(e.to_string()))??;
55        Ok(AsyncConnection::new(conn))
56    }
57
58    /// Opens an in-memory DuckDB connection.
59    ///
60    /// # Errors
61    ///
62    /// Returns an error if the connection cannot be established.
63    pub async fn open_in_memory() -> Result<AsyncConnection> {
64        let conn = tokio::task::spawn_blocking(Connection::open_in_memory)
65            .await
66            .map_err(|e| Error::BackgroundTaskFailed(e.to_string()))??;
67        Ok(AsyncConnection::new(conn))
68    }
69
70    /// Runs an arbitrary closure against the connection on a blocking thread.
71    ///
72    /// This is the primitive every other method on this type is built from. Use
73    /// it directly for transactions and anything the typed helpers don't cover.
74    ///
75    /// # Errors
76    ///
77    /// Returns an error if the closure returns one, or if the background task
78    /// panics or is cancelled.
79    pub async fn with_connection<F, T>(
80        &self,
81        f: F,
82    ) -> Result<T>
83    where
84        F: FnOnce(&mut Connection) -> Result<T> + Send + 'static,
85        T: Send + 'static,
86    {
87        let inner = Arc::clone(&self.inner);
88        tokio::task::spawn_blocking(move || {
89            let mut guard = inner.lock();
90            f(&mut guard)
91        })
92        .await
93        .map_err(|e| Error::BackgroundTaskFailed(e.to_string()))?
94    }
95
96    /// Executes one or more SQL statements separated by semicolons.
97    ///
98    /// # Errors
99    ///
100    /// Returns an error if any statement fails to execute.
101    pub async fn execute_batch<S>(
102        &self,
103        sql: S,
104    ) -> Result<()>
105    where
106        S: Into<String> + Send,
107    {
108        let sql = sql.into();
109        self.with_connection(move |conn| conn.execute_batch(&sql)).await
110    }
111
112    /// Prepares and executes a SQL statement, materializing the result.
113    ///
114    /// # Errors
115    ///
116    /// Returns an error if DuckDB cannot prepare or execute the statement.
117    pub async fn execute<S>(
118        &self,
119        sql: S,
120    ) -> Result<ResultSet>
121    where
122        S: Into<String> + Send,
123    {
124        let sql = sql.into();
125        self.with_connection(move |conn| conn.execute(&sql)?.materialize()).await
126    }
127
128    /// Prepares and executes a parameterized SQL statement, materializing the result.
129    ///
130    /// # Errors
131    ///
132    /// Returns an error if preparation, binding, or execution fails.
133    pub async fn execute_with<S>(
134        &self,
135        sql: S,
136        binds: Vec<DuckValue>,
137    ) -> Result<ResultSet>
138    where
139        S: Into<String> + Send,
140    {
141        let sql = sql.into();
142        self.with_connection(move |conn| {
143            let mut owned = binds;
144            let mut refs: Vec<&mut dyn AppendAble> =
145                owned.iter_mut().map(|v| v as &mut dyn AppendAble).collect();
146            conn.execute_with(&sql, &mut refs)?.materialize()
147        })
148        .await
149    }
150
151    /// Runs a closure against a bulk-insert [`Appender`] for `table`/`schema` on a
152    /// blocking thread.
153    ///
154    /// The appender never leaves the closure — it holds raw FFI pointers, so it
155    /// cannot be exposed as a standalone async handle.
156    ///
157    /// # Errors
158    ///
159    /// Returns an error if the table does not exist, the appender cannot be
160    /// created, or the closure returns an error.
161    pub async fn with_appender<F, T>(
162        &self,
163        table: impl Into<String> + Send,
164        schema: impl Into<String> + Send,
165        f: F,
166    ) -> Result<T>
167    where
168        F: FnOnce(&mut Appender) -> Result<T> + Send + 'static,
169        T: Send + 'static,
170    {
171        let table = table.into();
172        let schema = schema.into();
173        self.with_connection(move |conn| {
174            let mut appender = conn.appender(&table, &schema)?;
175            f(&mut appender)
176        })
177        .await
178    }
179}
180
181#[cfg(test)]
182mod tests {
183    use super::*;
184
185    fn assert_send<T: Send>(_: T) {}
186    fn assert_send_sync<T: Send + Sync>() {}
187    fn assert_clone<T: Clone>() {}
188
189    #[test]
190    fn async_connection_is_send_sync_clone() {
191        assert_send_sync::<AsyncConnection>();
192        assert_clone::<AsyncConnection>();
193    }
194
195    #[tokio::test]
196    async fn execute_future_is_send() {
197        let conn = AsyncConnection::open_in_memory().await.unwrap();
198        assert_send(conn.execute("SELECT 1"));
199    }
200
201    #[tokio::test]
202    async fn execute_batch_then_execute() {
203        let conn = AsyncConnection::open_in_memory().await.unwrap();
204        conn.execute_batch("CREATE TABLE t (id INTEGER)").await.unwrap();
205        conn.execute_batch("INSERT INTO t VALUES (1)").await.unwrap();
206        let result = conn.execute("SELECT id FROM t").await.unwrap();
207        assert_eq!(result.len(), 1);
208    }
209
210    #[tokio::test]
211    async fn execute_with_owned_binds() {
212        let conn = AsyncConnection::open_in_memory().await.unwrap();
213        conn.execute_batch("CREATE TABLE t (id INTEGER)").await.unwrap();
214        conn.execute_batch("INSERT INTO t VALUES (1), (2), (3)").await.unwrap();
215        let result = conn
216            .execute_with("SELECT id FROM t WHERE id = $1", vec![DuckValue::Int(2)])
217            .await
218            .unwrap();
219        assert_eq!(result.len(), 1);
220    }
221
222    #[tokio::test]
223    async fn error_variant_survives_boundary() {
224        let conn = AsyncConnection::open_in_memory().await.unwrap();
225        let err = conn.execute_batch("NOT VALID SQL").await.unwrap_err();
226        assert!(matches!(err, Error::DuckDBFailure(..)));
227    }
228
229    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
230    async fn concurrent_queries_on_cloned_handles() {
231        let conn = AsyncConnection::open_in_memory().await.unwrap();
232        conn.execute_batch("CREATE TABLE t (id INTEGER)").await.unwrap();
233        let mut handles = Vec::new();
234        for i in 0..8 {
235            let c = conn.clone();
236            handles.push(tokio::spawn(async move {
237                c.execute_batch(format!("INSERT INTO t VALUES ({i})")).await.unwrap();
238            }));
239        }
240        for h in handles {
241            h.await.unwrap();
242        }
243        let result = conn.execute("SELECT count(*) AS c FROM t").await.unwrap();
244        match result.rows()[0].get("c").unwrap() {
245            DuckValue::BigInt(n) => assert_eq!(*n, 8),
246            other => panic!("expected BigInt, got {other:?}"),
247        }
248    }
249
250    #[tokio::test]
251    async fn with_appender_bulk_insert() {
252        let conn = AsyncConnection::open_in_memory().await.unwrap();
253        conn.execute_batch("CREATE TABLE t (id INTEGER)").await.unwrap();
254        conn.with_appender("t", "main", |appender| {
255            for i in 0..100i32 {
256                appender.append(&mut DuckValue::Int(i))?;
257            }
258            Ok(())
259        })
260        .await
261        .unwrap();
262        let result = conn.execute("SELECT count(*) AS c FROM t").await.unwrap();
263        match result.rows()[0].get("c").unwrap() {
264            DuckValue::BigInt(n) => assert_eq!(*n, 100),
265            other => panic!("expected BigInt, got {other:?}"),
266        }
267    }
268
269    #[tokio::test]
270    async fn with_connection_transaction_rollback() {
271        let conn = AsyncConnection::open_in_memory().await.unwrap();
272        conn.execute_batch("CREATE TABLE t (id INTEGER)").await.unwrap();
273        conn.with_connection(|c| {
274            c.execute_batch("BEGIN")?;
275            c.execute_batch("INSERT INTO t VALUES (1)")?;
276            c.execute_batch("ROLLBACK")?;
277            Ok(())
278        })
279        .await
280        .unwrap();
281        let result = conn.execute("SELECT count(*) AS c FROM t").await.unwrap();
282        match result.rows()[0].get("c").unwrap() {
283            DuckValue::BigInt(n) => assert_eq!(*n, 0),
284            other => panic!("expected BigInt, got {other:?}"),
285        }
286    }
287}
better_duck_core/asynchronous/connection.rs

better_duck_core/asynchronous/
connection.rs
