Skip to content

Commit

Permalink
feat: add preupdate hook
Browse files Browse the repository at this point in the history
  • Loading branch information
aschey committed Dec 4, 2024
1 parent 42ce24d commit 5264ada
Show file tree
Hide file tree
Showing 6 changed files with 303 additions and 15 deletions.
12 changes: 6 additions & 6 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 2 additions & 1 deletion sqlx-sqlite/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,8 @@ default-features = false
features = [
"pkg-config",
"vcpkg",
"unlock_notify"
"unlock_notify",
"preupdate_hook"
]

[dependencies.sqlx-core]
Expand Down
1 change: 1 addition & 0 deletions sqlx-sqlite/src/connection/establish.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
})
Expand Down
202 changes: 199 additions & 3 deletions sqlx-sqlite/src/connection/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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_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;
Expand All @@ -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;
Expand Down Expand Up @@ -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 {}

Expand All @@ -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 libsqlite3_sys::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 { libsqlite3_sys::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 { libsqlite3_sys::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_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)))
}
}
}

pub(crate) struct ConnectionState {
pub(crate) handle: ConnectionHandle,

Expand All @@ -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>,
Expand All @@ -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 {
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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();
Expand All @@ -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 {
Expand Down
5 changes: 4 additions & 1 deletion sqlx-sqlite/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down
Loading

0 comments on commit 5264ada

Please sign in to comment.