Add begin_with methods to support database-specific transaction options

Cargo.toml

@@ -184,6 +184,7 @@ rand_xoshiro = "0.6.0"
 hex = "0.4.3"
 tempfile = "3.10.1"
 criterion = { version = "0.5.1", features = ["async_tokio"] }
+libsqlite3-sys = { version = "0.30.1" }
 
 # If this is an unconditional dev-dependency then Cargo will *always* try to build `libsqlite3-sys`,
 # even when SQLite isn't the intended test target, and fail if the build environment is not set up for compiling C code.

sqlx-core/src/acquire.rs

@@ -93,7 +93,7 @@ impl<'a, DB: Database> Acquire<'a> for &'_ Pool<DB> {
         let conn = self.acquire();
 
         Box::pin(async move {
-            Transaction::begin(MaybePoolConnection::PoolConnection(conn.await?)).await
+            Transaction::begin(MaybePoolConnection::PoolConnection(conn.await?), None).await
         })
     }
 }
@@ -121,7 +121,7 @@ macro_rules! impl_acquire {
                 'c,
                 Result<$crate::transaction::Transaction<'c, $DB>, $crate::error::Error>,
             > {
-                $crate::transaction::Transaction::begin(self)
+                $crate::transaction::Transaction::begin(self, None)
             }
         }
     };

sqlx-core/src/any/connection/backend.rs

@@ -3,6 +3,7 @@ use crate::describe::Describe;
 use either::Either;
 use futures_core::future::BoxFuture;
 use futures_core::stream::BoxStream;
+use std::borrow::Cow;
 use std::fmt::Debug;
 
 pub trait AnyConnectionBackend: std::any::Any + Debug + Send + 'static {
@@ -26,7 +27,13 @@ pub trait AnyConnectionBackend: std::any::Any + Debug + Send + 'static {
     fn ping(&mut self) -> BoxFuture<'_, crate::Result<()>>;
 
     /// Begin a new transaction or establish a savepoint within the active transaction.
