Repository: sfackler/rust-postgres (language: rust)

Issue:
`TransactionBuilder::start` Future dropped before completion
`batch_execute` submits the creation of the transaction (https://github.com/sfackler/rust-postgres/blob/master/tokio-postgres/src/transaction_builder.rs#L109), but the `Transaction` object itself may never be constructed in case the future is dropped.

https://github.com/sfackler/rust-postgres/commit/f6189a95f24af45068ecfd6e3c4e6b71ac8e43fe fixed it for `Client::transaction`.


Discussion / hints:
Good find. We should probably implement `transaction` with `TransactionBuilder` to avoid having two places that need to deal with this.
Makes sense. Anything else you need from my end? Is this something you'd knock out quickly or would you rather have a PR?
Either way!

Fix patch:
diff --git a/tokio-postgres/src/client.rs b/tokio-postgres/src/client.rs
index b04f05f88..92eabde36 100644
--- a/tokio-postgres/src/client.rs
+++ b/tokio-postgres/src/client.rs
@@ -1,4 +1,4 @@
-use crate::codec::{BackendMessages, FrontendMessage};
+use crate::codec::BackendMessages;
 use crate::config::SslMode;
 use crate::connection::{Request, RequestMessages};
 use crate::copy_out::CopyOutStream;
@@ -21,7 +21,7 @@ use fallible_iterator::FallibleIterator;
 use futures_channel::mpsc;
 use futures_util::{future, pin_mut, ready, StreamExt, TryStreamExt};
 use parking_lot::Mutex;
-use postgres_protocol::message::{backend::Message, frontend};
+use postgres_protocol::message::backend::Message;
 use postgres_types::BorrowToSql;
 use std::collections::HashMap;
 use std::fmt;
@@ -532,43 +532,7 @@ impl Client {
     ///
     /// The transaction will roll back by default - use the `commit` method to commit it.
     pub async fn transaction(&mut self) -> Result<Transaction<'_>, Error> {
-        struct RollbackIfNotDone<'me> {
-            client: &'me Client,
-            done: bool,
-        }
-
-        impl<'a> Drop for RollbackIfNotDone<'a> {
-            fn drop(&mut self) {
-                if self.done {
-                    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)));
-            }
-        }
-
-        // This is done, as `Future` created by this method can be dropped after
-        // `RequestMessages` is synchronously send to the `Connection` by
-        // `batch_execute()`, but before `Responses` is asynchronously polled to
-        // completion. In that case `Transaction` won't be created and thus
-        // won't be rolled back.
-        {
-            let mut cleaner = RollbackIfNotDone {
-                client: self,
-                done: false,
-            };
-            self.batch_execute("BEGIN").await?;
-            cleaner.done = true;
-        }
-
-        Ok(Transaction::new(self))
+        self.build_transaction().start().await
     }
 
     /// Returns a builder for a transaction with custom settings.
diff --git a/tokio-postgres/src/transaction_builder.rs b/tokio-postgres/src/transaction_builder.rs
index 9718ac588..93e9e9801 100644
--- a/tokio-postgres/src/transaction_builder.rs
+++ b/tokio-postgres/src/transaction_builder.rs
@@ -1,4 +1,6 @@
-use crate::{Client, Error, Transaction};
+use postgres_protocol::message::frontend;
+
+use crate::{codec::FrontendMessage, connection::RequestMessages, Client, Error, Transaction};
 
 /// The isolation level of a database transaction.
 #[derive(Debug, Copy, Clone)]
@@ -106,7 +108,41 @@ impl<'a> TransactionBuilder<'a> {
             query.push_str(s);
         }
 
-        self.client.batch_execute(&query).await?;
+        struct RollbackIfNotDone<'me> {
+            client: &'me Client,
+            done: bool,
+        }
+
+        impl<'a> Drop for RollbackIfNotDone<'a> {
+            fn drop(&mut self) {
+                if self.done {
+                    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)));
+            }
+        }
+
+        // This is done as `Future` created by this method can be dropped after
+        // `RequestMessages` is synchronously send to the `Connection` by
+        // `batch_execute()`, but before `Responses` is asynchronously polled to
+        // completion. In that case `Transaction` won't be created and thus
+        // won't be rolled back.
+        {
+            let mut cleaner = RollbackIfNotDone {
+                client: self.client,
+                done: false,
+            };
+            self.client.batch_execute(&query).await?;
+            cleaner.done = true;
+        }
 
         Ok(Transaction::new(self.client))
     }


Test patch: