feat: add preupdate hook

Author: aschey

Cargo.lock

@@ -58,7 +58,8 @@ default-features = false
 features = [
     "pkg-config",
     "vcpkg",
-    "unlock_notify"
+    "unlock_notify",
+    "preupdate_hook"
 ]
 
 [dependencies.sqlx-core]

sqlx-sqlite/Cargo.toml

@@ -296,6 +296,7 @@ impl EstablishParams {
             log_settings: self.log_settings.clone(),
             progress_handler_callback: None,
             update_hook_callback: None,
+            preupdate_hook_callback: None,
             commit_hook_callback: None,
             rollback_hook_callback: None,
         })

sqlx-sqlite/src/connection/establish.rs

@@ -11,8 +11,10 @@ use futures_core::future::BoxFuture;
 use futures_intrusive::sync::MutexGuard;
 use futures_util::future;
 use libsqlite3_sys::{
-    sqlite3, sqlite3_commit_hook, sqlite3_progress_handler, sqlite3_rollback_hook,
-    sqlite3_update_hook, SQLITE_DELETE, SQLITE_INSERT, SQLITE_UPDATE,
+    sqlite3, sqlite3_commit_hook, sqlite3_preupdate_count, sqlite3_preupdate_depth,
+    sqlite3_preupdate_hook, sqlite3_preupdate_new, sqlite3_preupdate_old, sqlite3_progress_handler,
+    sqlite3_rollback_hook, sqlite3_update_hook, sqlite3_value, sqlite3_value_type, SQLITE_DELETE,
+    SQLITE_INSERT, SQLITE_OK, SQLITE_UPDATE,
 };
 
 pub(crate) use handle::ConnectionHandle;
@@ -26,7 +28,8 @@ use crate::connection::establish::EstablishParams;
 use crate::connection::worker::ConnectionWorker;
 use crate::options::OptimizeOnClose;
 use crate::statement::VirtualStatement;
-use crate::{Sqlite, SqliteConnectOptions};
+use crate::type_info::DataType;
+use crate::{Sqlite, SqliteConnectOptions, SqliteError, SqliteTypeInfo, SqliteValue};
 
 pub(crate) mod collation;
 pub(crate) mod describe;
@@ -88,6 +91,14 @@ pub struct UpdateHookResult<'a> {
     pub table: &'a str,
     pub rowid: i64,
 }
