Fix transction rollback on Future drop

- fixes drop of `Future` after synchronous message send but before
  `Transaction` object creation
- follow-on to
  f6189a9
  for `TransactionBuilder` and savepoint `Transaction::transaction`
  creation

Author: wquark

tokio-postgres/src/client.rs

@@ -13,8 +13,9 @@ use crate::types::{Oid, ToSql, Type};
 #[cfg(feature = "runtime")]
 use crate::Socket;
 use crate::{
-    copy_in, copy_out, prepare, query, simple_query, slice_iter, CancelToken, CopyInSink, Error,
-    Row, SimpleQueryMessage, Statement, ToStatement, Transaction, TransactionBuilder,
+    copy_in, copy_out, prepare, query, simple_query, slice_iter, transaction::Savepoint,
+    CancelToken, CopyInSink, Error, Row, SimpleQueryMessage, Statement, ToStatement, Transaction,
+    TransactionBuilder,
 };
 use bytes::{Buf, BytesMut};
 use fallible_iterator::FallibleIterator;
@@ -469,8 +470,17 @@ 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> {
+        self.build_transaction().start().await
+    }
+
+    pub(crate) async fn start_transaction_with_rollback(
+        &mut self,
+        query: &str,
+        savepoint: Option<Savepoint>,
+    ) -> Result<Transaction<'_>, Error> {
         struct RollbackIfNotDone<'me> {
             client: &'me Client,
+            savepoint: Option<&'me Savepoint>,
             done: bool,
         }
 
@@ -480,8 +490,13 @@ impl Client {
                     return;
                 }
 
+                let query = if let Some(sp) = self.savepoint {
+                    format!("ROLLBACK TO {}", sp.name)
+                } else {
+                    "ROLLBACK".to_string()
+                };
                 let buf = self.client.inner().with_buf(|buf| {
-                    frontend::query("ROLLBACK", buf).unwrap();
+                    frontend::query(&query, buf).unwrap();
                     buf.split().freeze()
                 });
                 let _ = self
@@ -499,13 +514,14 @@ impl Client {
         {
             let mut cleaner = RollbackIfNotDone {
                 client: self,
+                savepoint: savepoint.as_ref(),
                 done: false,
             };
-            self.batch_execute("BEGIN").await?;
+            self.batch_execute(query).await?;
             cleaner.done = true;
         }
 
-        Ok(Transaction::new(self))
+        Ok(Transaction::new(self, savepoint))
     }
 
     /// Returns a builder for a transaction with custom settings.

tokio-postgres/src/transaction.rs

@@ -28,8 +28,8 @@ pub struct Transaction<'a> {
 }
 
 /// A representation of a PostgreSQL database savepoint.
-struct Savepoint {
-    name: String,
+pub(crate) struct Savepoint {
+    pub(crate) name: String,
     depth: u32,
 }
 
@@ -56,10 +56,10 @@ impl<'a> Drop for Transaction<'a> {
 }
 
 impl<'a> Transaction<'a> {
-    pub(crate) fn new(client: &'a mut Client) -> Transaction<'a> {
+    pub(crate) fn new(client: &'a mut Client, savepoint: Option<Savepoint>) -> Transaction<'a> {
         Transaction {
             client,
-            savepoint: None,
+            savepoint,
             done: false,
         }
     }
@@ -298,13 +298,11 @@ impl<'a> Transaction<'a> {
         let depth = self.savepoint.as_ref().map_or(0, |sp| sp.depth) + 1;
         let name = name.unwrap_or_else(|| format!("sp_{}", depth));
         let query = format!("SAVEPOINT {}", name);
-        self.batch_execute(&query).await?;
+        let savepoint = Savepoint { name, depth };
 
-        Ok(Transaction {
-            client: self.client,
-            savepoint: Some(Savepoint { name, depth }),
-            done: false,
-        })
+        self.client
+            .start_transaction_with_rollback(&query, Some(savepoint))
+            .await
     }
 
     /// Returns a reference to the underlying `Client`.

tokio-postgres/src/transaction_builder.rs

@@ -106,8 +106,8 @@ impl<'a> TransactionBuilder<'a> {
             query.push_str(s);
         }
 
-        self.client.batch_execute(&query).await?;
-
-        Ok(Transaction::new(self.client))
+        self.client
+            .start_transaction_with_rollback(&query, None)
+            .await
     }
 }