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