From 7d5252218a66e223bd9b52fa65143c6b15c434ac Mon Sep 17 00:00:00 2001
From: Fredrik Enestad <fredrik@enestad.com>
Date: Wed, 2 Apr 2025 14:46:03 +0200
Subject: [PATCH] expose sync rollback as a private api method

---
 tokio-postgres/src/client.rs              | 18 +++++++++++++++++-
 tokio-postgres/src/transaction.rs         | 18 ++----------------
 tokio-postgres/src/transaction_builder.rs | 13 ++-----------
 3 files changed, 21 insertions(+), 28 deletions(-)

diff --git a/tokio-postgres/src/client.rs b/tokio-postgres/src/client.rs
index b38bbba37..530995855 100644
--- a/tokio-postgres/src/client.rs
+++ b/tokio-postgres/src/client.rs
@@ -1,4 +1,4 @@
-use crate::codec::BackendMessages;
+use crate::codec::{BackendMessages, FrontendMessage};
 use crate::config::{SslMode, SslNegotiation};
 use crate::connection::{Request, RequestMessages};
 use crate::copy_out::CopyOutStream;
@@ -22,6 +22,7 @@ use futures_channel::mpsc;
 use futures_util::{future, pin_mut, ready, StreamExt, TryStreamExt};
 use parking_lot::Mutex;
 use postgres_protocol::message::backend::Message;
+use postgres_protocol::message::frontend;
 use postgres_types::BorrowToSql;
 use std::collections::HashMap;
 use std::fmt;
@@ -605,6 +606,21 @@ impl Client {
     pub fn __private_api_close(&mut self) {
         self.inner.sender.close_channel()
     }
+
+    #[doc(hidden)]
+    pub fn __private_api_rollback(&self, name: Option<&str>) {
+        let buf = self.inner().with_buf(|buf| {
+            if let Some(name) = name {
+                frontend::query(&format!("ROLLBACK TO {}", name), buf).unwrap();
+            } else {
+                frontend::query("ROLLBACK", buf).unwrap();
+            }
+            buf.split().freeze()
+        });
+        let _ = self
+            .inner()
+            .send(RequestMessages::Single(FrontendMessage::Raw(buf)));
+    }
 }
 
 impl fmt::Debug for Client {
diff --git a/tokio-postgres/src/transaction.rs b/tokio-postgres/src/transaction.rs
index 782c476c4..36a1d994f 100644
--- a/tokio-postgres/src/transaction.rs
+++ b/tokio-postgres/src/transaction.rs
@@ -1,5 +1,3 @@
-use crate::codec::FrontendMessage;
-use crate::connection::RequestMessages;
 use crate::copy_out::CopyOutStream;
 use crate::query::RowStream;
 #[cfg(feature = "runtime")]
@@ -14,7 +12,6 @@ use crate::{
 };
 use bytes::Buf;
 use futures_util::TryStreamExt;
-use postgres_protocol::message::frontend;
 use tokio::io::{AsyncRead, AsyncWrite};
 
 /// A representation of a PostgreSQL database transaction.
@@ -39,19 +36,8 @@ impl Drop for Transaction<'_> {
             return;
         }
 
-        let query = if let Some(sp) = self.savepoint.as_ref() {
-            format!("ROLLBACK TO {}", sp.name)
-        } else {
-            "ROLLBACK".to_string()
-        };
-        let buf = self.client.inner().with_buf(|buf| {
-            frontend::query(&query, buf).unwrap();
-            buf.split().freeze()
-        });
-        let _ = self
-            .client
-            .inner()
-            .send(RequestMessages::Single(FrontendMessage::Raw(buf)));
+        let name = self.savepoint.as_ref().map(|sp| sp.name.as_str());
+        self.client.__private_api_rollback(name);
     }
 }
 
diff --git a/tokio-postgres/src/transaction_builder.rs b/tokio-postgres/src/transaction_builder.rs
index 88c883176..8dec56afc 100644
--- a/tokio-postgres/src/transaction_builder.rs
+++ b/tokio-postgres/src/transaction_builder.rs
@@ -1,6 +1,4 @@
-use postgres_protocol::message::frontend;
-
-use crate::{codec::FrontendMessage, connection::RequestMessages, Client, Error, Transaction};
+use crate::{Client, Error, Transaction};
 
 /// The isolation level of a database transaction.
 #[derive(Debug, Copy, Clone)]
@@ -119,14 +117,7 @@ impl<'a> TransactionBuilder<'a> {
                     return;
                 }
 
-                let buf = self.client.inner().with_buf(|buf| {
-                    frontend::query("ROLLBACK", buf).unwrap();
-                    buf.split().freeze()
-                });
-                let _ = self
-                    .client
-                    .inner()
-                    .send(RequestMessages::Single(FrontendMessage::Raw(buf)));
+                self.client.__private_api_rollback(None);
             }
         }