+
+pub struct PreupdateHookResult<'a> {
+    pub operation: SqliteOperation,
+    pub database: &'a str,
+    pub table: &'a str,
+    pub case: PreupdateCase,
+}
+
 pub(crate) struct UpdateHookHandler(NonNull<dyn FnMut(UpdateHookResult) + Send + 'static>);
 unsafe impl Send for UpdateHookHandler {}
 
@@ -97,6 +108,108 @@ unsafe impl Send for CommitHookHandler {}
 pub(crate) struct RollbackHookHandler(NonNull<dyn FnMut() + Send + 'static>);
 unsafe impl Send for RollbackHookHandler {}
 
+pub(crate) struct PreupdateHookHandler(NonNull<dyn FnMut(PreupdateHookResult) + Send + 'static>);
+unsafe impl Send for PreupdateHookHandler {}
+
+/// The possible cases for when a PreUpdate Hook gets triggered. Allows access to the relevant
+/// functions for each case through the contained values.
+pub enum PreupdateCase {
+    /// Pre-update hook was triggered by an insert.
+    Insert(PreupdateNewValueAccessor),
+    /// Pre-update hook was triggered by a delete.
+    Delete(PreupdateOldValueAccessor),
+    /// Pre-update hook was triggered by an update.
+    Update {
+        old_value_accessor: PreupdateOldValueAccessor,
+        new_value_accessor: PreupdateNewValueAccessor,
+    },
+    /// This variant is not normally produced by SQLite. You may encounter it
+    /// if you're using a different version than what's supported by this library.
+    Unknown,
+}
+
+/// An accessor for the new values of the row being inserted/updated during the preupdate callback.
+#[derive(Debug)]
+pub struct PreupdateNewValueAccessor {
+    db: *mut sqlite3,
+    new_row_id: i64,
+}
+
+impl PreupdateNewValueAccessor {
+    /// Gets the amount of columns in the row being inserted/updated.
+    pub fn get_column_count(&self) -> i32 {
+        unsafe { sqlite3_preupdate_count(self.db) }
+    }
+
+    /// Gets the depth of the query that triggered the preupdate hook.
+    /// Returns 0 if the preupdate callback was invoked as a result of
+    /// a direct insert, update, or delete operation;
+    /// 1 for inserts, updates, or deletes invoked by top-level triggers;
+    /// 2 for changes resulting from triggers called by top-level triggers; and so forth.
+    pub fn get_query_depth(&self) -> i32 {
+        unsafe { sqlite3_preupdate_depth(self.db) }
+    }
+
+    /// Gets the row id of the row being inserted/updated.
+    pub fn get_new_row_id(&self) -> i64 {
+        self.new_row_id
+    }
+
+    /// Gets the value of the row being updated/deleted at the specified index.
+    pub fn get_new_column_value(&self, i: i32) -> Result<SqliteValue, Error> {
+        let mut p_value: *mut sqlite3_value = ptr::null_mut();
+        unsafe {
+            let ret = sqlite3_preupdate_new(self.db, i, &mut p_value);
+            if ret != SQLITE_OK {
+                return Err(Error::Database(Box::new(SqliteError::new(self.db))));
+            }
+            let data_type = DataType::from_code(sqlite3_value_type(p_value));
+            Ok(SqliteValue::new(p_value, SqliteTypeInfo(data_type)))
+        }
+    }
+}
+
+/// An accessor for the old values of the row being deleted/updated during the preupdate callback.
+#[derive(Debug)]
+pub struct PreupdateOldValueAccessor {
+    db: *mut sqlite3,
+    old_row_id: i64,
+}
+
+impl PreupdateOldValueAccessor {
+    /// Gets the amount of columns in the row being deleted/updated.
+    pub fn get_column_count(&self) -> i32 {
+        unsafe { sqlite3_preupdate_count(self.db) }
+    }
+
+    /// Gets the depth of the query that triggered the preupdate hook.
+    /// Returns 0 if the preupdate callback was invoked as a result of
+    /// a direct insert, update, or delete operation;
+    /// 1 for inserts, updates, or deletes invoked by top-level triggers;
+    /// 2 for changes resulting from triggers called by top-level triggers; and so forth.
+    pub fn get_query_depth(&self) -> i32 {
+        unsafe { sqlite3_preupdate_depth(self.db) }
+    }
+
+    /// Gets the row id of the row being updated/deleted.
+    pub fn get_old_row_id(&self) -> i64 {
+        self.old_row_id
+    }
+
+    /// Gets the value of the row being updated/deleted at the specified index.
+    pub fn get_old_column_value(&self, i: i32) -> Result<SqliteValue, Error> {
+        let mut p_value: *mut sqlite3_value = ptr::null_mut();
+        unsafe {
+            let ret = sqlite3_preupdate_old(self.db, i, &mut p_value);
+            if ret != SQLITE_OK {
+                return Err(Error::Database(Box::new(SqliteError::new(self.db))));
+            }
+            let data_type = DataType::from_code(sqlite3_value_type(p_value));
+            Ok(SqliteValue::new(p_value, SqliteTypeInfo(data_type)))
+        }
+    }
+}
+
 pub(crate) struct ConnectionState {
     pub(crate) handle: ConnectionHandle,
 
@@ -113,6 +226,8 @@ pub(crate) struct ConnectionState {
 
     update_hook_callback: Option<UpdateHookHandler>,
 
+    preupdate_hook_callback: Option<PreupdateHookHandler>,
+
     commit_hook_callback: Option<CommitHookHandler>,
 
     rollback_hook_callback: Option<RollbackHookHandler>,
@@ -138,6 +253,15 @@ impl ConnectionState {
         }
     }
 
+    pub(crate) fn remove_preupdate_hook(&mut self) {
+        if let Some(mut handler) = self.preupdate_hook_callback.take() {
+            unsafe {
+                sqlite3_preupdate_hook(self.handle.as_ptr(), None, ptr::null_mut());
+                let _ = { Box::from_raw(handler.0.as_mut()) };
+            }
+        }
+    }
+
     pub(crate) fn remove_commit_hook(&mut self) {
         if let Some(mut handler) = self.commit_hook_callback.take() {
             unsafe {
@@ -312,6 +436,47 @@ extern "C" fn update_hook<F>(
     }
 }
 
+extern "C" fn preupdate_hook<F>(
+    callback: *mut c_void,
+    db: *mut sqlite3,
+    op_code: c_int,
+    database: *const i8,
+    table: *const i8,
+    old_row_id: i64,
+    new_row_id: i64,
+) where
+    F: FnMut(PreupdateHookResult),
+{
+    unsafe {
+        let _ = catch_unwind(|| {
+            let callback: *mut F = callback.cast::<F>();
+            let operation: SqliteOperation = op_code.into();
+            let database = CStr::from_ptr(database).to_str().unwrap_or_default();
+            let table = CStr::from_ptr(table).to_str().unwrap_or_default();
+
+            let preupdate_case = match operation {
+                SqliteOperation::Insert => {
+                    PreupdateCase::Insert(PreupdateNewValueAccessor { db, new_row_id })
+                }
+                SqliteOperation::Delete => {
+                    PreupdateCase::Delete(PreupdateOldValueAccessor { db, old_row_id })
+                }
+                SqliteOperation::Update => PreupdateCase::Update {
+                    old_value_accessor: PreupdateOldValueAccessor { db, old_row_id },
+                    new_value_accessor: PreupdateNewValueAccessor { db, new_row_id },
+                },
+                SqliteOperation::Unknown(_) => PreupdateCase::Unknown,
+            };
+            (*callback)(PreupdateHookResult {
+                operation,
+                database,
+                table,
+                case: preupdate_case,
+            })
+        });
+    }
+}
+
 extern "C" fn commit_hook<F>(callback: *mut c_void) -> c_int
 where
     F: FnMut() -> bool,
@@ -476,6 +641,33 @@ impl LockedSqliteHandle<'_> {
         }
     }
 
+    /// Registers a hook that is invoked prior to each `INSERT`, `UPDATE`, and `DELETE` operation on a database table.
+    /// At most one preupdate hook may be registered at a time on a single database connection.
+    ///
+    /// The preupdate hook only fires for changes to real database tables;
+    /// it is not invoked for changes to virtual tables or to system tables like sqlite_sequence or sqlite_stat1.
+    ///
+    /// See https://sqlite.org/c3ref/preupdate_count.html
+    pub fn set_preupdate_hook<F>(&mut self, callback: F)
+    where
+        F: FnMut(PreupdateHookResult) + Send + 'static,
+    {
+        unsafe {
+            let callback_boxed = Box::new(callback);
+            // SAFETY: `Box::into_raw()` always returns a non-null pointer.
+            let callback = NonNull::new_unchecked(Box::into_raw(callback_boxed));
+            let handler = callback.as_ptr() as *mut _;
+            self.guard.remove_preupdate_hook();
+            self.guard.preupdate_hook_callback = Some(PreupdateHookHandler(callback));
+
+            sqlite3_preupdate_hook(
+                self.as_raw_handle().as_mut(),
+                Some(preupdate_hook::<F>),
+                handler,
+            );
+        }
+    }
+
     /// Removes the progress handler on a database connection. The method does nothing if no handler was set.
     pub fn remove_progress_handler(&mut self) {
         self.guard.remove_progress_handler();
@@ -492,6 +684,10 @@ impl LockedSqliteHandle<'_> {
     pub fn remove_rollback_hook(&mut self) {
         self.guard.remove_rollback_hook();
     }
+
+    pub fn remove_preupdate_hook(&mut self) {
+        self.guard.remove_preupdate_hook();
+    }
 }
 
 impl Drop for ConnectionState {

sqlx-sqlite/src/connection/mod.rs

@@ -46,7 +46,10 @@ use std::sync::atomic::AtomicBool;
 
 pub use arguments::{SqliteArgumentValue, SqliteArguments};
 pub use column::SqliteColumn;
-pub use connection::{LockedSqliteHandle, SqliteConnection, SqliteOperation, UpdateHookResult};
+pub use connection::{
+    LockedSqliteHandle, PreupdateCase, PreupdateHookResult, PreupdateNewValueAccessor,
+    PreupdateOldValueAccessor, SqliteConnection, SqliteOperation, UpdateHookResult,
+};
 pub use database::Sqlite;
 pub use error::SqliteError;
 pub use options::{