feat: add begin_with to start transaction with custom SQL

sqlx-core/src/acquire.rs

@@ -1,6 +1,7 @@
 use crate::database::Database;
 use crate::error::Error;
 use crate::pool::{MaybePoolConnection, Pool, PoolConnection};
+use std::borrow::Cow;
 
 use crate::transaction::Transaction;
 use futures_core::future::BoxFuture;
@@ -78,6 +79,10 @@ pub trait Acquire<'c> {
     fn acquire(self) -> BoxFuture<'c, Result<Self::Connection, Error>>;
 
     fn begin(self) -> BoxFuture<'c, Result<Transaction<'c, Self::Database>, Error>>;
+
+    fn begin_with<S>(self, sql: S) -> BoxFuture<'c, Result<Transaction<'c, Self::Database>, Error>>
+    where
+        S: Into<Cow<'static, str>> + Send + 'c;
 }
 
 impl<'a, DB: Database> Acquire<'a> for &'_ Pool<DB> {
@@ -96,6 +101,17 @@ impl<'a, DB: Database> Acquire<'a> for &'_ Pool<DB> {
             Transaction::begin(MaybePoolConnection::PoolConnection(conn.await?)).await
         })
     }
+
+    fn begin_with<S>(self, sql: S) -> BoxFuture<'a, Result<Transaction<'a, Self::Database>, Error>>
+    where
+        S: Into<Cow<'static, str>> + Send + 'a,
+    {
+        let conn = self.acquire();
+
+        Box::pin(async move {
+            Transaction::begin_with(MaybePoolConnection::PoolConnection(conn.await?), sql).await
+        })
+    }
 }
 
 #[macro_export]
@@ -123,6 +139,20 @@ macro_rules! impl_acquire {
             > {
                 $crate::transaction::Transaction::begin(self)
             }
+
+            #[inline]
+            fn begin_with<S>(
+                self,
+                stmt: S,
+            ) -> futures_core::future::BoxFuture<
+                'c,
+                Result<$crate::transaction::Transaction<'c, $DB>, $crate::error::Error>,
+            >
+            where
+                S: Into<std::borrow::Cow<'static, str>> + Send + 'c,
+            {
+                $crate::transaction::Transaction::begin_with(self, stmt)
+            }
         }
     };
 }

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 {
@@ -30,6 +31,11 @@ pub trait AnyConnectionBackend: std::any::Any + Debug + Send + 'static {
     /// Returns a [`Transaction`] for controlling and tracking the new transaction.
     fn begin(&mut self) -> BoxFuture<'_, crate::Result<()>>;
 
+    /// Begin a new transaction or establish a savepoint within the active transaction.
+    ///
+    /// Returns a [`Transaction`] for controlling and tracking the new transaction.
+    fn begin_with(&mut self, sql: Cow<'static, str>) -> BoxFuture<'_, crate::Result<()>>;
+
     fn commit(&mut self) -> BoxFuture<'_, crate::Result<()>>;
 
     fn rollback(&mut self) -> BoxFuture<'_, crate::Result<()>>;

sqlx-core/src/any/transaction.rs

@@ -1,6 +1,8 @@
 use futures_util::future::BoxFuture;
+use std::borrow::Cow;
 
 use crate::any::{Any, AnyConnection};
+use crate::database::Database;
 use crate::error::Error;
 use crate::transaction::TransactionManager;
 
@@ -13,6 +15,16 @@ impl TransactionManager for AnyTransactionManager {
         conn.backend.begin()
     }
 
+    fn begin_with<'a, S>(
+        conn: &'a mut <Self::Database as Database>::Connection,
+        sql: S,
+    ) -> BoxFuture<'a, Result<(), Error>>
+    where
+        S: Into<Cow<'static, str>> + Send + 'a,
+    {
+        conn.backend.begin_with(sql.into())
+    }
+
     fn commit(conn: &mut AnyConnection) -> BoxFuture<'_, Result<(), Error>> {
         conn.backend.commit()
     }

sqlx-core/src/pool/connection.rs

@@ -1,3 +1,4 @@
+use std::borrow::Cow;
 use std::fmt::{self, Debug, Formatter};
 use std::ops::{Deref, DerefMut};
 use std::sync::Arc;
@@ -11,6 +12,8 @@ use crate::error::Error;
 
 use super::inner::{is_beyond_max_lifetime, DecrementSizeGuard, PoolInner};
 use crate::pool::options::PoolConnectionMetadata;