-    fn begin(&mut self) -> BoxFuture<'_, crate::Result<()>>;
+    ///
+    /// If this is a new transaction, `statement` may be used instead of the
+    /// default "BEGIN" statement.
+    ///
+    /// If we are already inside a transaction and `statement.is_some()`, then
+    /// `Error::InvalidSavePoint` is returned without running any statements.
+    fn begin(&mut self, statement: Option<Cow<'static, str>>) -> BoxFuture<'_, crate::Result<()>>;
 
     fn commit(&mut self) -> BoxFuture<'_, crate::Result<()>>;
 

sqlx-core/src/any/connection/mod.rs

@@ -1,4 +1,5 @@
 use futures_core::future::BoxFuture;
+use std::borrow::Cow;
 
 use crate::any::{Any, AnyConnectOptions};
 use crate::connection::{ConnectOptions, Connection};
@@ -87,7 +88,17 @@ impl Connection for AnyConnection {
     where
         Self: Sized,
     {
-        Transaction::begin(self)
+        Transaction::begin(self, None)
+    }
+
+    fn begin_with(
+        &mut self,
+        statement: impl Into<Cow<'static, str>>,
+    ) -> BoxFuture<'_, Result<Transaction<'_, Self::Database>, Error>>
+    where
+        Self: Sized,
+    {
+        Transaction::begin(self, Some(statement.into()))
     }
 
     fn cached_statements_size(&self) -> usize {

sqlx-core/src/any/transaction.rs

@@ -1,4 +1,5 @@
 use futures_util::future::BoxFuture;
+use std::borrow::Cow;
 
 use crate::any::{Any, AnyConnection};
 use crate::error::Error;
@@ -9,8 +10,11 @@ pub struct AnyTransactionManager;
 impl TransactionManager for AnyTransactionManager {
     type Database = Any;
 
-    fn begin(conn: &mut AnyConnection) -> BoxFuture<'_, Result<(), Error>> {
-        conn.backend.begin()
+    fn begin<'conn>(
+        conn: &'conn mut AnyConnection,
+        statement: Option<Cow<'static, str>>,
+    ) -> BoxFuture<'conn, Result<(), Error>> {
+        conn.backend.begin(statement)
     }
 
     fn commit(conn: &mut AnyConnection) -> BoxFuture<'_, Result<(), Error>> {

sqlx-core/src/connection.rs

@@ -4,6 +4,7 @@ use crate::error::Error;
 use crate::transaction::Transaction;
 use futures_core::future::BoxFuture;
 use log::LevelFilter;
+use std::borrow::Cow;
 use std::fmt::Debug;
 use std::str::FromStr;
 use std::time::Duration;
@@ -49,6 +50,22 @@ pub trait Connection: Send {
     where
         Self: Sized;
 
+    /// Begin a new transaction with a custom statement.
+    ///
+    /// Returns a [`Transaction`] for controlling and tracking the new transaction.
+    ///
+    /// Returns an error if the connection is already in a transaction or if
+    /// `statement` does not put the connection into a transaction.
+    fn begin_with(
+        &mut self,
+        statement: impl Into<Cow<'static, str>>,
+    ) -> BoxFuture<'_, Result<Transaction<'_, Self::Database>, Error>>
+    where
+        Self: Sized,
+    {
+        Transaction::begin(self, Some(statement.into()))
+    }
+
     /// Execute the function inside a transaction.
     ///
     /// If the function returns an error, the transaction will be rolled back. If it does not

sqlx-core/src/error.rs

@@ -111,6 +111,12 @@ pub enum Error {
     #[cfg(feature = "migrate")]
     #[error("{0}")]
     Migrate(#[source] Box<crate::migrate::MigrateError>),
+
+    #[error("attempted to call begin_with at non-zero transaction depth")]
+    InvalidSavePointStatement,
+
+    #[error("got unexpected connection status after attempting to begin transaction")]
+    BeginFailed,
 }
 
 impl StdError for Box<dyn DatabaseError> {}

sqlx-core/src/pool/connection.rs

@@ -191,7 +191,7 @@ impl<'c, DB: Database> crate::acquire::Acquire<'c> for &'c mut PoolConnection<DB
         self,
     ) -> futures_core::future::BoxFuture<'c, Result<crate::transaction::Transaction<'c, DB>, Error>>
     {
-        crate::transaction::Transaction::begin(&mut **self)
+        crate::transaction::Transaction::begin(&mut **self, None)
     }
 }
 

sqlx-core/src/pool/mod.rs

@@ -54,6 +54,7 @@
 //! [`Pool::acquire`] or
 //! [`Pool::begin`].
 
+use std::borrow::Cow;
 use std::fmt;
 use std::future::Future;
 use std::pin::Pin;
@@ -367,20 +368,54 @@ impl<DB: Database> Pool<DB> {
 
     /// Retrieves a connection and immediately begins a new transaction.
     pub async fn begin(&self) -> Result<Transaction<'static, DB>, Error> {
-        Transaction::begin(MaybePoolConnection::PoolConnection(self.acquire().await?)).await
+        Transaction::begin(
+            MaybePoolConnection::PoolConnection(self.acquire().await?),
+            None,
+        )
+        .await
     }
 
     /// Attempts to retrieve a connection and immediately begins a new transaction if successful.
     pub async fn try_begin(&self) -> Result<Option<Transaction<'static, DB>>, Error> {
         match self.try_acquire() {
-            Some(conn) => Transaction::begin(MaybePoolConnection::PoolConnection(conn))
+            Some(conn) => Transaction::begin(MaybePoolConnection::PoolConnection(conn), None)
                 .await
                 .map(Some),
 
             None => Ok(None),
         }
     }
 
+    /// Retrieves a connection and immediately begins a new transaction using `statement`.
+    pub async fn begin_with(
+        &self,
+        statement: impl Into<Cow<'static, str>>,
+    ) -> Result<Transaction<'static, DB>, Error> {
+        Transaction::begin(
+            MaybePoolConnection::PoolConnection(self.acquire().await?),
+            Some(statement.into()),
+        )
+        .await
+    }
+
+    /// Attempts to retrieve a connection and, if successful, immediately begins a new
+    /// transaction using `statement`.
+    pub async fn try_begin_with(
+        &self,
+        statement: impl Into<Cow<'static, str>>,
+    ) -> Result<Option<Transaction<'static, DB>>, Error> {
+        match self.try_acquire() {
+            Some(conn) => Transaction::begin(
+                MaybePoolConnection::PoolConnection(conn),
+                Some(statement.into()),
+            )
+            .await
+            .map(Some),
+
+            None => Ok(None),
+        }
+    }
+
     /// Shut down the connection pool, immediately waking all tasks waiting for a connection.
     ///
     /// Upon calling this method, any currently waiting or subsequent calls to [`Pool::acquire`] and

