Support savepoints (#547)

* Support savepoints

* Fix spelling

Author: levkk

integration/rust/tests/integration/mod.rs

@@ -9,6 +9,7 @@ pub mod notify;
 pub mod per_stmt_routing;
 pub mod prepared;
 pub mod reload;
+pub mod savepoint;
 pub mod set_sharding_key;
 pub mod shard_consistency;
 pub mod stddev;

integration/rust/tests/integration/savepoint.rs

@@ -0,0 +1,26 @@
+use rust::setup::{admin_sqlx, connections_sqlx};
+use sqlx::Executor;
+
+#[tokio::test]
+async fn test_savepoint() {
+    let conns = connections_sqlx().await;
+
+    for conn in conns {
+        let mut transaction = conn.begin().await.unwrap();
+        transaction
+            .execute("CREATE TABLE test_savepoint (id BIGINT)")
+            .await
+            .unwrap();
+        transaction.execute("SAVEPOINT test").await.unwrap();
+        assert!(transaction.execute("SELECT sdfsf").await.is_err());
+        transaction
+            .execute("ROLLBACK TO SAVEPOINT test")
+            .await
+            .unwrap();
+        transaction
+            .execute("SELECT * FROM test_savepoint")
+            .await
+            .unwrap();
+        transaction.rollback().await.unwrap();
+    }
+}

pgdog/src/frontend/client/mod.rs

@@ -56,7 +56,8 @@ pub enum TransactionType {
     ReadOnly,
     #[default]
     ReadWrite,
-    Error,
+    ErrorReadWrite,
+    ErrorReadOnly,
 }
 
 impl TransactionType {
@@ -69,7 +70,7 @@ impl TransactionType {
     }
 
     pub fn error(&self) -> bool {
-        matches!(self, Self::Error)
+        matches!(self, Self::ErrorReadWrite | Self::ErrorReadOnly)
     }
 }
 

pgdog/src/frontend/client/query_engine/query.rs

@@ -21,7 +21,7 @@ impl QueryEngine {
         route: &Route,
     ) -> Result<(), Error> {
         // Check that we're not in a transaction error state.
-        if !self.transaction_error_check(context).await? {
+        if !self.transaction_error_check(context, route).await? {
             return Ok(());
         }
 
@@ -115,7 +115,12 @@ impl QueryEngine {
 
             match state {
                 TransactionState::Error => {
-                    context.transaction = Some(TransactionType::Error);
+                    let error_state = match context.transaction {
+                        Some(TransactionType::ReadOnly) => Some(TransactionType::ErrorReadOnly),
+                        Some(TransactionType::ReadWrite) => Some(TransactionType::ErrorReadWrite),
+                        _ => None,
+                    };
+                    context.transaction = error_state;
                     if self.two_pc.auto() {
                         self.end_two_pc(true).await?;
                         // TODO: this records a 2pc transaction in client
@@ -133,10 +138,23 @@ impl QueryEngine {
                         self.end_two_pc(false).await?;
                         two_pc_auto = true;
                     }
-                    if context.transaction.is_none() {
+                    match context.transaction {
                         // Query parser is disabled, so the server is responsible for telling us
                         // we started a transaction.
-                        context.transaction = Some(TransactionType::ReadWrite);
+                        None => {
+                            context.transaction = Some(TransactionType::ReadWrite);
+                        }
+
+                        // Restore transaction state after rollback to savepoint.
+                        Some(TransactionType::ErrorReadOnly) => {
+                            context.transaction = Some(TransactionType::ReadOnly);
+                        }
+
+                        Some(TransactionType::ErrorReadWrite) => {
+                            context.transaction = Some(TransactionType::ReadWrite);
+                        }
+
+                        _ => (),
                     }
                 }
             }
@@ -264,8 +282,13 @@ impl QueryEngine {
     async fn transaction_error_check(
         &mut self,
         context: &mut QueryEngineContext<'_>,
+        route: &Route,
     ) -> Result<bool, Error> {
-        if context.in_error() && !context.rollback && context.client_request.executable() {
+        if context.in_error()
+            && !context.rollback
+            && context.client_request.executable()
+            && !route.rollback_savepoint()
+        {
             let bytes_sent = context
                 .stream
                 .error(

pgdog/src/frontend/router/parser/query/transaction.rs

@@ -16,6 +16,7 @@ impl QueryParser {
         context: &QueryParserContext,
     ) -> Result<Command, Error> {
         let extended = !context.query()?.simple();
+        let mut rollback_savepoint = false;
 
         if context.rw_conservative() && !context.read_only {
             self.write_override = true;
@@ -37,6 +38,7 @@ impl QueryParser {
                     extended,
                 });
             }
+            TransactionStmtKind::TransStmtRollbackTo => rollback_savepoint = true,
             TransactionStmtKind::TransStmtPrepare
             | TransactionStmtKind::TransStmtCommitPrepared
             | TransactionStmtKind::TransStmtRollbackPrepared => {
@@ -47,7 +49,9 @@ impl QueryParser {
             _ => (),
         }
 
-        Ok(Command::Query(Route::write(None)))
+        Ok(Command::Query(
+            Route::write(None).set_rollback_savepoint(rollback_savepoint),
+        ))
     }
 
     #[inline]

pgdog/src/frontend/router/parser/route.rs

@@ -60,6 +60,7 @@ pub struct Route {
     maintenance: bool,
     rewrite_plan: RewritePlan,
     rewritten_sql: Option<String>,
+    rollback_savepoint: bool,
 }
 
 impl Display for Route {
@@ -187,6 +188,15 @@ impl Route {
         self.read = read;
     }
 
+    pub fn set_rollback_savepoint(mut self, rollback: bool) -> Self {
+        self.rollback_savepoint = rollback;
+        self
+    }
+
+    pub fn rollback_savepoint(&self) -> bool {
+        self.rollback_savepoint
+    }
+
     pub fn set_write(mut self, write: FunctionBehavior) -> Self {
         self.set_write_mut(write);
         self