+use crate::transaction::Transaction;
+use futures_core::future::BoxFuture;
 use std::future::Future;
 
 /// A connection managed by a [`Pool`][crate::pool::Pool].
@@ -159,6 +162,13 @@ impl<'c, DB: Database> crate::acquire::Acquire<'c> for &'c mut PoolConnection<DB
     {
         crate::transaction::Transaction::begin(&mut **self)
     }
+
+    fn begin_with<S>(self, sql: S) -> BoxFuture<'c, Result<Transaction<'c, Self::Database>, Error>>
+    where
+        S: Into<Cow<'static, str>> + Send + 'c,
+    {
+        crate::transaction::Transaction::begin_with(&mut **self, sql)
+    }
 }
 
 /// Returns the connection to the [`Pool`][crate::pool::Pool] it was checked-out from.

sqlx-core/src/transaction.rs

@@ -20,6 +20,13 @@ pub trait TransactionManager {
         conn: &mut <Self::Database as Database>::Connection,
     ) -> BoxFuture<'_, Result<(), Error>>;
 
+    fn begin_with<'a, S>(
+        conn: &'a mut <Self::Database as Database>::Connection,
+        sql: S,
+    ) -> BoxFuture<'a, Result<(), Error>>
+    where
+        S: Into<Cow<'static, str>> + Send + 'a;
+
     /// Commit the active transaction or release the most recent savepoint.
     fn commit(
         conn: &mut <Self::Database as Database>::Connection,
@@ -78,6 +85,26 @@ where
         })
     }
 
+    #[doc(hidden)]
+    pub fn begin_with<S>(
+        conn: impl Into<MaybePoolConnection<'c, DB>>,
+        sql: S,
+    ) -> BoxFuture<'c, Result<Self, Error>>
+    where
+        S: Into<Cow<'static, str>> + Send + 'c,
+    {
+        let mut conn = conn.into();
+
+        Box::pin(async move {
+            DB::TransactionManager::begin_with(&mut conn, sql).await?;
+
+            Ok(Self {
+                connection: conn,
+                open: true,
+            })
+        })
+    }
+
     /// Commits this transaction or savepoint.
     pub async fn commit(mut self) -> Result<(), Error> {
         DB::TransactionManager::commit(&mut self.connection).await?;
@@ -221,6 +248,13 @@ impl<'c, 't, DB: Database> crate::acquire::Acquire<'t> for &'t mut Transaction<'
     fn begin(self) -> BoxFuture<'t, Result<Transaction<'t, DB>, Error>> {
         Transaction::begin(&mut **self)
     }
+
+    fn begin_with<S>(self, sql: S) -> BoxFuture<'t, Result<Transaction<'t, Self::Database>, Error>>
+    where
+        S: Into<Cow<'static, str>> + Send + 't,
+    {
+        Transaction::begin_with(&mut **self, sql)
+    }
 }
 
 impl<'c, DB> Drop for Transaction<'c, DB>

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;
 
 sqlx_core::declare_driver_with_optional_migrate!(DRIVER = MySql);
 
@@ -40,6 +41,10 @@ impl AnyConnectionBackend for MySqlConnection {
         MySqlTransactionManager::begin(self)
     }
 
+    fn begin_with(&mut self, sql: Cow<'static, str>) -> BoxFuture<'_, sqlx_core::Result<()>> {
+        MySqlTransactionManager::begin_with(self, sql)
+    }
+
     fn commit(&mut self) -> BoxFuture<'_, sqlx_core::Result<()>> {
         MySqlTransactionManager::commit(self)
     }

sqlx-mysql/src/transaction.rs

@@ -1,10 +1,11 @@
-use futures_core::future::BoxFuture;
-
 use crate::connection::Waiting;
 use crate::error::Error;
 use crate::executor::Executor;
 use crate::protocol::text::Query;
 use crate::{MySql, MySqlConnection};
+use futures_core::future::BoxFuture;
+use sqlx_core::database::Database;
+use std::borrow::Cow;
 
 pub(crate) use sqlx_core::transaction::*;
 
@@ -25,6 +26,23 @@ impl TransactionManager for MySqlTransactionManager {
         })
     }
 
