feat: add support for nested transaction rollbacks via savepoints in sql

quaint/src/connector/mssql.rs

@@ -11,6 +11,7 @@ use crate::{
 use async_trait::async_trait;
 use connection_string::JdbcString;
 use futures::lock::Mutex;
+use std::sync::Arc;
 use std::{
     convert::TryFrom,
     fmt,
@@ -106,11 +107,23 @@ impl TransactionCapable for Mssql {
             .or(self.url.query_params.transaction_isolation_level)
             .or(Some(SQL_SERVER_DEFAULT_ISOLATION));
 
-        let opts = TransactionOptions::new(isolation, self.requires_isolation_first());
+        let mut transaction_depth = self.transaction_depth.lock().await;
+        *transaction_depth += 1;
+        let st_depth = *transaction_depth;
 
-        Ok(Box::new(
-            DefaultTransaction::new(self, self.begin_statement(), opts).await?,
-        ))
+        let begin_statement = self.begin_statement(st_depth).await;
+        let commit_stmt = self.commit_statement(st_depth).await;
+        let rollback_stmt = self.rollback_statement(st_depth).await;
+
+        let opts = TransactionOptions::new(
+            isolation,
+            self.requires_isolation_first(),
+            self.transaction_depth.clone(),
+            commit_stmt,
+            rollback_stmt,
+        );
+
+        Ok(Box::new(DefaultTransaction::new(self, &begin_statement, opts).await?))
     }
 }
 
@@ -273,6 +286,7 @@ pub struct Mssql {
     url: MssqlUrl,
     socket_timeout: Option<Duration>,
     is_healthy: AtomicBool,
+    transaction_depth: Arc<Mutex<i32>>,
 }
 
 impl Mssql {
@@ -304,6 +318,7 @@ impl Mssql {
             url,
             socket_timeout,
             is_healthy: AtomicBool::new(true),
+            transaction_depth: Arc::new(Mutex::new(0)),
         };
 
         if let Some(isolation) = this.url.transaction_isolation_level() {
@@ -443,8 +458,41 @@ impl Queryable for Mssql {
         Ok(())
     }
 
-    fn begin_statement(&self) -> &'static str {
-        "BEGIN TRAN"
+    /// Statement to begin a transaction
+    async fn begin_statement(&self, depth: i32) -> String {
+        let savepoint_stmt = format!("SAVE TRANSACTION savepoint{}", depth);
+        let ret = if depth > 1 {
+            savepoint_stmt
+        } else {
+            "BEGIN TRAN".to_string()
+        };
+
+        return ret;
+    }
+
+    /// Statement to commit a transaction
+    async fn commit_statement(&self, depth: i32) -> String {
+        // MSSQL doesn't have a "RELEASE SAVEPOINT" equivalent, so in a nested
+        // transaction we just continue onwards
+        let ret = if depth > 1 {
+            " ".to_string()
+        } else {
+            "COMMIT".to_string()
+        };
+
+        return ret;
+    }
+
+    /// Statement to rollback a transaction
+    async fn rollback_statement(&self, depth: i32) -> String {
+        let savepoint_stmt = format!("ROLLBACK TRANSACTION savepoint{}", depth);
+        let ret = if depth > 1 {
+            savepoint_stmt
+        } else {
+            "ROLLBACK".to_string()
+        };
+
+        return ret;
     }
 
     fn requires_isolation_first(&self) -> bool {

quaint/src/connector/mysql.rs

@@ -14,6 +14,7 @@ use mysql_async::{
     prelude::{Query as _, Queryable as _},
 };
 use percent_encoding::percent_decode;
+use std::sync::Arc;
 use std::{
     borrow::Cow,
     future::Future,
@@ -41,6 +42,7 @@ pub struct Mysql {
     socket_timeout: Option<Duration>,
     is_healthy: AtomicBool,
     statement_cache: Mutex<LruCache<String, my::Statement>>,
+    transaction_depth: Arc<futures::lock::Mutex<i32>>,
 }
 
 /// Wraps a connection url and exposes the parsing logic used by quaint, including default values.
@@ -376,6 +378,7 @@ impl Mysql {
             statement_cache: Mutex::new(url.cache()),
             url,
             is_healthy: AtomicBool::new(true),
+            transaction_depth: Arc::new(futures::lock::Mutex::new(0)),
         })
     }
 
@@ -583,6 +586,38 @@ impl Queryable for Mysql {
     fn requires_isolation_first(&self) -> bool {
         true
     }
+
+    /// Statement to begin a transaction
+    async fn begin_statement(&self, depth: i32) -> String {
+        let savepoint_stmt = format!("SAVEPOINT savepoint{}", depth);
+        let ret = if depth > 1 { savepoint_stmt } else { "BEGIN".to_string() };
+
+        return ret;
+    }
+
+    /// Statement to commit a transaction
+    async fn commit_statement(&self, depth: i32) -> String {
+        let savepoint_stmt = format!("RELEASE SAVEPOINT savepoint{}", depth);
+        let ret = if depth > 1 {
+            savepoint_stmt
+        } else {
+            "COMMIT".to_string()
+        };
+
+        return ret;
+    }
+
+    /// Statement to rollback a transaction
+    async fn rollback_statement(&self, depth: i32) -> String {
+        let savepoint_stmt = format!("ROLLBACK TO savepoint{}", depth);
+        let ret = if depth > 1 {
+            savepoint_stmt
+        } else {
+            "ROLLBACK".to_string()
+        };
+
+        return ret;
+    }
 }
 
 #[cfg(test)]

quaint/src/connector/postgres.rs

@@ -13,6 +13,7 @@ use lru_cache::LruCache;
 use native_tls::{Certificate, Identity, TlsConnector};
 use percent_encoding::percent_decode;
 use postgres_native_tls::MakeTlsConnector;
+use std::sync::Arc;
 use std::{
     borrow::{Borrow, Cow},
     fmt::{Debug, Display},
@@ -63,6 +64,7 @@ pub struct PostgreSql {
     socket_timeout: Option<Duration>,
     statement_cache: Mutex<LruCache<String, Statement>>,
     is_healthy: AtomicBool,
+    transaction_depth: Arc<Mutex<i32>>,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -652,6 +654,7 @@ impl PostgreSql {
             pg_bouncer: url.query_params.pg_bouncer,
             statement_cache: Mutex::new(url.cache()),
             is_healthy: AtomicBool::new(true),
+            transaction_depth: Arc::new(Mutex::new(0)),
         })
     }
 
@@ -932,6 +935,38 @@ impl Queryable for PostgreSql {
     fn requires_isolation_first(&self) -> bool {
         false
     }
+
+    /// Statement to begin a transaction
+    async fn begin_statement(&self, depth: i32) -> String {
+        let savepoint_stmt = format!("SAVEPOINT savepoint{}", depth);
+        let ret = if depth > 1 { savepoint_stmt } else { "BEGIN".to_string() };
+
+        return ret;
+    }
+
+    /// Statement to commit a transaction
+    async fn commit_statement(&self, depth: i32) -> String {
+        let savepoint_stmt = format!("RELEASE SAVEPOINT savepoint{}", depth);
+        let ret = if depth > 1 {
+            savepoint_stmt
+        } else {
+            "COMMIT".to_string()
+        };
+
+        return ret;
+    }
+
+    /// Statement to rollback a transaction
+    async fn rollback_statement(&self, depth: i32) -> String {
+        let savepoint_stmt = format!("ROLLBACK TO SAVEPOINT savepoint{}", depth);
+        let ret = if depth > 1 {
+            savepoint_stmt
+        } else {
+            "ROLLBACK".to_string()
+        };
+
+        return ret;
+    }
 }
 
 /// Sorted list of CockroachDB's reserved keywords.

quaint/src/connector/queryable.rs

@@ -87,8 +87,18 @@ pub trait Queryable: Send + Sync {
     }
 
     /// Statement to begin a transaction
-    fn begin_statement(&self) -> &'static str {
-        "BEGIN"
+    async fn begin_statement(&self, _depth: i32) -> String {
+        "BEGIN".to_string()
+    }
+
+    /// Statement to commit a transaction
+    async fn commit_statement(&self, _depth: i32) -> String {
+        "COMMIT".to_string()
+    }
+
+    /// Statement to rollback a transaction
+    async fn rollback_statement(&self, _depth: i32) -> String {
+        "ROLLBACK".to_string()
     }
 
     /// Sets the transaction isolation level to given value.
@@ -117,10 +127,26 @@ macro_rules! impl_default_TransactionCapable {
                 &'a self,
                 isolation: Option<IsolationLevel>,
             ) -> crate::Result<Box<dyn crate::connector::Transaction + 'a>> {
-                let opts = crate::connector::TransactionOptions::new(isolation, self.requires_isolation_first());
+                let depth = self.transaction_depth.clone();
+                let mut depth_guard = self.transaction_depth.lock().await;
+                *depth_guard += 1;
+
+                let st_depth = *depth_guard;
+
+                let begin_statement = self.begin_statement(st_depth).await;
+                let commit_stmt = self.commit_statement(st_depth).await;
+                let rollback_stmt = self.rollback_statement(st_depth).await;
+
+                let opts = crate::connector::TransactionOptions::new(
+                    isolation,
+                    self.requires_isolation_first(),
+                    depth,
+                    commit_stmt,
+                    rollback_stmt,
+                );
 
                 Ok(Box::new(
-                    crate::connector::DefaultTransaction::new(self, self.begin_statement(), opts).await?,
+                    crate::connector::DefaultTransaction::new(self, &begin_statement, opts).await?,
                 ))
             }
         }

quaint/src/connector/sqlite.rs

@@ -13,6 +13,7 @@ use crate::{
     visitor::{self, Visitor},
 };
 use async_trait::async_trait;
+use std::sync::Arc;
 use std::{convert::TryFrom, path::Path, time::Duration};
 use tokio::sync::Mutex;
 
@@ -25,6 +26,7 @@ pub use rusqlite;
 /// A connector interface for the SQLite database
 pub struct Sqlite {
     pub(crate) client: Mutex<rusqlite::Connection>,
+    transaction_depth: Arc<futures::lock::Mutex<i32>>,
 }
 
 /// Wraps a connection url and exposes the parsing logic used by Quaint,
@@ -141,7 +143,10 @@ impl TryFrom<&str> for Sqlite {
 
         let client = Mutex::new(conn);
 
-        Ok(Sqlite { client })
+        Ok(Sqlite {
+            client,
+            transaction_depth: Arc::new(futures::lock::Mutex::new(0)),
+        })
     }
 }
 
