duckdb/
cache.rs

1//! Prepared statements cache for faster execution.
2
3use crate::{Connection, Result, Statement, raw_statement::RawStatement};
4use hashlink::LruCache;
5use std::{
6    cell::RefCell,
7    ops::{Deref, DerefMut},
8    sync::Arc,
9};
10
11impl Connection {
12    /// Prepare a SQL statement for execution, returning a previously prepared
13    /// (but not currently in-use) statement if one is available. The
14    /// returned statement will be cached for reuse by future calls to
15    /// [`prepare_cached`](Connection::prepare_cached) once it is dropped.
16    ///
17    /// ```rust,no_run
18    /// # use duckdb::{Connection, Result};
19    /// fn insert_new_people(conn: &Connection) -> Result<()> {
20    ///     {
21    ///         let mut stmt = conn.prepare_cached("INSERT INTO People (name) VALUES (?)")?;
22    ///         stmt.execute(["Joe Smith"])?;
23    ///     }
24    ///     {
25    ///         // This will return the same underlying DuckDB statement handle without
26    ///         // having to prepare it again.
27    ///         let mut stmt = conn.prepare_cached("INSERT INTO People (name) VALUES (?)")?;
28    ///         stmt.execute(["Bob Jones"])?;
29    ///     }
30    ///     Ok(())
31    /// }
32    /// ```
33    ///
34    /// # Failure
35    ///
36    /// Will return `Err` if `sql` cannot be converted to a C-compatible string
37    /// or if the underlying DuckDB call fails.
38    #[inline]
39    pub fn prepare_cached(&self, sql: &str) -> Result<CachedStatement<'_>> {
40        self.cache.get(self, sql)
41    }
42
43    /// Set the maximum number of cached prepared statements this connection
44    /// will hold. By default, a connection will hold a relatively small
45    /// number of cached statements. If you need more, or know that you
46    /// will not use cached statements, you
47    /// can set the capacity manually using this method.
48    #[inline]
49    pub fn set_prepared_statement_cache_capacity(&self, capacity: usize) {
50        self.cache.set_capacity(capacity);
51    }
52
53    /// Remove/finalize all prepared statements currently in the cache.
54    #[inline]
55    pub fn flush_prepared_statement_cache(&self) {
56        self.cache.flush();
57    }
58}
59
60/// Prepared statements LRU cache.
61#[derive(Debug)]
62pub struct StatementCache(RefCell<LruCache<Arc<str>, RawStatement>>);
63
64#[allow(clippy::non_send_fields_in_send_ty)]
65unsafe impl Send for StatementCache {}
66
67/// Cacheable statement.
68///
69/// Statement will return automatically to the cache by default.
70/// If you want the statement to be discarded, call
71/// [`discard()`](CachedStatement::discard) on it.
72pub struct CachedStatement<'conn> {
73    stmt: Option<Statement<'conn>>,
74    cache: &'conn StatementCache,
75}
76
77impl<'conn> Deref for CachedStatement<'conn> {
78    type Target = Statement<'conn>;
79
80    #[inline]
81    fn deref(&self) -> &Statement<'conn> {
82        self.stmt.as_ref().unwrap()
83    }
84}
85
86impl<'conn> DerefMut for CachedStatement<'conn> {
87    #[inline]
88    fn deref_mut(&mut self) -> &mut Statement<'conn> {
89        self.stmt.as_mut().unwrap()
90    }
91}
92
93impl Drop for CachedStatement<'_> {
94    #[inline]
95    fn drop(&mut self) {
96        if let Some(stmt) = self.stmt.take() {
97            self.cache.cache_stmt(unsafe { stmt.into_raw() });
98        }
99    }
100}
101
102impl CachedStatement<'_> {
103    #[inline]
104    fn new<'conn>(stmt: Statement<'conn>, cache: &'conn StatementCache) -> CachedStatement<'conn> {
105        CachedStatement {
106            stmt: Some(stmt),
107            cache,
108        }
109    }
110
111    /// Discard the statement, preventing it from being returned to its
112    /// [`Connection`]'s collection of cached statements.
113    #[inline]
114    pub fn discard(mut self) {
115        self.stmt = None;
116    }
117}
118
119impl StatementCache {
120    /// Create a statement cache.
121    #[inline]
122    pub fn with_capacity(capacity: usize) -> Self {
123        Self(RefCell::new(LruCache::new(capacity)))
124    }
125
126    #[inline]
127    fn set_capacity(&self, capacity: usize) {
128        self.0.borrow_mut().set_capacity(capacity);
129    }
130
131    // Search the cache for a prepared-statement object that implements `sql`.
132    // If no such prepared-statement can be found, allocate and prepare a new one.
133    //
134    // # Failure
135    //
136    // Will return `Err` if no cached statement can be found and the underlying
137    // DuckDB prepare call fails.
138    fn get<'conn>(&'conn self, conn: &'conn Connection, sql: &str) -> Result<CachedStatement<'conn>> {
139        let trimmed = sql.trim();
140        let mut cache = self.0.borrow_mut();
141        let stmt = match cache.remove(trimmed) {
142            Some(raw_stmt) => Ok(Statement::new(conn, raw_stmt)),
143            None => conn.prepare(trimmed),
144        };
145        stmt.map(|mut stmt| {
146            stmt.stmt.set_statement_cache_key(trimmed);
147            CachedStatement::new(stmt, self)
148        })
149    }
150
151    // Return a statement to the cache.
152    fn cache_stmt(&self, stmt: RawStatement) {
153        if stmt.is_null() {
154            return;
155        }
156        let mut cache = self.0.borrow_mut();
157        stmt.clear_bindings();
158        if let Some(sql) = stmt.statement_cache_key() {
159            cache.insert(sql, stmt);
160        } else {
161            debug_assert!(
162                false,
163                "bug in statement cache code, statement returned to cache that without key"
164            );
165        }
166    }
167
168    #[inline]
169    fn flush(&self) {
170        let mut cache = self.0.borrow_mut();
171        cache.clear();
172    }
173}
174
175#[cfg(test)]
176mod test {
177    use super::StatementCache;
178    use crate::{Connection, Result, core::LogicalTypeId, types::Value};
179    use arrow::array::Int32Array;
180    use fallible_iterator::FallibleIterator;
181
182    impl StatementCache {
183        fn clear(&self) {
184            self.0.borrow_mut().clear();
185        }
186
187        fn len(&self) -> usize {
188            self.0.borrow().len()
189        }
190
191        fn capacity(&self) -> usize {
192            self.0.borrow().capacity()
193        }
194    }
195
196    #[test]
197    fn test_cache() -> Result<()> {
198        let db = Connection::open_in_memory()?;
199        let cache = &db.cache;
200        let initial_capacity = cache.capacity();
201        assert_eq!(0, cache.len());
202        assert!(initial_capacity > 0);
203
204        let sql = "PRAGMA database_list";
205        {
206            let mut stmt = db.prepare_cached(sql)?;
207            assert_eq!(0, cache.len());
208            assert_eq!("memory", stmt.query_row([], |r| r.get::<_, String>(1))?);
209        }
210        assert_eq!(1, cache.len());
211
212        {
213            let mut stmt = db.prepare_cached(sql)?;
214            assert_eq!(0, cache.len());
215            assert_eq!("memory", stmt.query_row([], |r| r.get::<_, String>(1))?);
216        }
217        assert_eq!(1, cache.len());
218
219        cache.clear();
220        assert_eq!(0, cache.len());
221        assert_eq!(initial_capacity, cache.capacity());
222        Ok(())
223    }
224
225    #[test]
226    fn test_set_capacity() -> Result<()> {
227        let db = Connection::open_in_memory()?;
228        let cache = &db.cache;
229
230        let sql = "PRAGMA database_list";
231        {
232            let mut stmt = db.prepare_cached(sql)?;
233            assert_eq!(0, cache.len());
234            assert_eq!("memory", stmt.query_row([], |r| r.get::<_, String>(1))?);
235        }
236        assert_eq!(1, cache.len());
237
238        db.set_prepared_statement_cache_capacity(0);
239        assert_eq!(0, cache.len());
240
241        {
242            let mut stmt = db.prepare_cached(sql)?;
243            assert_eq!(0, cache.len());
244            assert_eq!("memory", stmt.query_row([], |r| r.get::<_, String>(1))?);
245        }
246        assert_eq!(0, cache.len());
247
248        db.set_prepared_statement_cache_capacity(8);
249        {
250            let mut stmt = db.prepare_cached(sql)?;
251            assert_eq!(0, cache.len());
252            assert_eq!("memory", stmt.query_row([], |r| r.get::<_, String>(1))?);
253        }
254        assert_eq!(1, cache.len());
255        Ok(())
256    }
257
258    #[test]
259    fn test_discard() -> Result<()> {
260        let db = Connection::open_in_memory()?;
261        let cache = &db.cache;
262
263        let sql = "PRAGMA database_list";
264        {
265            let mut stmt = db.prepare_cached(sql)?;
266            assert_eq!(0, cache.len());
267            assert_eq!("memory", stmt.query_row([], |r| r.get::<_, String>(1))?);
268            stmt.discard();
269        }
270        assert_eq!(0, cache.len());
271        Ok(())
272    }
273
274    #[test]
275    fn test_ddl() -> Result<()> {
276        let db = Connection::open_in_memory()?;
277        db.execute_batch(
278            r#"
279            CREATE TABLE foo (x INT);
280            INSERT INTO foo VALUES (1);
281        "#,
282        )?;
283
284        let sql = "SELECT * FROM foo";
285
286        {
287            let mut stmt = db.prepare_cached(sql)?;
288            assert_eq!(Ok(Some(1i32)), stmt.query([])?.map(|r| r.get(0)).next());
289        }
290
291        db.execute_batch(
292            r#"
293            ALTER TABLE foo ADD COLUMN y INT;
294            UPDATE foo SET y = 2;
295        "#,
296        )?;
297
298        {
299            // Rebinding statement after catalog change resulted in change of types
300            let mut stmt = db.prepare_cached(sql)?;
301            assert_eq!(
302                Ok(Some((1i32, 2i32))),
303                stmt.query([])?.map(|r| <(i32, i32)>::try_from(r)).next()
304            );
305        }
306        Ok(())
307    }
308
309    #[test]
310    fn test_ddl_with_trailing_geometry_uses_executed_metadata() -> Result<()> {
311        let db = Connection::open_in_memory()?;
312        db.execute_batch(
313            r#"
314            CREATE TABLE foo (x INT);
315            INSERT INTO foo VALUES (1);
316        "#,
317        )?;
318
319        let sql = "SELECT * FROM foo";
320
321        {
322            let mut stmt = db.prepare_cached(sql)?;
323            assert_eq!(Ok(Some(1i32)), stmt.query([])?.map(|r| r.get(0)).next());
324        }
325
326        db.execute_batch(
327            r#"
328            ALTER TABLE foo ADD COLUMN geom GEOMETRY;
329            UPDATE foo SET geom = 'POINT EMPTY'::GEOMETRY;
330        "#,
331        )?;
332
333        {
334            let mut stmt = db.prepare_cached(sql)?;
335            let got = stmt.query([])?.map(|r| r.get_ref(1).map(|v| v.to_owned())).next()?;
336            let point_empty_wkb = vec![
337                0x01, 0x01, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xF8, 0x7F, 0, 0, 0, 0, 0, 0, 0xF8, 0x7F,
338            ];
339            match got {
340                Some(Value::Geometry(wkb)) => assert_eq!(wkb, point_empty_wkb),
341                other => panic!("expected trailing geometry column, got {other:?}"),
342            }
343        }
344        Ok(())
345    }
346
347    #[test]
348    fn test_ddl_with_trailing_ambiguous_decimal_uses_executed_metadata() -> Result<()> {
349        let db = Connection::open_in_memory()?;
350        db.execute_batch(
351            r#"
352            CREATE TABLE foo (x INT);
353            INSERT INTO foo VALUES (1);
354        "#,
355        )?;
356
357        let sql = "SELECT * FROM foo";
358
359        {
360            let mut stmt = db.prepare_cached(sql)?;
361            assert_eq!(Ok(Some(1i32)), stmt.query([])?.map(|r| r.get(0)).next());
362        }
363
364        db.execute_batch(
365            r#"
366            ALTER TABLE foo ADD COLUMN y UHUGEINT;
367            UPDATE foo SET y = 340282366920938463463374607431768211455::UHUGEINT;
368        "#,
369        )?;
370
371        {
372            let mut stmt = db.prepare_cached(sql)?;
373            assert_eq!(
374                Ok(Some((Value::Int(1), Value::UHugeInt(u128::MAX)))),
375                stmt.query([])?
376                    .map(|r| Ok((r.get_ref(0)?.to_owned(), r.get_ref(1)?.to_owned())))
377                    .next()
378            );
379        }
380        Ok(())
381    }
382
383    #[test]
384    fn test_ddl_with_same_index_ambiguous_logical_type_uses_executed_metadata() -> Result<()> {
385        let db = Connection::open_in_memory()?;
386        db.execute_batch(
387            r#"
388            CREATE TABLE foo (x HUGEINT);
389            INSERT INTO foo VALUES (0);
390        "#,
391        )?;
392
393        let sql = "SELECT * FROM foo";
394
395        {
396            let mut stmt = db.prepare_cached(sql)?;
397            assert_eq!(
398                Ok(Some(Value::HugeInt(0))),
399                stmt.query([])?.map(|r| r.get_ref(0).map(|v| v.to_owned())).next()
400            );
401        }
402
403        db.execute_batch(
404            r#"
405            DROP TABLE foo;
406            CREATE TABLE foo (x UHUGEINT);
407            INSERT INTO foo VALUES (340282366920938463463374607431768211455::UHUGEINT);
408        "#,
409        )?;
410
411        {
412            let mut stmt = db.prepare_cached(sql)?;
413            assert_eq!(
414                Ok(Some(Value::UHugeInt(u128::MAX))),
415                stmt.query([])?.map(|r| r.get_ref(0).map(|v| v.to_owned())).next()
416            );
417        }
418        Ok(())
419    }
420
421    #[test]
422    fn test_ddl_with_same_index_ambiguous_logical_type_updates_column_logical_type() -> Result<()> {
423        let db = Connection::open_in_memory()?;
424        db.execute_batch(
425            r#"
426            CREATE TABLE foo (x HUGEINT);
427            INSERT INTO foo VALUES (0);
428        "#,
429        )?;
430
431        let sql = "SELECT * FROM foo";
432
433        {
434            let mut stmt = db.prepare_cached(sql)?;
435            let value = {
436                let mut rows = stmt.query([])?;
437                rows.next()?.unwrap().get_ref(0)?.to_owned()
438            };
439            assert_eq!(Value::HugeInt(0), value);
440            assert_eq!(LogicalTypeId::Hugeint, stmt.column_logical_type(0).id());
441        }
442
443        db.execute_batch(
444            r#"
445            DROP TABLE foo;
446            CREATE TABLE foo (x UHUGEINT);
447            INSERT INTO foo VALUES (340282366920938463463374607431768211455::UHUGEINT);
448        "#,
449        )?;
450
451        {
452            let mut stmt = db.prepare_cached(sql)?;
453            let value = {
454                let mut rows = stmt.query([])?;
455                rows.next()?.unwrap().get_ref(0)?.to_owned()
456            };
457            assert_eq!(Value::UHugeInt(u128::MAX), value);
458            assert_eq!(LogicalTypeId::UHugeint, stmt.column_logical_type(0).id());
459        }
460        Ok(())
461    }
462
463    #[test]
464    fn test_ddl_with_stale_variant_uses_executed_streaming_metadata() -> Result<()> {
465        let db = Connection::open_in_memory()?;
466        db.execute_batch("CREATE TABLE foo (x VARIANT);")?;
467
468        let sql = "SELECT * FROM foo";
469
470        {
471            let stmt = db.prepare_cached(sql)?;
472            assert_eq!(LogicalTypeId::Variant, stmt.column_logical_type(0).id());
473        }
474
475        db.execute_batch(
476            r#"
477            DROP TABLE foo;
478            CREATE TABLE foo (x INTEGER);
479            INSERT INTO foo VALUES (42);
480        "#,
481        )?;
482
483        {
484            let mut stmt = db.prepare_cached(sql)?;
485            let batches = stmt.stream_arrow([])?.collect::<Vec<_>>();
486            assert_eq!(1, batches.len());
487
488            let values = batches[0].column(0).as_any().downcast_ref::<Int32Array>().unwrap();
489            assert_eq!(42, values.value(0));
490            assert_eq!(LogicalTypeId::Integer, stmt.column_logical_type(0).id());
491        }
492        Ok(())
493    }
494
495    #[test]
496    fn test_prepare_cached_multichunk_reexecution_resets_result_chunks() -> Result<()> {
497        let db = Connection::open_in_memory()?;
498        let sql = "SELECT i FROM range(3000) AS t(i) ORDER BY i";
499
500        for _ in 0..2 {
501            let mut stmt = db.prepare_cached(sql)?;
502            let values = stmt
503                .query_map([], |row| row.get::<_, i64>(0))?
504                .collect::<Result<Vec<_>>>()?;
505
506            assert_eq!(values.len(), 3000);
507            assert_eq!(values[0], 0);
508            assert_eq!(values[2048], 2048);
509            assert_eq!(values[2999], 2999);
510        }
511
512        Ok(())
513    }
514
515    #[test]
516    fn test_connection_close() -> Result<()> {
517        let conn = Connection::open_in_memory()?;
518        conn.prepare_cached("SELECT * FROM sqlite_master;")?;
519
520        conn.close().expect("connection not closed");
521        Ok(())
522    }
523
524    #[test]
525    fn test_cache_key() -> Result<()> {
526        let db = Connection::open_in_memory()?;
527        let cache = &db.cache;
528        assert_eq!(0, cache.len());
529
530        let sql = "PRAGMA database_list; ";
531        {
532            let mut stmt = db.prepare_cached(sql)?;
533            assert_eq!(0, cache.len());
534            assert_eq!("memory", stmt.query_row([], |r| r.get::<_, String>(1))?);
535        }
536        assert_eq!(1, cache.len());
537
538        {
539            let mut stmt = db.prepare_cached(sql)?;
540            assert_eq!(0, cache.len());
541            assert_eq!("memory", stmt.query_row([], |r| r.get::<_, String>(1))?);
542        }
543        assert_eq!(1, cache.len());
544        Ok(())
545    }
546
547    #[test]
548    fn test_cannot_prepare_empty_stmt() -> Result<()> {
549        let conn = Connection::open_in_memory()?;
550        let result = conn.prepare_cached("");
551        assert!(result.is_err());
552        Ok(())
553    }
554}
duckdb/cache.rs

duckdb/
cache.rs
