From cee97e97298a2bce777b2d027cd49946c6735cc4 Mon Sep 17 00:00:00 2001
From: gwenn <gtreguier@gmail.com>
Date: Sun, 10 Nov 2024 15:40:57 +0100
Subject: [PATCH] Make possible to checkpoint a database from `wal_hook`

---
 src/hooks/mod.rs | 58 +++++++++++++++++++++++++++++++++++-------------
 1 file changed, 42 insertions(+), 16 deletions(-)

diff --git a/src/hooks/mod.rs b/src/hooks/mod.rs
index 9d14aad..38d79dc 100644
--- a/src/hooks/mod.rs
+++ b/src/hooks/mod.rs
@@ -7,7 +7,7 @@ use std::ptr;
 
 use crate::ffi;
 
-use crate::{Connection, DatabaseName, InnerConnection};
+use crate::{error, Connection, DatabaseName, InnerConnection, Result};
 
 #[cfg(feature = "preupdate_hook")]
 pub use preupdate_hook::*;
@@ -389,23 +389,22 @@ impl Connection {
     /// Calling `wal_hook` replaces any previously registered write-ahead log callback.
     /// Note that the `sqlite3_wal_autocheckpoint()` interface and the `wal_autocheckpoint` pragma
     /// both invoke `sqlite3_wal_hook()` and will overwrite any prior `sqlite3_wal_hook()` settings.
-    pub fn wal_hook(&self, hook: Option<fn(DatabaseName<'_>, c_int) -> c_int>) {
+    pub fn wal_hook(&self, hook: Option<fn(&Wal, c_int) -> Result<()>>) {
         unsafe extern "C" fn wal_hook_callback(
             client_data: *mut c_void,
-            _db: *mut ffi::sqlite3,
+            db: *mut ffi::sqlite3,
             db_name: *const c_char,
             pages: c_int,
         ) -> c_int {
-            let hook_fn: fn(DatabaseName<'_>, c_int) -> c_int = std::mem::transmute(client_data);
-            c_int::from(
-                catch_unwind(|| {
-                    hook_fn(
-                        DatabaseName::from_cstr(std::ffi::CStr::from_ptr(db_name)),
-                        pages,
-                    )
-                })
-                .unwrap_or_default(),
-            )
+            let hook_fn: fn(&Wal, c_int) -> Result<()> = std::mem::transmute(client_data);
+            let wal = Wal { db, db_name };
+            catch_unwind(|| match hook_fn(&wal, pages) {
+                Ok(_) => ffi::SQLITE_OK,
+                Err(e) => e
+                    .sqlite_error()
+                    .map_or(ffi::SQLITE_ERROR, |x| x.extended_code),
+            })
+            .unwrap_or_default()
         }
         let c = self.db.borrow_mut();
         match hook {
@@ -442,6 +441,33 @@ impl Connection {
     }
 }
 
+/// Write-Ahead Log
+pub struct Wal {
+    db: *mut ffi::sqlite3,
+    db_name: *const c_char,
+}
+
+impl Wal {
+    /// Checkpoint a database
+    pub fn checkpoint(&self) -> Result<()> {
+        error::check(unsafe { ffi::sqlite3_wal_checkpoint(self.db, self.db_name) })
+    }
+    /// Checkpoint a database
+    pub fn checkpoint_v2(&self, mode: c_int) -> Result<(c_int, c_int)> {
+        let mut n_log = 0;
+        let mut n_ckpt = 0;
+        error::check(unsafe {
+            ffi::sqlite3_wal_checkpoint_v2(self.db, self.db_name, mode, &mut n_log, &mut n_ckpt)
+        })?;
+        Ok((n_log, n_ckpt))
+    }
+
+    /// Name of the database that was written to
+    pub fn name(&self) -> DatabaseName<'_> {
+        DatabaseName::from_cstr(unsafe { std::ffi::CStr::from_ptr(self.db_name) })
+    }
+}
+
 impl InnerConnection {
     #[inline]
     pub fn remove_hooks(&mut self) {
@@ -942,11 +968,11 @@ mod test {
         assert_eq!(journal_mode, "wal");
 
         static CALLED: AtomicBool = AtomicBool::new(false);
-        db.wal_hook(Some(|db_name, pages| {
-            assert_eq!(db_name, DatabaseName::Main);
+        db.wal_hook(Some(|wal, pages| {
+            assert_eq!(wal.name(), DatabaseName::Main);
             assert!(pages > 0);
             CALLED.swap(true, Ordering::Relaxed);
-            crate::ffi::SQLITE_OK
+            wal.checkpoint()
         }));
         db.execute_batch("CREATE TABLE x(c);")?;
         assert!(CALLED.load(Ordering::Relaxed));