From f3777bed7670bfbefdc29eda5ce3a6ffd19d8352 Mon Sep 17 00:00:00 2001
From: Darren Tsung <darren@onesignal.com>
Date: Tue, 6 Nov 2018 10:14:32 -0800
Subject: [PATCH] Add a TransactionBuilder struct

The TransactionBuilder is useful when trying to create a Transaction
around a future which takes ownership of the client. See doc comment
for more details.
---
 tokio-postgres/src/lib.rs    | 25 +++++++++++++++++-
 tokio-postgres/tests/test.rs | 49 ++++++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/tokio-postgres/src/lib.rs b/tokio-postgres/src/lib.rs
index a16dd15e..0d918664 100644
--- a/tokio-postgres/src/lib.rs
+++ b/tokio-postgres/src/lib.rs
@@ -116,7 +116,17 @@ impl Client {
         // FIXME error type?
         T::Error: From<Error>,
     {
-        Transaction(proto::TransactionFuture::new(self.0.clone(), future))
+        self.transaction_builder().build(future)
+    }
+
+    /// Creates a TransactionBuilder, which can later be used to create
+    /// a Transaction around a future.
+    ///
+    /// Use this when Client is moved into the future being built.
+    /// For example, when executing multiple statements that depend
+    /// on the previous statement's result.
+    pub fn transaction_builder(&self) -> TransactionBuilder {
+        TransactionBuilder(self.0.clone())
     }
 
     pub fn batch_execute(&mut self, query: &str) -> BatchExecute {
@@ -333,6 +343,19 @@ impl Row {
     }
 }
 
+pub struct TransactionBuilder(proto::Client);
+
+impl TransactionBuilder {
+    pub fn build<T>(self, future: T) -> Transaction<T>
+    where
+        T: Future,
+        // FIXME error type?
+        T::Error: From<Error>,
+    {
+        Transaction(proto::TransactionFuture::new(self.0, future))
+    }
+}
+
 #[must_use = "futures do nothing unless polled"]
 pub struct Transaction<T>(proto::TransactionFuture<T, T::Item, T::Error>)
 where
diff --git a/tokio-postgres/tests/test.rs b/tokio-postgres/tests/test.rs
index 015a6b8e..6ef5111a 100644
--- a/tokio-postgres/tests/test.rs
+++ b/tokio-postgres/tests/test.rs
@@ -692,3 +692,52 @@ fn copy_out() {
         ).unwrap();
     assert_eq!(&data[..], b"1\tjim\n2\tjoe\n");
 }
+
+#[test]
+fn transaction_builder_around_moved_client() {
+    let _ = env_logger::try_init();
+    let mut runtime = Runtime::new().unwrap();
+
+    let (mut client, connection) = runtime
+        .block_on(tokio_postgres::connect(
+            "postgres://postgres@localhost:5433".parse().unwrap(),
+            TlsMode::None,
+        )).unwrap();
+    let connection = connection.map_err(|e| panic!("{}", e));
+    runtime.handle().spawn(connection).unwrap();
+
+    let transaction_builder = client.transaction_builder();
+    let work = future::lazy(move || {
+        let execute =
+            client.batch_execute(
+                "CREATE TEMPORARY TABLE transaction_foo (
+                    id SERIAL,
+                    name TEXT
+                )");
+
+        execute.and_then(move |_| {
+            client
+                .prepare("INSERT INTO transaction_foo (name) VALUES ($1), ($2)")
+                .map(|statement| (client, statement))
+        })
+    }).and_then(|(mut client, statement)| {
+        client
+            .query(&statement, &[&"jim", &"joe"])
+            .collect()
+            .map(|_res| client)
+    });
+
+    let transaction = transaction_builder.build(work);
+    let mut client = runtime.block_on(transaction).unwrap();
+
+    let data = runtime
+        .block_on(
+            client
+                .prepare("COPY transaction_foo TO STDOUT")
+                .and_then(|s| client.copy_out(&s, &[]).concat2()),
+        ).unwrap();
+    assert_eq!(&data[..], b"1\tjim\n2\tjoe\n");
+
+    drop(client);
+    runtime.run().unwrap();
+}