Skip to content

Commit

Permalink
address some PR comments
Browse files Browse the repository at this point in the history
  • Loading branch information
aschey committed Dec 14, 2024
1 parent 7f1fdd9 commit 46df6f8
Show file tree
Hide file tree
Showing 7 changed files with 276 additions and 302 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

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

1 change: 1 addition & 0 deletions sqlx-sqlite/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ atoi = "2.0"

log = "0.4.18"
tracing = { version = "0.1.37", features = ["log"] }
thiserror = "2.0.0"

serde = { version = "1.0.145", features = ["derive"], optional = true }
regex = { version = "1.5.5", optional = true }
Expand Down
261 changes: 45 additions & 216 deletions sqlx-sqlite/src/connection/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,8 @@ mod executor;
mod explain;
mod handle;
pub(crate) mod intmap;
#[cfg(feature = "preupdate-hook")]
mod preupdate_hook;

mod worker;

Expand Down Expand Up @@ -143,6 +145,16 @@ impl ConnectionState {
}
}

#[cfg(feature = "preupdate-hook")]
pub(crate) fn remove_preupdate_hook(&mut self) {
if let Some(mut handler) = self.preupdate_hook_callback.take() {
unsafe {
libsqlite3_sys::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 @@ -426,6 +438,34 @@ 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
#[cfg(feature = "preupdate-hook")]
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));

libsqlite3_sys::sqlite3_preupdate_hook(
self.as_raw_handle().as_mut(),
Some(preupdate_hook::<F>),
handler,
);
}
}

/// Sets a commit hook that is invoked whenever a transaction is committed. If the commit hook callback
/// returns `false`, then the operation is turned into a ROLLBACK.
///
Expand Down Expand Up @@ -490,6 +530,11 @@ impl LockedSqliteHandle<'_> {
self.guard.remove_update_hook();
}

#[cfg(feature = "preupdate-hook")]
pub fn remove_preupdate_hook(&mut self) {
self.guard.remove_preupdate_hook();
}

pub fn remove_commit_hook(&mut self) {
self.guard.remove_commit_hook();
}
Expand Down Expand Up @@ -549,219 +594,3 @@ impl Statements {
self.temp = None;
}
}

#[cfg(feature = "preupdate-hook")]
mod preupdate_hook {
use super::ConnectionState;
use super::LockedSqliteHandle;
use super::SqliteOperation;
use crate::type_info::DataType;
use crate::{SqliteError, SqliteTypeInfo, SqliteValue};
use libsqlite3_sys::{
sqlite3, sqlite3_preupdate_count, sqlite3_preupdate_depth, sqlite3_preupdate_hook,
sqlite3_preupdate_new, sqlite3_preupdate_old, sqlite3_value, sqlite3_value_type, SQLITE_OK,
};
use sqlx_core::error::Error;
use std::ffi::CStr;
use std::fmt::Debug;
use std::os::raw::{c_char, c_int, c_void};
use std::panic::catch_unwind;
use std::ptr;
use std::ptr::NonNull;

pub struct PreupdateHookResult<'a> {
pub operation: SqliteOperation,
pub database: &'a str,
pub table: &'a str,
pub case: PreupdateCase,
}

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 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)))
}
}
}

/// 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)))
}
}
}

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()) };
}
}
}
}

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,
);
}
}

pub fn remove_preupdate_hook(&mut self) {
self.guard.remove_preupdate_hook();
}
}

extern "C" fn preupdate_hook<F>(
callback: *mut c_void,
db: *mut sqlite3,
op_code: c_int,
database: *const c_char,
table: *const c_char,
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,
})
});
}
}
}
Loading

0 comments on commit 46df6f8

Please sign in to comment.