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}