diff --git a/pgx-tests/src/tests/subxact_tests.rs b/pgx-tests/src/tests/subxact_tests.rs index 1a600b61a..ce6b70edd 100644 --- a/pgx-tests/src/tests/subxact_tests.rs +++ b/pgx-tests/src/tests/subxact_tests.rs @@ -57,7 +57,7 @@ mod tests { Spi::execute(|c| { c.update("CREATE TABLE a (v INTEGER)", None, None); // The type below is explicit to ensure it's commit on drop by default - c.sub_transaction(|xact: SubTransaction| { + c.sub_transaction(|xact: SubTransaction| { xact.update("INSERT INTO a VALUES (0)", None, None); // Dropped explicitly for illustration purposes drop(xact); @@ -81,7 +81,7 @@ mod tests { Spi::execute(|c| { c.update("CREATE TABLE a (v INTEGER)", None, None); // The type below is explicit to ensure it's commit on drop by default - c.sub_transaction(|xact: SubTransaction| { + c.sub_transaction(|xact: SubTransaction| { xact.update("INSERT INTO a VALUES (0)", None, None); let xact = xact.rollback_on_drop(); // Dropped explicitly for illustration purposes diff --git a/pgx/src/subxact.rs b/pgx/src/subxact.rs index f97d837c1..2e33a1608 100644 --- a/pgx/src/subxact.rs +++ b/pgx/src/subxact.rs @@ -2,152 +2,200 @@ use crate::{pg_sys, PgMemoryContexts, SpiClient}; use std::fmt::Debug; use std::ops::Deref; -/// Sub-transaction -/// -/// Can be created by calling `SpiClient::sub_transaction`, `SubTransaction::sub_transaction` -/// or any other implementation of `SubTransactionExt` and obtaining it as an argument to the provided closure. -/// -/// Unless rolled back or committed explicitly, it'll commit if `COMMIT` generic parameter is `true` -/// (default) or roll back if it is `false`. -#[derive(Debug)] -pub struct SubTransaction { +/// Releases a sub-transaction on Drop +pub trait ReleaseOnDrop {} + +/// Sub-transaction's contextual information +#[derive(Clone, Copy)] +pub struct Context { memory_context: pg_sys::MemoryContext, // Resource ownership before the transaction // // Based on information from src/backend/utils/resowner/README // as well as practical use of it in src/pl/plpython/plpy_spi.c resource_owner: pg_sys::ResourceOwner, - // Should the transaction be released, or was it already committed or rolled back? - // - // The reason we are not calling this `released` as we're also using this flag when - // we convert between commit_on_drop and rollback_on_drop to ensure it doesn't get released - // on the drop of the original value. - should_release: bool, - parent: Option, } -impl SubTransaction { - /// Create a new sub-transaction. - fn new(parent: Parent) -> Self { +impl Context { + /// Captures the context + fn capture() -> Self { // Remember the memory context before starting the sub-transaction - let ctx = PgMemoryContexts::CurrentMemoryContext.value(); + let memory_context = PgMemoryContexts::CurrentMemoryContext.value(); // Remember resource owner before starting the sub-transaction let resource_owner = unsafe { pg_sys::CurrentResourceOwner }; - unsafe { - pg_sys::BeginInternalSubTransaction(std::ptr::null() /* [no] transaction name */); - } - // Switch to the outer memory context so that all allocations remain - // there instead of the sub-transaction's context - PgMemoryContexts::For(ctx).set_as_current(); - Self { memory_context: ctx, should_release: true, resource_owner, parent: Some(parent) } + Self { memory_context, resource_owner } } +} - /// Commit the transaction, returning its parent - pub fn commit(mut self) -> Parent { - self.internal_commit(); - self.should_release = false; - self.parent.take().unwrap() +impl From for CommitOnDrop { + fn from(context: Context) -> Self { + CommitOnDrop(context) } +} - /// Rollback the transaction, returning its parent - pub fn rollback(mut self) -> Parent { - self.internal_rollback(); - self.should_release = false; - self.parent.take().unwrap() +impl From for RollbackOnDrop { + fn from(context: Context) -> Self { + RollbackOnDrop(context) } +} - fn internal_rollback(&self) { +/// Commits a sub-transaction on Drop +pub struct CommitOnDrop(Context); + +impl Drop for CommitOnDrop { + fn drop(&mut self) { unsafe { - pg_sys::RollbackAndReleaseCurrentSubTransaction(); - pg_sys::CurrentResourceOwner = self.resource_owner; + pg_sys::ReleaseCurrentSubTransaction(); + pg_sys::CurrentResourceOwner = self.0.resource_owner; } - PgMemoryContexts::For(self.memory_context).set_as_current(); + PgMemoryContexts::For(self.0.memory_context).set_as_current(); } +} - fn internal_commit(&self) { +impl ReleaseOnDrop for CommitOnDrop {} + +/// Rolls back a sub-transaction on Drop +pub struct RollbackOnDrop(Context); + +impl Drop for RollbackOnDrop { + fn drop(&mut self) { unsafe { - pg_sys::ReleaseCurrentSubTransaction(); - pg_sys::CurrentResourceOwner = self.resource_owner; + pg_sys::RollbackAndReleaseCurrentSubTransaction(); + pg_sys::CurrentResourceOwner = self.0.resource_owner; } - PgMemoryContexts::For(self.memory_context).set_as_current(); + PgMemoryContexts::For(self.0.memory_context).set_as_current(); } } -impl SubTransaction { - /// Make this sub-transaction roll back on drop - pub fn rollback_on_drop(self) -> SubTransaction { - self.into() +impl ReleaseOnDrop for RollbackOnDrop {} + +impl Into for CommitOnDrop { + fn into(self) -> RollbackOnDrop { + let result = RollbackOnDrop(self.0); + // IMPORTANT: avoid running Drop (that would commit) + std::mem::forget(self); + result } } -impl SubTransaction { - /// Make this sub-transaction commit on drop - pub fn commit_on_drop(self) -> SubTransaction { - self.into() +impl Into for RollbackOnDrop { + fn into(self) -> CommitOnDrop { + let result = CommitOnDrop(self.0); + // IMPORTANT: avoid running Drop (that would roll back) + std::mem::forget(self); + result } } -impl Into> - for SubTransaction +struct NoOpOnDrop; + +impl ReleaseOnDrop for NoOpOnDrop {} + +/// Sub-transaction +/// +/// Can be created by calling `SpiClient::sub_transaction`, `SubTransaction::sub_transaction` +/// or any other implementation of `SubTransactionExt` and obtaining it as an argument to the provided closure. +/// +/// Unless rolled back or committed explicitly, it'll commit if `Release` generic parameter is `CommitOnDrop` +/// (default) or roll back if it is `RollbackOnDrop`. +#[derive(Debug)] +pub struct SubTransaction { + release: Release, + parent: Parent, +} + +impl SubTransaction +where + Release: From, { - fn into(mut self) -> SubTransaction { - let result = SubTransaction { - memory_context: self.memory_context, - resource_owner: self.resource_owner, - should_release: self.should_release, - parent: self.parent.take(), - }; - // Make sure original sub-transaction won't commit - self.should_release = false; - result + /// Create a new sub-transaction. + fn new(parent: Parent) -> Self { + let context = Context::capture(); + let memory_context = context.memory_context; + let release = context.into(); + unsafe { + pg_sys::BeginInternalSubTransaction(std::ptr::null() /* [no] transaction name */); + } + // Switch to the outer memory context so that all allocations remain + // there instead of the sub-transaction's context + PgMemoryContexts::For(memory_context).set_as_current(); + Self { release, parent } } } -impl Into> - for SubTransaction -{ - fn into(mut self) -> SubTransaction { - let result = SubTransaction { - memory_context: self.memory_context, - resource_owner: self.resource_owner, - should_release: self.should_release, - parent: self.parent.take(), - }; - // Make sure original sub-transaction won't roll back - self.should_release = false; - result +impl SubTransaction { + /// Commit the transaction, returning its parent + pub fn commit(self) -> Parent { + // `Self::do_nothing_on_drop()` will commit as `Release` is `CommitOnDrop` + self.do_nothing_on_drop().parent } } -impl Drop for SubTransaction { - fn drop(&mut self) { - if self.should_release { - if COMMIT { - self.internal_commit(); - } else { - self.internal_rollback(); - } - } +impl SubTransaction { + /// Commit the transaction, returning its parent + pub fn commit(self) -> Parent { + // Make sub-transaction commit on drop and then use `commit` + self.commit_on_drop().commit() + } +} + +impl SubTransaction { + /// Rollback the transaction, returning its parent + pub fn rollback(self) -> Parent { + // `Self::do_nothing_on_drop()` will roll back as `Release` is `RollbackOnDrop` + self.do_nothing_on_drop().parent + } +} + +impl SubTransaction { + /// Rollback the transaction, returning its parent + pub fn rollback(self) -> Parent { + // Make sub-transaction roll back on drop and then use `rollback` + self.rollback_on_drop().rollback() + } +} + +impl SubTransaction { + /// Make this sub-transaction roll back on drop + pub fn rollback_on_drop(self) -> SubTransaction { + SubTransaction { parent: self.parent, release: self.release.into() } + } +} + +impl SubTransaction { + /// Make this sub-transaction commit on drop + pub fn commit_on_drop(self) -> SubTransaction { + SubTransaction { parent: self.parent, release: self.release.into() } + } +} + +impl SubTransaction { + /// Make this sub-transaction do nothing on drop + /// + /// Releases the sub-transaction based on `Release` generic parameter. Further + /// dropping of the sub-transaction will not do anything. + fn do_nothing_on_drop(self) -> SubTransaction { + SubTransaction { parent: self.parent, release: NoOpOnDrop } } } // This allows SubTransaction to be de-referenced to SpiClient -impl<'conn, const COMMIT: bool> Deref for SubTransaction, COMMIT> { +impl<'conn, Release: ReleaseOnDrop> Deref for SubTransaction, Release> { type Target = SpiClient<'conn>; fn deref(&self) -> &Self::Target { - self.parent.as_ref().unwrap() + &self.parent } } // This allows a SubTransaction of a SubTransaction to be de-referenced to SpiClient -impl Deref - for SubTransaction, COMMIT> +impl Deref + for SubTransaction, Release> { type Target = Parent; fn deref(&self) -> &Self::Target { - self.parent.as_ref().and_then(|p| p.parent.as_ref()).unwrap() + &self.parent.parent } }