Skip to main content

diesel/sqlite/connection/
hooks.rs

1use crate::sqlite::BorrowedSqliteConnection;
2
3use super::SqliteConnection;
4use core::num::NonZeroU32;
5
6pub(super) use super::{BusyDecision, CommitDecision, ProgressDecision};
7
8impl SqliteConnection {
9    /// Registers a callback invoked when a transaction is about to be
10    /// committed.
11    ///
12    /// The callback returns a [`CommitDecision`]: `Proceed` lets the commit
13    /// complete, `Rollback` converts it into a rollback.
14    ///
15    /// Only one commit hook can be active at a time per connection.
16    /// Registering a new one replaces the previous.
17    ///
18    /// The callback runs synchronously as part of the committing
19    /// `sqlite3_step()` call, on the thread performing the commit, so it is
20    /// never invoked concurrently. Per SQLite, the callback must not use the
21    /// connection that triggered it (running any SQL, including a `SELECT`,
22    /// counts as use) and is not reentrant. A panic in the callback aborts the
23    /// process.
24    ///
25    /// See: [`sqlite3_commit_hook`](https://www.sqlite.org/c3ref/commit_hook.html)
26    ///
27    /// # Example
28    ///
29    /// ```rust
30    /// # // This test registers Rust callbacks that sqlite invokes through the
31    /// # // FFI boundary, which is not supported by miri `-Zmiri-native-lib`.
32    /// # #[cfg(not(miri))]
33    /// # fn main() {
34    /// use diesel::prelude::*;
35    /// use diesel::sqlite::{SqliteConnection, CommitDecision};
36    /// use std::sync::{Arc, Mutex};
37    ///
38    /// diesel::table! {
39    ///     users (id) {
40    ///         id -> Integer,
41    ///         name -> Text,
42    ///     }
43    /// }
44    ///
45    /// # use diesel::connection::SimpleConnection;
46    /// # let conn = &mut SqliteConnection::establish(":memory:").unwrap();
47    /// # conn.batch_execute("CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT NOT NULL)").unwrap();
48    /// let commits = Arc::new(Mutex::new(0u32));
49    /// let commits2 = commits.clone();
50    ///
51    /// conn.on_commit(move || {
52    ///     *commits2.lock().unwrap() += 1;
53    ///     CommitDecision::Proceed
54    /// });
55    ///
56    /// conn.immediate_transaction(|conn| {
57    ///     diesel::insert_into(users::table)
58    ///         .values(users::name.eq("Alice"))
59    ///         .execute(conn)?;
60    ///     Ok::<_, diesel::result::Error>(())
61    /// }).unwrap();
62    ///
63    /// assert_eq!(*commits.lock().unwrap(), 1);
64    /// # }
65    /// # #[cfg(miri)]
66    /// # fn main() {}
67    /// ```
68    pub fn on_commit<F>(&mut self, hook: F)
69    where
70        F: FnMut() -> CommitDecision + Send + 'static,
71    {
72        self.raw_connection.set_commit_hook(hook);
73    }
74
75    /// Removes the commit hook. Subsequent commits will not invoke any
76    /// callback.
77    ///
78    /// See [`on_commit`](Self::on_commit) for usage example.
79    pub fn remove_commit_hook(&mut self) {
80        self.raw_connection.remove_commit_hook();
81    }
82
83    /// Registers a callback invoked after a transaction is rolled back.
84    ///
85    /// This is **not** invoked for the implicit rollback that occurs when
86    /// the connection is closed. It **is** invoked when a commit hook forces
87    /// a rollback by returning [`CommitDecision::Rollback`].
88    ///
89    /// Only one rollback hook can be active at a time per connection.
90    /// Registering a new one replaces the previous.
91    ///
92    /// The callback must not use the database connection. It is invoked
93    /// synchronously on the thread driving the connection, so it is never
94    /// called concurrently, and like the commit hook it is not reentrant.
95    /// Panics in the callback abort the process.
96    ///
97    /// See: [`sqlite3_rollback_hook`](https://www.sqlite.org/c3ref/commit_hook.html)
98    ///
99    /// # Example
100    ///
101    /// ```rust
102    /// # // This test registers Rust callbacks that sqlite invokes through the
103    /// # // FFI boundary, which is not supported by miri `-Zmiri-native-lib`.
104    /// # #[cfg(not(miri))]
105    /// # fn main() {
106    /// use diesel::prelude::*;
107    /// use diesel::sqlite::SqliteConnection;
108    /// use std::sync::Arc;
109    /// use std::sync::atomic::{AtomicU32, Ordering};
110    ///
111    /// # use diesel::connection::SimpleConnection;
112    /// # let conn = &mut SqliteConnection::establish(":memory:").unwrap();
113    /// # conn.batch_execute("CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT NOT NULL)").unwrap();
114    /// let rollbacks = Arc::new(AtomicU32::new(0));
115    /// let rb2 = rollbacks.clone();
116    ///
117    /// conn.on_rollback(move || {
118    ///     rb2.fetch_add(1, Ordering::Relaxed);
119    /// });
120    ///
121    /// // Force a rollback by returning an error.
122    /// let _ = conn.immediate_transaction(|_conn| {
123    ///     Err::<(), _>(diesel::result::Error::RollbackTransaction)
124    /// });
125    ///
126    /// assert_eq!(rollbacks.load(Ordering::Relaxed), 1);
127    /// # }
128    /// # #[cfg(miri)]
129    /// # fn main() {}
130    /// ```
131    pub fn on_rollback<F>(&mut self, hook: F)
132    where
133        F: FnMut() + Send + 'static,
134    {
135        self.raw_connection.set_rollback_hook(hook);
136    }
137
138    /// Removes the rollback hook. Subsequent rollbacks will not invoke any
139    /// callback.
140    ///
141    /// See [`on_rollback`](Self::on_rollback) for usage example.
142    pub fn remove_rollback_hook(&mut self) {
143        self.raw_connection.remove_rollback_hook();
144    }
145
146    /// Registers a progress handler that can interrupt long-running queries.
147    ///
148    /// The callback is invoked periodically while a query runs. `n` is the
149    /// approximate number of virtual-machine instructions between callbacks. It
150    /// is a [`NonZeroU32`] so the handler cannot be disabled implicitly by
151    /// passing zero. Use [`remove_progress_handler`](Self::remove_progress_handler)
152    /// to disable it. Since SQLite 3.41.0 the callback may also fire during
153    /// statement preparation.
154    ///
155    /// The callback returns a [`ProgressDecision`]: `Continue` lets the query
156    /// keep executing, `Interrupt` aborts it (causes `SQLITE_INTERRUPT`).
157    ///
158    /// Only one progress handler can be active at a time per connection.
159    /// Registering a new one replaces the previous.
160    ///
161    /// The callback must not use the database connection. It is invoked
162    /// synchronously on the thread driving the connection, so it is never
163    /// called concurrently. Panics in the callback abort the process.
164    ///
165    /// See: [`sqlite3_progress_handler`](https://www.sqlite.org/c3ref/progress_handler.html)
166    ///
167    /// # Example
168    ///
169    /// ```rust
170    /// use diesel::prelude::*;
171    /// use diesel::sqlite::{SqliteConnection, ProgressDecision};
172    /// use std::num::NonZeroU32;
173    /// use std::sync::Arc;
174    /// use std::sync::atomic::{AtomicBool, Ordering};
175    ///
176    /// # let conn = &mut SqliteConnection::establish(":memory:").unwrap();
177    /// let cancelled = Arc::new(AtomicBool::new(false));
178    /// let cancelled2 = cancelled.clone();
179    ///
180    /// conn.on_progress(NonZeroU32::new(1000).unwrap(), move || {
181    ///     if cancelled2.load(Ordering::Relaxed) {
182    ///         ProgressDecision::Interrupt
183    ///     } else {
184    ///         ProgressDecision::Continue
185    ///     }
186    /// });
187    ///
188    /// // Later: remove the handler.
189    /// conn.remove_progress_handler();
190    /// ```
191    pub fn on_progress<F>(&mut self, n: NonZeroU32, hook: F)
192    where
193        F: FnMut() -> ProgressDecision + Send + 'static,
194    {
195        self.raw_connection.set_progress_handler(n, hook);
196    }
197
198    /// Removes the progress handler. Subsequent queries will not invoke any
199    /// callback.
200    ///
201    /// See [`on_progress`](Self::on_progress) for usage example.
202    pub fn remove_progress_handler(&mut self) {
203        self.raw_connection.remove_progress_handler();
204    }
205
206    /// Registers a callback invoked after each commit of a transaction in
207    /// [WAL mode](https://www.sqlite.org/wal.html). It receives a borrowed
208    /// connection, the database name (`"main"`, `"temp"`, or an `ATTACH`
209    /// alias), and the current WAL page count.
210    ///
211    /// The hook fires after the commit completes and the write-lock is
212    /// released, so the callback may read, write, or
213    /// [checkpoint](https://www.sqlite.org/wal.html#ckpt) through `conn`,
214    /// provided it leaves no open transaction. A write that commits inside the
215    /// callback re-fires the hook re-entrantly, so a callback that writes
216    /// unconditionally will recurse until the stack overflows. Guard against
217    /// that yourself if the callback writes.
218    ///
219    /// Only one WAL hook is active at a time, and re-registering replaces it.
220    /// `PRAGMA wal_autocheckpoint` installs its own WAL hook and overwrites
221    /// this one. A panic in the callback aborts the process.
222    ///
223    /// See: [`sqlite3_wal_hook`](https://www.sqlite.org/c3ref/wal_hook.html)
224    ///
225    /// # Example
226    ///
227    /// ```rust,no_run
228    /// # use diesel::prelude::*;
229    /// # use diesel::connection::SimpleConnection;
230    /// # use diesel::sqlite::SqliteConnection;
231    /// # use diesel::sqlite::WalCheckpointMode;
232    ///
233    /// # let conn = &mut SqliteConnection::establish("test.db").unwrap();
234    /// # conn.batch_execute("PRAGMA journal_mode = WAL;").unwrap();
235    /// conn.on_wal(|conn, db_name, n_pages| {
236    ///     println!("WAL for {db_name}: {n_pages} pages");
237    ///     if n_pages > 1000 {
238    ///         // The connection may be used here, e.g. to force a checkpoint.
239    ///         unsafe {
240    ///             // SAFETY: This is safe as long as we don't move or replace out the
241    ///             // inner SqliteConnection. None of the functions defined on SqliteConnection
242    ///             // do that
243    ///             let _ = conn.inner_mut().wal_checkpoint(None, WalCheckpointMode::Truncate);
244    ///         }
245    ///     }
246    /// });
247    /// ```
248    pub fn on_wal<F>(&mut self, hook: F)
249    where
250        F: Fn(&'_ mut BorrowedSqliteConnection, &str, u32) + Send + 'static,
251    {
252        self.raw_connection.set_wal_hook(hook);
253    }
254
255    /// Removes the WAL hook. Subsequent commits will not invoke any callback.
256    ///
257    /// See [`on_wal`](Self::on_wal) for usage example.
258    pub fn remove_wal_hook(&mut self) {
259        self.raw_connection.remove_wal_hook();
260    }
261
262    /// Registers a custom busy handler for lock contention.
263    ///
264    /// The callback receives the retry count (starting from 0) and returns a
265    /// [`BusyDecision`]: `Retry` retries the locked operation, `GiveUp` aborts
266    /// and returns `SQLITE_BUSY` to the caller.
267    ///
268    /// Setting this clears any timeout previously set with
269    /// [`set_busy_timeout`](Self::set_busy_timeout). Conversely, calling
270    /// `set_busy_timeout` clears this handler. Only one busy handler can be
271    /// active at a time per connection.
272    ///
273    /// The callback must not use the database connection. If the callback
274    /// modifies the database, behavior is undefined. SQLite may return
275    /// `SQLITE_BUSY` instead of calling the handler to prevent deadlocks.
276    ///
277    /// The callback is invoked synchronously on the thread driving the
278    /// connection, so it is never called concurrently, and per SQLite it is
279    /// not reentrant.
280    ///
281    /// Panics in the callback abort the process.
282    ///
283    /// See: [`sqlite3_busy_handler`](https://www.sqlite.org/c3ref/busy_handler.html)
284    ///
285    /// # Example
286    ///
287    /// ```rust
288    /// use diesel::prelude::*;
289    /// use diesel::sqlite::{SqliteConnection, BusyDecision};
290    /// use std::thread;
291    /// use std::time::Duration;
292    ///
293    /// # let conn = &mut SqliteConnection::establish(":memory:").unwrap();
294    /// conn.on_busy(|retry_count| {
295    ///     if retry_count < 5 {
296    ///         thread::sleep(Duration::from_millis(100));
297    ///         BusyDecision::Retry
298    ///     } else {
299    ///         BusyDecision::GiveUp
300    ///     }
301    /// });
302    ///
303    /// // Later: remove the handler
304    /// conn.remove_busy_handler();
305    /// ```
306    pub fn on_busy<F>(&mut self, hook: F)
307    where
308        F: FnMut(i32) -> BusyDecision + Send + 'static,
309    {
310        self.raw_connection.set_busy_handler(hook);
311    }
312
313    /// Removes the custom busy handler.
314    ///
315    /// See [`on_busy`](Self::on_busy) for usage example.
316    pub fn remove_busy_handler(&mut self) {
317        self.raw_connection.remove_busy_handler();
318    }
319
320    /// Sets a simple timeout-based busy handler.
321    ///
322    /// When a table is locked, SQLite will sleep and retry until `timeout`
323    /// has elapsed, rounded up to whole milliseconds and capped at `i32::MAX`
324    /// milliseconds. Pass [`Duration::ZERO`](core::time::Duration::ZERO) to
325    /// disable (return `SQLITE_BUSY` immediately).
326    ///
327    /// Setting this clears any custom [`on_busy`](Self::on_busy) handler.
328    /// Conversely, calling `on_busy` clears this timeout. For most use cases,
329    /// this is simpler than a custom busy handler.
330    ///
331    /// See: [`sqlite3_busy_timeout`](https://www.sqlite.org/c3ref/busy_timeout.html)
332    ///
333    /// # Example
334    ///
335    /// ```rust
336    /// use diesel::prelude::*;
337    /// use diesel::sqlite::SqliteConnection;
338    ///
339    /// # let conn = &mut SqliteConnection::establish(":memory:").unwrap();
340    /// // Wait up to 5 seconds for locked tables
341    /// conn.set_busy_timeout(std::time::Duration::from_secs(5));
342    /// ```
343    pub fn set_busy_timeout(&mut self, timeout: core::time::Duration) {
344        self.raw_connection.set_busy_timeout(timeout);
345    }
346}
347
348#[cfg(test)]
349mod tests {
350    use super::*;
351    use crate::connection::Connection;
352    use crate::query_dsl::RunQueryDsl;
353    use std::sync::Arc;
354    use std::sync::atomic::{AtomicU32, Ordering};
355
356    fn connection() -> SqliteConnection {
357        SqliteConnection::establish(":memory:").unwrap()
358    }
359
360    // Only used by tests that are compiled out under miri (see the
361    // individual tests below).
362    #[cfg(not(miri))]
363    #[derive(crate::QueryableByName)]
364    struct CountResult {
365        #[diesel(sql_type = crate::sql_types::BigInt)]
366        c: i64,
367    }
368
369    // Registers a callback that is invoked by the native library, which is not
370    // supported when running under miri with a native libsqlite3 (`-Zmiri-native-lib`).
371    #[cfg(not(miri))] // ffi call
372    #[diesel_test_helper::test]
373    fn on_commit_fires_on_commit() {
374        let conn = &mut connection();
375
376        let count = Arc::new(AtomicU32::new(0));
377        let c2 = count.clone();
378
379        conn.on_commit(move || {
380            c2.fetch_add(1, Ordering::Relaxed);
381            CommitDecision::Proceed
382        });
383
384        conn.immediate_transaction(|conn| {
385            crate::sql_query("CREATE TABLE t1 (id INTEGER PRIMARY KEY)")
386                .execute(conn)
387                .unwrap();
388            Ok::<_, crate::result::Error>(())
389        })
390        .unwrap();
391
392        assert_eq!(count.load(Ordering::Relaxed), 1);
393    }
394
395    // Registers a callback that is invoked by the native library, which is not
396    // supported when running under miri with a native libsqlite3 (`-Zmiri-native-lib`).
397    #[cfg(not(miri))] // ffi call
398    #[diesel_test_helper::test]
399    fn on_commit_returning_true_forces_rollback() {
400        let conn = &mut connection();
401
402        crate::sql_query("CREATE TABLE t_commit (id INTEGER PRIMARY KEY)")
403            .execute(conn)
404            .unwrap();
405
406        conn.on_commit(|| CommitDecision::Rollback);
407
408        // The transaction will attempt to commit, but the hook will convert
409        // it to a rollback. diesel's AnsiTransactionManager will see the
410        // failure from the COMMIT statement (sqlite returns an error when
411        // the commit hook returns non-zero and the commit is aborted).
412        let result = conn.immediate_transaction(|conn| {
413            crate::sql_query("INSERT INTO t_commit (id) VALUES (1)")
414                .execute(conn)
415                .unwrap();
416            Ok::<_, crate::result::Error>(())
417        });
418
419        // The transaction should have been rolled back.
420        assert!(
421            matches!(result, Err(crate::result::Error::DatabaseError(..))),
422            "{result:?}"
423        );
424        assert_eq!(
425            Ok(None),
426            <crate::connection::AnsiTransactionManager as crate::connection::TransactionManager<
427                SqliteConnection,
428            >>::transaction_manager_status_mut(conn)
429            .transaction_depth()
430        );
431
432        // Remove the hook so subsequent queries don't fail.
433        conn.remove_commit_hook();
434        conn.immediate_transaction(|_| Ok::<_, crate::result::Error>(()))
435            .unwrap();
436
437        // Verify the row was not persisted.
438        let cnt: i64 = crate::sql_query("SELECT COUNT(*) as c FROM t_commit")
439            .get_result::<CountResult>(conn)
440            .unwrap()
441            .c;
442        assert_eq!(cnt, 0);
443    }
444
445    // Registers a callback that is invoked by the native library, which is not
446    // supported when running under miri with a native libsqlite3 (`-Zmiri-native-lib`).
447    #[cfg(not(miri))] // ffi call
448    #[diesel_test_helper::test]
449    fn replacing_commit_hook_drops_old() {
450        let conn = &mut connection();
451
452        let old_count = Arc::new(AtomicU32::new(0));
453        let new_count = Arc::new(AtomicU32::new(0));
454        let oc = old_count.clone();
455        let nc = new_count.clone();
456
457        conn.on_commit(move || {
458            oc.fetch_add(1, Ordering::Relaxed);
459            CommitDecision::Proceed
460        });
461
462        // Replace with a new hook.
463        conn.on_commit(move || {
464            nc.fetch_add(1, Ordering::Relaxed);
465            CommitDecision::Proceed
466        });
467
468        conn.immediate_transaction(|conn| {
469            crate::sql_query("CREATE TABLE t_replace (id INTEGER PRIMARY KEY)")
470                .execute(conn)
471                .unwrap();
472            Ok::<_, crate::result::Error>(())
473        })
474        .unwrap();
475
476        assert_eq!(old_count.load(Ordering::Relaxed), 0);
477        assert_eq!(new_count.load(Ordering::Relaxed), 1);
478    }
479
480    #[diesel_test_helper::test]
481    fn remove_commit_hook_disables_callback() {
482        let conn = &mut connection();
483
484        let count = Arc::new(AtomicU32::new(0));
485        let c2 = count.clone();
486
487        conn.on_commit(move || {
488            c2.fetch_add(1, Ordering::Relaxed);
489            CommitDecision::Proceed
490        });
491
492        conn.remove_commit_hook();
493
494        conn.immediate_transaction(|conn| {
495            crate::sql_query("CREATE TABLE t_rem (id INTEGER PRIMARY KEY)")
496                .execute(conn)
497                .unwrap();
498            Ok::<_, crate::result::Error>(())
499        })
500        .unwrap();
501
502        assert_eq!(count.load(Ordering::Relaxed), 0);
503    }
504
505    // Registers a callback that is invoked by the native library, which is not
506    // supported when running under miri with a native libsqlite3 (`-Zmiri-native-lib`).
507    #[cfg(not(miri))] // ffi call
508    #[diesel_test_helper::test]
509    fn on_rollback_fires_on_explicit_rollback() {
510        let conn = &mut connection();
511
512        crate::sql_query("CREATE TABLE t_rb (id INTEGER PRIMARY KEY)")
513            .execute(conn)
514            .unwrap();
515
516        let count = Arc::new(AtomicU32::new(0));
517        let c2 = count.clone();
518
519        conn.on_rollback(move || {
520            c2.fetch_add(1, Ordering::Relaxed);
521        });
522
523        // Force a rollback by returning Err from the transaction closure.
524        let _ = conn.immediate_transaction(|conn| {
525            crate::sql_query("INSERT INTO t_rb (id) VALUES (1)")
526                .execute(conn)
527                .unwrap();
528            Err::<(), _>(crate::result::Error::RollbackTransaction)
529        });
530
531        assert_eq!(count.load(Ordering::Relaxed), 1);
532    }
533
534    // Registers a callback that is invoked by the native library, which is not
535    // supported when running under miri with a native libsqlite3 (`-Zmiri-native-lib`).
536    #[cfg(not(miri))] // ffi call
537    #[diesel_test_helper::test]
538    fn on_rollback_fires_when_commit_hook_forces_rollback() {
539        let conn = &mut connection();
540
541        crate::sql_query("CREATE TABLE t_rb2 (id INTEGER PRIMARY KEY)")
542            .execute(conn)
543            .unwrap();
544
545        let rb_count = Arc::new(AtomicU32::new(0));
546        let rb2 = rb_count.clone();
547
548        conn.on_commit(|| CommitDecision::Rollback);
549        conn.on_rollback(move || {
550            rb2.fetch_add(1, Ordering::Relaxed);
551        });
552
553        let _ = conn.immediate_transaction(|conn| {
554            crate::sql_query("INSERT INTO t_rb2 (id) VALUES (1)")
555                .execute(conn)
556                .unwrap();
557            Ok::<_, crate::result::Error>(())
558        });
559
560        // Rollback hook should have fired.
561        assert_eq!(rb_count.load(Ordering::Relaxed), 1);
562
563        conn.remove_commit_hook();
564        conn.remove_rollback_hook();
565
566        // Verify the row was not persisted.
567        let cnt: i64 = crate::sql_query("SELECT COUNT(*) as c FROM t_rb2")
568            .get_result::<CountResult>(conn)
569            .unwrap()
570            .c;
571        assert_eq!(cnt, 0);
572    }
573
574    #[diesel_test_helper::test]
575    fn on_rollback_does_not_fire_on_connection_close() {
576        let count = Arc::new(AtomicU32::new(0));
577        let c2 = count.clone();
578
579        {
580            let conn = &mut connection();
581            conn.on_rollback(move || {
582                c2.fetch_add(1, Ordering::Relaxed);
583            });
584            // conn is dropped here: implicit close, not a rollback.
585        }
586
587        assert_eq!(count.load(Ordering::Relaxed), 0);
588    }
589
590    #[diesel_test_helper::test]
591    fn remove_rollback_hook_disables_callback() {
592        let conn = &mut connection();
593
594        crate::sql_query("CREATE TABLE t_rem_rb (id INTEGER PRIMARY KEY)")
595            .execute(conn)
596            .unwrap();
597
598        let count = Arc::new(AtomicU32::new(0));
599        let c2 = count.clone();
600
601        conn.on_rollback(move || {
602            c2.fetch_add(1, Ordering::Relaxed);
603        });
604
605        conn.remove_rollback_hook();
606
607        let _ = conn.immediate_transaction(|conn| {
608            crate::sql_query("INSERT INTO t_rem_rb (id) VALUES (1)")
609                .execute(conn)
610                .unwrap();
611            Err::<(), _>(crate::result::Error::RollbackTransaction)
612        });
613
614        assert_eq!(count.load(Ordering::Relaxed), 0);
615    }
616
617    // A recursive CTE heavy enough that the progress handler fires while it runs.
618    const HEAVY_QUERY: &str = "WITH RECURSIVE c(x) AS \
619        (SELECT 1 UNION ALL SELECT x + 1 FROM c WHERE x < 100000) SELECT count(*) FROM c";
620
621    // Registers a callback that is invoked by the native library, which is not
622    // supported when running under miri with a native libsqlite3 (`-Zmiri-native-lib`).
623    #[cfg(not(miri))] // ffi call
624    #[diesel_test_helper::test]
625    fn on_progress_interrupts_query() {
626        let conn = &mut connection();
627
628        conn.on_progress(NonZeroU32::new(1).unwrap(), || ProgressDecision::Interrupt);
629
630        let result = crate::sql_query(HEAVY_QUERY).execute(conn);
631        assert!(
632            result.is_err(),
633            "the query should be interrupted by the progress handler"
634        );
635    }
636
637    #[diesel_test_helper::test]
638    fn remove_progress_handler_stops_interruption() {
639        let conn = &mut connection();
640
641        conn.on_progress(NonZeroU32::new(1).unwrap(), || ProgressDecision::Interrupt);
642        conn.remove_progress_handler();
643
644        // With the handler removed the same query runs to completion.
645        let result = crate::sql_query(HEAVY_QUERY).execute(conn);
646        assert!(
647            result.is_ok(),
648            "the query should complete after the handler is removed"
649        );
650    }
651
652    // WAL hook tests
653    //
654    // Gated out on WASM because these tests need a file-backed database
655    // (WAL mode does not work with `:memory:`), and `tempfile::tempdir()`
656    // panics on WASM due to the lack of a filesystem.
657    //
658    // The WAL API itself (`on_wal`, `remove_wal_hook`) is available on all
659    // platforms, including WASM. The file-backed databases used here cannot
660    // be created on WASM, so only the tests are gated out.
661
662    #[cfg(not(any(all(target_family = "wasm", target_os = "unknown"), miri)))]
663    /// Helper: create a file-backed connection in WAL mode (WAL requires a real
664    /// file, and every WAL test below wants the connection already in WAL mode).
665    /// Its only consumers are also compiled out under miri (see the
666    /// individual tests below).
667    fn wal_connection() -> (SqliteConnection, tempfile::TempDir) {
668        let dir = tempfile::tempdir().unwrap();
669        let path = dir.path().join("test.db");
670        let mut conn = SqliteConnection::establish(path.to_str().unwrap()).unwrap();
671        crate::sql_query("PRAGMA journal_mode=WAL")
672            .execute(&mut conn)
673            .unwrap();
674        (conn, dir)
675    }
676
677    #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
678    // Registers a callback that is invoked by the native library, which is not
679    // supported when running under miri with a native libsqlite3 (`-Zmiri-native-lib`).
680    #[cfg(not(miri))] // ffi call
681    #[diesel_test_helper::test]
682    fn on_wal_fires_in_wal_mode() {
683        let (conn, _dir) = &mut wal_connection();
684
685        crate::sql_query("CREATE TABLE t_wal (id INTEGER PRIMARY KEY)")
686            .execute(conn)
687            .unwrap();
688
689        let events: Arc<std::sync::Mutex<Vec<(String, u32)>>> =
690            Arc::new(std::sync::Mutex::new(Vec::new()));
691        let events2 = events.clone();
692
693        conn.on_wal(move |_, db_name, n_pages| {
694            events2.lock().unwrap().push((db_name.to_owned(), n_pages));
695        });
696
697        crate::sql_query("INSERT INTO t_wal (id) VALUES (1)")
698            .execute(conn)
699            .unwrap();
700
701        let events = events.lock().unwrap();
702        assert!(
703            !events.is_empty(),
704            "WAL hook should have fired at least once"
705        );
706        assert!(
707            events.iter().all(|(db_name, _)| db_name == "main"),
708            "db_name should always be \"main\""
709        );
710        assert!(
711            events.iter().any(|(_, n_pages)| *n_pages > 0),
712            "n_pages should be positive"
713        );
714    }
715
716    #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
717    // Registers a callback that is invoked by the native library, which is not
718    // supported when running under miri with a native libsqlite3 (`-Zmiri-native-lib`).
719    #[cfg(not(miri))] // ffi call
720    #[diesel_test_helper::test]
721    fn replacing_wal_hook_drops_old() {
722        let (conn, _dir) = &mut wal_connection();
723
724        crate::sql_query("CREATE TABLE t_wal2 (id INTEGER PRIMARY KEY)")
725            .execute(conn)
726            .unwrap();
727
728        let old_count = Arc::new(AtomicU32::new(0));
729        let new_count = Arc::new(AtomicU32::new(0));
730
731        let c_old = old_count.clone();
732        conn.on_wal(move |_, _, _| {
733            c_old.fetch_add(1, Ordering::Relaxed);
734        });
735
736        // Replace with a new hook.
737        let c_new = new_count.clone();
738        conn.on_wal(move |_, _, _| {
739            c_new.fetch_add(1, Ordering::Relaxed);
740        });
741
742        crate::sql_query("INSERT INTO t_wal2 (id) VALUES (1)")
743            .execute(conn)
744            .unwrap();
745
746        // Old hook should NOT have fired after replacement.
747        let old_before = old_count.load(Ordering::Relaxed);
748        crate::sql_query("INSERT INTO t_wal2 (id) VALUES (2)")
749            .execute(conn)
750            .unwrap();
751        assert_eq!(
752            old_count.load(Ordering::Relaxed),
753            old_before,
754            "old WAL hook should not fire after replacement"
755        );
756        assert!(
757            new_count.load(Ordering::Relaxed) > 0,
758            "new WAL hook should fire"
759        );
760    }
761
762    #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
763    // Registers a callback that is invoked by the native library, which is not
764    // supported when running under miri with a native libsqlite3 (`-Zmiri-native-lib`).
765    #[cfg(not(miri))] // ffi call
766    #[diesel_test_helper::test]
767    fn remove_wal_hook_disables_callback() {
768        let (conn, _dir) = &mut wal_connection();
769
770        crate::sql_query("CREATE TABLE t_wal3 (id INTEGER PRIMARY KEY)")
771            .execute(conn)
772            .unwrap();
773
774        let count = Arc::new(AtomicU32::new(0));
775        let c2 = count.clone();
776
777        conn.on_wal(move |_, _, _| {
778            c2.fetch_add(1, Ordering::Relaxed);
779        });
780
781        conn.remove_wal_hook();
782
783        crate::sql_query("INSERT INTO t_wal3 (id) VALUES (1)")
784            .execute(conn)
785            .unwrap();
786
787        assert_eq!(count.load(Ordering::Relaxed), 0);
788    }
789
790    #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
791    // Registers a callback that is invoked by the native library, which is not
792    // supported when running under miri with a native libsqlite3 (`-Zmiri-native-lib`).
793    #[cfg(not(miri))] // ffi call
794    #[diesel_test_helper::test]
795    fn wal_hook_does_not_fire_in_default_journal_mode() {
796        // A plain file connection left in the default journal mode ("delete"),
797        // not WAL, so `wal_connection()` (which enables WAL) is not used here.
798        let dir = tempfile::tempdir().unwrap();
799        let path = dir.path().join("test.db");
800        let conn = &mut SqliteConnection::establish(path.to_str().unwrap()).unwrap();
801
802        crate::sql_query("CREATE TABLE t_wal4 (id INTEGER PRIMARY KEY)")
803            .execute(conn)
804            .unwrap();
805
806        let count = Arc::new(AtomicU32::new(0));
807        let c2 = count.clone();
808
809        conn.on_wal(move |_, _, _| {
810            c2.fetch_add(1, Ordering::Relaxed);
811        });
812
813        crate::sql_query("INSERT INTO t_wal4 (id) VALUES (1)")
814            .execute(conn)
815            .unwrap();
816
817        assert_eq!(
818            count.load(Ordering::Relaxed),
819            0,
820            "WAL hook should not fire when not in WAL mode"
821        );
822    }
823
824    #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
825    // Registers a callback that is invoked by the native library, which is not
826    // supported when running under miri with a native libsqlite3 (`-Zmiri-native-lib`).
827    #[cfg(not(miri))] // ffi call
828    #[diesel_test_helper::test]
829    fn on_wal_can_use_borrowed_connection() {
830        let (conn, _dir) = &mut wal_connection();
831
832        crate::sql_query("CREATE TABLE t_wal_use (id INTEGER PRIMARY KEY)")
833            .execute(conn)
834            .unwrap();
835
836        let counts: Arc<std::sync::Mutex<Vec<i64>>> = Arc::new(std::sync::Mutex::new(Vec::new()));
837        let counts2 = counts.clone();
838
839        conn.on_wal(move |conn, _db_name, _n_pages| {
840            // Read through the borrowed connection, a fresh implicit read
841            // transaction that finalizes on return.
842            let c = crate::sql_query("SELECT COUNT(*) AS c FROM t_wal_use")
843                .get_result::<CountResult>(conn)
844                .unwrap()
845                .c;
846            counts2.lock().unwrap().push(c);
847        });
848
849        crate::sql_query("INSERT INTO t_wal_use (id) VALUES (1)")
850            .execute(conn)
851            .unwrap();
852
853        let observed = counts.lock().unwrap();
854        assert!(!observed.is_empty(), "WAL hook should have fired");
855        assert!(
856            observed.contains(&1),
857            "callback should observe the committed row through the connection"
858        );
859    }
860
861    #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
862    // Registers a callback that is invoked by the native library, which is not
863    // supported when running under miri with a native libsqlite3 (`-Zmiri-native-lib`).
864    #[cfg(not(miri))] // ffi call
865    #[diesel_test_helper::test]
866    fn on_wal_leaked_failed_cursor_keeps_connection_open() {
867        let (conn, _dir) = &mut wal_connection();
868        crate::sql_query("CREATE TABLE t_wal_leak (id INTEGER PRIMARY KEY)")
869            .execute(conn)
870            .unwrap();
871        #[allow(unsafe_code, reason = "Test code")]
872        conn.on_wal(|conn, _db_name, _n_pages| unsafe {
873            crate::sqlite::connection::leak_failed_cursor(conn.inner_mut());
874        });
875
876        for _ in 0..2 {
877            crate::sql_query("INSERT INTO t_wal_leak DEFAULT VALUES")
878                .execute(conn)
879                .unwrap();
880        }
881    }
882
883    #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
884    // Registers a callback that is invoked by the native library, which is not
885    // supported when running under miri with a native libsqlite3 (`-Zmiri-native-lib`).
886    #[cfg(not(miri))] // ffi call
887    #[diesel_test_helper::test]
888    fn on_wal_callback_write_re_enters_hook() {
889        let (conn, _dir) = &mut wal_connection();
890
891        crate::sql_query("CREATE TABLE t_wal_re (id INTEGER PRIMARY KEY)")
892            .execute(conn)
893            .unwrap();
894        crate::sql_query("CREATE TABLE t_wal_log (id INTEGER PRIMARY KEY AUTOINCREMENT)")
895            .execute(conn)
896            .unwrap();
897
898        let calls = Arc::new(AtomicU32::new(0));
899        let calls2 = calls.clone();
900
901        conn.on_wal(move |conn, _db_name, _n_pages| {
902            let n = calls2.fetch_add(1, Ordering::Relaxed);
903            // Write on the first invocation only. The write commits in WAL mode
904            // and re-fires the hook re-entrantly. Gating on `n == 0` bounds the
905            // recursion to a single nested call instead of overflowing the stack.
906            if n == 0 {
907                crate::sql_query("INSERT INTO t_wal_log DEFAULT VALUES")
908                    .execute(conn)
909                    .unwrap();
910            }
911        });
912
913        crate::sql_query("INSERT INTO t_wal_re (id) VALUES (1)")
914            .execute(conn)
915            .unwrap();
916
917        // The outer commit fired the hook, and the write inside it re-fired the
918        // hook once more, so the `Fn` callback ran twice.
919        assert_eq!(
920            calls.load(Ordering::Relaxed),
921            2,
922            "a committing write inside the callback re-enters the hook"
923        );
924
925        // The write performed inside the callback took effect.
926        let logged: i64 = crate::sql_query("SELECT COUNT(*) AS c FROM t_wal_log")
927            .get_result::<CountResult>(conn)
928            .unwrap()
929            .c;
930        assert_eq!(
931            logged, 1,
932            "the write performed inside the callback should persist"
933        );
934    }
935
936    #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
937    // Registers a callback that is invoked by the native library, which is not
938    // supported when running under miri with a native libsqlite3 (`-Zmiri-native-lib`).
939    #[cfg(not(miri))] // ffi call
940    #[diesel_test_helper::test]
941    fn on_wal_db_name_survives_detach_in_hook() {
942        // Raw SQL because the typed DSL cannot name a schema attached at runtime.
943        for schema in [
944            "wal_detach_arc".to_string(),
945            format!("wal_{}", "a".repeat(2000)),
946        ] {
947            let dir = tempfile::tempdir().unwrap();
948            let mut conn =
949                SqliteConnection::establish(dir.path().join("main.db").to_str().unwrap()).unwrap();
950            let archive = dir.path().join("archive.db");
951            crate::sql_query(format!(
952                "ATTACH DATABASE '{}' AS {schema}",
953                archive.display()
954            ))
955            .execute(&mut conn)
956            .unwrap();
957            crate::sql_query(format!("PRAGMA {schema}.journal_mode = WAL"))
958                .execute(&mut conn)
959                .unwrap();
960            crate::sql_query(format!("CREATE TABLE {schema}.log (line TEXT NOT NULL)"))
961                .execute(&mut conn)
962                .unwrap();
963
964            let observed: Arc<std::sync::Mutex<Option<(String, String)>>> = Default::default();
965            let observed2 = observed.clone();
966            conn.on_wal(move |conn, db_name, _| {
967                let at_entry = db_name.to_owned();
968                crate::sql_query(format!("DETACH DATABASE {at_entry}"))
969                    .execute(conn)
970                    .unwrap();
971                // `db_name` must stay readable after the DETACH ran.
972                *observed2.lock().unwrap() = Some((at_entry, db_name.to_owned()));
973            });
974
975            crate::sql_query(format!("INSERT INTO {schema}.log (line) VALUES ('hello')"))
976                .execute(&mut conn)
977                .unwrap();
978
979            let (at_entry, after) = observed.lock().unwrap().take().expect("WAL hook fired");
980            assert_eq!(at_entry, schema, "the hook fires for the attached schema");
981            assert_eq!(
982                after, at_entry,
983                "db_name must stay valid after a DETACH inside the hook"
984            );
985        }
986    }
987
988    #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
989    // Registers a callback that is invoked by the native library, which is not
990    // supported when running under miri with a native libsqlite3 (`-Zmiri-native-lib`).
991    #[cfg(not(miri))] // ffi call
992    #[diesel_test_helper::test]
993    fn on_wal_fires_once_per_transaction_commit() {
994        let (conn, _dir) = &mut wal_connection();
995
996        crate::sql_query("CREATE TABLE t_wal_txn (id INTEGER PRIMARY KEY)")
997            .execute(conn)
998            .unwrap();
999
1000        let count = Arc::new(AtomicU32::new(0));
1001        let c2 = count.clone();
1002
1003        conn.on_wal(move |_, _, _| {
1004            c2.fetch_add(1, Ordering::Relaxed);
1005        });
1006
1007        // Several writes inside one explicit transaction commit together, so the
1008        // WAL hook fires once for the single commit, not once per statement.
1009        conn.immediate_transaction(|conn| {
1010            crate::sql_query("INSERT INTO t_wal_txn (id) VALUES (1)").execute(conn)?;
1011            crate::sql_query("INSERT INTO t_wal_txn (id) VALUES (2)").execute(conn)?;
1012            crate::sql_query("INSERT INTO t_wal_txn (id) VALUES (3)").execute(conn)?;
1013            Ok::<_, crate::result::Error>(())
1014        })
1015        .unwrap();
1016
1017        assert_eq!(
1018            count.load(Ordering::Relaxed),
1019            1,
1020            "the WAL hook should fire once per transaction commit, not per statement"
1021        );
1022    }
1023
1024    // Busy handler test.
1025    //
1026    // Gated out on WASM: it needs two connections to a shared file-backed
1027    // database (`:memory:` connections do not share a lock), and
1028    // `tempfile::tempdir()` panics on WASM due to the lack of a filesystem.
1029    #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
1030    // Registers a callback that is invoked by the native library, which is not
1031    // supported when running under miri with a native libsqlite3 (`-Zmiri-native-lib`).
1032    #[cfg(not(miri))] // ffi call
1033    #[diesel_test_helper::test]
1034    fn on_busy_handler_is_invoked_on_lock_contention() {
1035        let dir = tempfile::tempdir().unwrap();
1036        let path = dir.path().join("busy.db");
1037        let url = path.to_str().unwrap();
1038
1039        // One connection acquires and holds a write lock.
1040        let mut holder = SqliteConnection::establish(url).unwrap();
1041        crate::sql_query("CREATE TABLE t_busy (id INTEGER PRIMARY KEY)")
1042            .execute(&mut holder)
1043            .unwrap();
1044        crate::sql_query("BEGIN IMMEDIATE")
1045            .execute(&mut holder)
1046            .unwrap();
1047
1048        // A second connection registers a busy handler that records the call
1049        // and gives up.
1050        let mut contender = SqliteConnection::establish(url).unwrap();
1051        let calls = Arc::new(AtomicU32::new(0));
1052        let calls2 = calls.clone();
1053        contender.on_busy(move |_retry_count| {
1054            calls2.fetch_add(1, Ordering::Relaxed);
1055            BusyDecision::GiveUp
1056        });
1057
1058        // The write contends with the held lock, so the busy handler fires.
1059        // Because it gives up, the write fails instead of blocking.
1060        let result = crate::sql_query("INSERT INTO t_busy (id) VALUES (1)").execute(&mut contender);
1061
1062        assert!(
1063            result.is_err(),
1064            "the contended write should fail once the busy handler gives up"
1065        );
1066        assert!(
1067            calls.load(Ordering::Relaxed) >= 1,
1068            "the busy handler should have been invoked at least once"
1069        );
1070    }
1071
1072    // A hook closure whose destructor panics while a normal connection is
1073    // dropped must not leave the handle open (and the database locked).
1074    // Cannot run this test with miri due to callbacks
1075    #[cfg(not(any(miri, all(target_family = "wasm", target_os = "unknown"))))]
1076    #[diesel_test_helper::test]
1077    fn panicking_hook_drop_on_connection_drop_still_closes_handle() {
1078        use crate::connection::SimpleConnection;
1079        use crate::sqlite::{SqliteChangeOps, SqliteUpdateRouter};
1080        use std::sync::Arc;
1081        use std::sync::atomic::{AtomicU32, Ordering};
1082
1083        struct PanicOnDrop;
1084        impl Drop for PanicOnDrop {
1085            fn drop(&mut self) {
1086                panic!("hook closure dropped");
1087            }
1088        }
1089
1090        let dir = tempfile::tempdir().unwrap();
1091        let path = dir.path().join("test.db");
1092        let path = path.to_str().unwrap();
1093        let mut conn = SqliteConnection::establish(path).unwrap();
1094        conn.batch_execute("CREATE TABLE t_panic_hook (x)").unwrap();
1095
1096        // The update hook is the first one `RawConnection::drop` frees.
1097        let guard = PanicOnDrop;
1098        conn.on_update(
1099            SqliteUpdateRouter::new().on_any(SqliteChangeOps::ALL, move |_| {
1100                let _ = &guard;
1101            }),
1102        );
1103        let rollbacks = Arc::new(AtomicU32::new(0));
1104        let r = rollbacks.clone();
1105        conn.on_rollback(move || {
1106            r.fetch_add(1, Ordering::SeqCst);
1107        });
1108        // Leave a write transaction open, so a leaked handle keeps the file locked.
1109        conn.batch_execute("BEGIN; INSERT INTO t_panic_hook VALUES (1)")
1110            .unwrap();
1111
1112        let dropped = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| drop(conn)));
1113        assert!(dropped.is_err(), "the closure destructor panics");
1114
1115        let mut other = SqliteConnection::establish(path).unwrap();
1116        let write = other.batch_execute("INSERT INTO t_panic_hook VALUES (2)");
1117        assert!(
1118            write.is_ok(),
1119            "the first handle was never closed: {write:?}"
1120        );
1121    }
1122}