sqlx-core/src/transaction.rs

@@ -16,9 +16,16 @@ pub trait TransactionManager {
     type Database: Database;
 
     /// Begin a new transaction or establish a savepoint within the active transaction.
-    fn begin(
-        conn: &mut <Self::Database as Database>::Connection,
-    ) -> BoxFuture<'_, Result<(), Error>>;
+    ///
+    /// If this is a new transaction, `statement` may be used instead of the
+    /// default "BEGIN" statement.
+    ///
+    /// If we are already inside a transaction and `statement.is_some()`, then
+    /// `Error::InvalidSavePoint` is returned without running any statements.
+    fn begin<'conn>(
+        conn: &'conn mut <Self::Database as Database>::Connection,
+        statement: Option<Cow<'static, str>>,
+    ) -> BoxFuture<'conn, Result<(), Error>>;
 
     /// Commit the active transaction or release the most recent savepoint.
     fn commit(
@@ -83,11 +90,12 @@ where
     #[doc(hidden)]
     pub fn begin(
         conn: impl Into<MaybePoolConnection<'c, DB>>,
+        statement: Option<Cow<'static, str>>,
     ) -> BoxFuture<'c, Result<Self, Error>> {
         let mut conn = conn.into();
 
         Box::pin(async move {
-            DB::TransactionManager::begin(&mut conn).await?;
+            DB::TransactionManager::begin(&mut conn, statement).await?;
 
             Ok(Self {
                 connection: conn,
@@ -237,7 +245,7 @@ impl<'c, 't, DB: Database> crate::acquire::Acquire<'t> for &'t mut Transaction<'
 
     #[inline]
     fn begin(self) -> BoxFuture<'t, Result<Transaction<'t, DB>, Error>> {
-        Transaction::begin(&mut **self)
+        Transaction::begin(&mut **self, None)
     }
 }
 

sqlx-mysql/src/any.rs

@@ -16,6 +16,7 @@ use sqlx_core::database::Database;
 use sqlx_core::describe::Describe;
 use sqlx_core::executor::Executor;
 use sqlx_core::transaction::TransactionManager;
+use std::borrow::Cow;
 use std::future;
 
 sqlx_core::declare_driver_with_optional_migrate!(DRIVER = MySql);
@@ -37,8 +38,11 @@ impl AnyConnectionBackend for MySqlConnection {
         Connection::ping(self)
     }
 
-    fn begin(&mut self) -> BoxFuture<'_, sqlx_core::Result<()>> {
-        MySqlTransactionManager::begin(self)
+    fn begin(
+        &mut self,
+        statement: Option<Cow<'static, str>>,
+    ) -> BoxFuture<'_, sqlx_core::Result<()>> {
+        MySqlTransactionManager::begin(self, statement)
     }
 
     fn commit(&mut self) -> BoxFuture<'_, sqlx_core::Result<()>> {

sqlx-mysql/src/connection/establish.rs

@@ -28,6 +28,7 @@ impl MySqlConnection {
             inner: Box::new(MySqlConnectionInner {
                 stream,
                 transaction_depth: 0,
+                status_flags: Default::default(),
                 cache_statement: StatementCache::new(options.statement_cache_capacity),
                 log_settings: options.log_settings.clone(),
             }),

sqlx-mysql/src/connection/executor.rs

@@ -166,6 +166,8 @@ impl MySqlConnection {
                     // this indicates either a successful query with no rows at all or a failed query
                     let ok = packet.ok()?;
 
+                    self.inner.status_flags = ok.status;
+
                     let rows_affected = ok.affected_rows;
                     logger.increase_rows_affected(rows_affected);
                     let done = MySqlQueryResult {
@@ -208,6 +210,8 @@ impl MySqlConnection {
                     if packet[0] == 0xfe && packet.len() < 9 {
                         let eof = packet.eof(self.inner.stream.capabilities)?;
 
+                        self.inner.status_flags = eof.status;
+
                         r#yield!(Either::Left(MySqlQueryResult {
                             rows_affected: 0,
                             last_insert_id: 0,