@@ -156,6 +161,7 @@ impl Sqlite {
 
         Ok(Sqlite {
             client: Mutex::new(client),
+            transaction_depth: Arc::new(futures::lock::Mutex::new(0)),
         })
     }
 
@@ -252,6 +258,38 @@ impl Queryable for Sqlite {
     fn requires_isolation_first(&self) -> bool {
         false
     }
+
+    /// Statement to begin a transaction
+    async fn begin_statement(&self, depth: i32) -> String {
+        let savepoint_stmt = format!("SAVEPOINT savepoint{}", depth);
+        let ret = if depth > 1 { savepoint_stmt } else { "BEGIN".to_string() };
+
+        return ret;
+    }
+
+    /// Statement to commit a transaction
+    async fn commit_statement(&self, depth: i32) -> String {
+        let savepoint_stmt = format!("RELEASE SAVEPOINT savepoint{}", depth);
+        let ret = if depth > 1 {
+            savepoint_stmt
+        } else {
+            "COMMIT".to_string()
+        };
+
+        return ret;
+    }
+
+    /// Statement to rollback a transaction
+    async fn rollback_statement(&self, depth: i32) -> String {
+        let savepoint_stmt = format!("ROLLBACK TO savepoint{}", depth);
+        let ret = if depth > 1 {
+            savepoint_stmt
+        } else {
+            "ROLLBACK".to_string()
+        };
+
+        return ret;
+    }
 }
 
 #[cfg(test)]