+    fn begin_with<'a, S>(
+        conn: &'a mut <Self::Database as Database>::Connection,
+        sql: S,
+    ) -> BoxFuture<'a, Result<(), Error>>
+    where
+        S: Into<Cow<'static, str>> + Send + 'a,
+    {
+        Box::pin(async move {
+            let depth = conn.transaction_depth;
+
+            conn.execute(&*sql.into()).await?;
+            conn.transaction_depth = depth + 1;
+
+            Ok(())
+        })
+    }
+
     fn commit(conn: &mut MySqlConnection) -> BoxFuture<'_, Result<(), Error>> {
         Box::pin(async move {
             let depth = conn.transaction_depth;

sqlx-postgres/src/any.rs

@@ -5,6 +5,7 @@ use crate::{
 use futures_core::future::BoxFuture;
 use futures_core::stream::BoxStream;
 use futures_util::{StreamExt, TryFutureExt, TryStreamExt};
+use std::borrow::Cow;
 
 pub use sqlx_core::any::*;
 
@@ -39,6 +40,10 @@ impl AnyConnectionBackend for PgConnection {
         PgTransactionManager::begin(self)
     }
 
+    fn begin_with(&mut self, sql: Cow<'static, str>) -> BoxFuture<'_, sqlx_core::Result<()>> {
+        PgTransactionManager::begin_with(self, sql)
+    }
+
     fn commit(&mut self) -> BoxFuture<'_, sqlx_core::Result<()>> {
         PgTransactionManager::commit(self)
     }

sqlx-postgres/src/transaction.rs

@@ -1,7 +1,8 @@
-use futures_core::future::BoxFuture;
-
 use crate::error::Error;
 use crate::executor::Executor;
+use futures_core::future::BoxFuture;
+use sqlx_core::database::Database;
+use std::borrow::Cow;
 
 use crate::{PgConnection, Postgres};
 
@@ -26,6 +27,24 @@ impl TransactionManager for PgTransactionManager {
         })
     }
 
+    fn begin_with<'a, S>(
+        conn: &'a mut <Self::Database as Database>::Connection,
+        sql: S,
+    ) -> BoxFuture<'a, Result<(), Error>>
+    where
+        S: Into<Cow<'static, str>> + Send + 'a,
+    {
+        Box::pin(async move {
+            let rollback = Rollback::new(conn);
+            rollback.conn.queue_simple_query(&sql.into());
+            rollback.conn.transaction_depth += 1;
+            rollback.conn.wait_until_ready().await?;
+            rollback.defuse();
+
+            Ok(())
+        })
+    }
+
     fn commit(conn: &mut PgConnection) -> BoxFuture<'_, Result<(), Error>> {
         Box::pin(async move {
             if conn.transaction_depth > 0 {

sqlx-sqlite/src/any.rs

@@ -5,6 +5,7 @@ use crate::{
 use futures_core::future::BoxFuture;
 use futures_core::stream::BoxStream;
 use futures_util::{StreamExt, TryFutureExt, TryStreamExt};
+use std::borrow::Cow;
 
 use sqlx_core::any::{
     Any, AnyArguments, AnyColumn, AnyConnectOptions, AnyConnectionBackend, AnyQueryResult, AnyRow,
@@ -41,6 +42,10 @@ impl AnyConnectionBackend for SqliteConnection {
         SqliteTransactionManager::begin(self)
     }
 
+    fn begin_with(&mut self, sql: Cow<'static, str>) -> BoxFuture<'_, sqlx_core::Result<()>> {
+        todo!()
+    }
+
     fn commit(&mut self) -> BoxFuture<'_, sqlx_core::Result<()>> {
         SqliteTransactionManager::commit(self)
     }

sqlx-sqlite/src/transaction.rs

@@ -1,8 +1,9 @@
-use futures_core::future::BoxFuture;
-
 use crate::{Sqlite, SqliteConnection};
+use futures_core::future::BoxFuture;
+use sqlx_core::database::Database;
 use sqlx_core::error::Error;
 use sqlx_core::transaction::TransactionManager;
+use std::borrow::Cow;
 
 /// Implementation of [`TransactionManager`] for SQLite.
 pub struct SqliteTransactionManager;
@@ -14,6 +15,16 @@ impl TransactionManager for SqliteTransactionManager {
         Box::pin(conn.worker.begin())
     }
 
+    fn begin_with<'a, S>(
+        _conn: &'a mut <Self::Database as Database>::Connection,
+        sql: S,
+    ) -> BoxFuture<'a, Result<(), Error>>
+    where
+        S: Into<Cow<'static, str>> + Send + 'a,
+    {
+        unimplemented!()
+    }
+
     fn commit(conn: &mut SqliteConnection) -> BoxFuture<'_, Result<(), Error>> {
         Box::pin(conn.worker.commit())
     }