Basic transaction support

Author: sfackler

tokio-postgres/src/lib.rs

@@ -95,6 +95,14 @@ impl Client {
         Query(self.0.query(&statement.0, params))
     }
 
+    pub fn transaction<T>(&mut self, future: T) -> Transaction<T>
+    where
+        T: Future,
+        T::Error: From<Error>,
+    {
+        Transaction(proto::TransactionFuture::new(self.0.clone(), future))
+    }
+
     pub fn batch_execute(&mut self, query: &str) -> BatchExecute {
         BatchExecute(self.0.batch_execute(query))
     }
@@ -242,6 +250,25 @@ impl Row {
     }
 }
 
+#[must_use = "futures do nothing unless polled"]
+pub struct Transaction<T>(proto::TransactionFuture<T, T::Item, T::Error>)
+where
+    T: Future,
+    T::Error: From<Error>;
+
+impl<T> Future for Transaction<T>
+where
+    T: Future,
+    T::Error: From<Error>,
+{
+    type Item = T::Item;
+    type Error = T::Error;
+
+    fn poll(&mut self) -> Poll<T::Item, T::Error> {
+        self.0.poll()
+    }
+}
+
 #[must_use = "futures do nothing unless polled"]
 pub struct BatchExecute(proto::SimpleQueryFuture);
 

tokio-postgres/src/proto/mod.rs

@@ -21,6 +21,7 @@ mod row;
 mod simple_query;
 mod socket;
 mod statement;
+mod transaction;
 mod typeinfo;
 mod typeinfo_composite;
 mod typeinfo_enum;
@@ -37,3 +38,4 @@ pub use proto::row::Row;
 pub use proto::simple_query::SimpleQueryFuture;
 pub use proto::socket::Socket;
 pub use proto::statement::Statement;
+pub use proto::transaction::TransactionFuture;

tokio-postgres/src/proto/transaction.rs

@@ -0,0 +1,104 @@
+use futures::{Async, Future, Poll};
+use proto::client::Client;
+use proto::simple_query::SimpleQueryFuture;
+use state_machine_future::RentToOwn;
+
+use Error;
+
+#[derive(StateMachineFuture)]
+pub enum Transaction<F, T, E>
+where
+    F: Future<Item = T, Error = E>,
+    E: From<Error>,
+{
+    #[state_machine_future(start, transitions(Beginning))]
+    Start { client: Client, future: F },
+    #[state_machine_future(transitions(Running))]
+    Beginning {
+        client: Client,
+        begin: SimpleQueryFuture,
+        future: F,
+    },
+    #[state_machine_future(transitions(Finishing))]
+    Running { client: Client, future: F },
+    #[state_machine_future(transitions(Finished))]
+    Finishing {
+        future: SimpleQueryFuture,
+        result: Result<T, E>,
+    },
+    #[state_machine_future(ready)]
+    Finished(T),
+    #[state_machine_future(error)]
+    Failed(E),
+}
+
+impl<F, T, E> PollTransaction<F, T, E> for Transaction<F, T, E>
+where
+    F: Future<Item = T, Error = E>,
+    E: From<Error>,
+{
+    fn poll_start<'a>(
+        state: &'a mut RentToOwn<'a, Start<F, T, E>>,
+    ) -> Poll<AfterStart<F, T, E>, E> {
+        let state = state.take();
+        transition!(Beginning {
+            begin: state.client.batch_execute("BEGIN"),
+            client: state.client,
+            future: state.future,
+        })
+    }
+
+    fn poll_beginning<'a>(
+        state: &'a mut RentToOwn<'a, Beginning<F, T, E>>,
+    ) -> Poll<AfterBeginning<F, T, E>, E> {
+        try_ready!(state.begin.poll());
+        let state = state.take();
+        transition!(Running {
+            client: state.client,
+            future: state.future,
+        })
+    }
+
+    fn poll_running<'a>(
+        state: &'a mut RentToOwn<'a, Running<F, T, E>>,
+    ) -> Poll<AfterRunning<T, E>, E> {
+        match state.future.poll() {
+            Ok(Async::NotReady) => return Ok(Async::NotReady),
+            Ok(Async::Ready(t)) => transition!(Finishing {
+                future: state.client.batch_execute("COMMIT"),
+                result: Ok(t),
+            }),
+            Err(e) => transition!(Finishing {
+                future: state.client.batch_execute("ROLLBACK"),
+                result: Err(e),
+            }),
+        }
+    }
+
+    fn poll_finishing<'a>(
+        state: &'a mut RentToOwn<'a, Finishing<T, E>>,
+    ) -> Poll<AfterFinishing<T>, E> {
+        match state.future.poll() {
+            Ok(Async::NotReady) => return Ok(Async::NotReady),
+            Ok(Async::Ready(())) => {
+                let t = state.take().result?;
+                transition!(Finished(t))
+            }
+            Err(e) => match state.take().result {
+                Ok(_) => Err(e.into()),
+                // prioritize the future's error over the rollback error
+                Err(e) => Err(e),
+            },
+        }
+    }
+}
+
+impl<F, T, E> TransactionFuture<F, T, E>
+where
+    F: Future<Item = T, Error = E>,
+    E: From<Error>,
+{
+    pub fn new(client: Client, future: F) -> TransactionFuture<F, T, E> {
+        Transaction::start(client, future)
+    }
+}

tokio-postgres/tests/test.rs

@@ -9,6 +9,7 @@ extern crate log;
 
 use futures::future;
 use futures::sync::mpsc;
+use std::error::Error;
 use std::time::{Duration, Instant};
 use tokio::prelude::*;
 use tokio::runtime::current_thread::Runtime;
@@ -477,3 +478,81 @@ fn notifications() {
     assert_eq!(notifications[1].channel, "test_notifications");
     assert_eq!(notifications[1].payload, "world");
 }
+
+#[test]
+fn test_transaction_commit() {
+    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();
+
+    runtime
+        .block_on(client.batch_execute(
+            "CREATE TEMPORARY TABLE foo (
+                id SERIAL,
+                name TEXT
+            )",
+        ))
+        .unwrap();
+
+    let f = client.batch_execute("INSERT INTO foo (name) VALUES ('steven')");
+    runtime.block_on(client.transaction(f)).unwrap();
+
+    let rows = runtime
+        .block_on(
+            client
+                .prepare("SELECT name FROM foo")
+                .and_then(|s| client.query(&s, &[]).collect()),
+        )
+        .unwrap();
+
+    assert_eq!(rows.len(), 1);
+    assert_eq!(rows[0].get::<_, &str>(0), "steven");
+}
+
+#[test]
+fn test_transaction_abort() {
+    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();
+
+    runtime
+        .block_on(client.batch_execute(
+            "CREATE TEMPORARY TABLE foo (
+                id SERIAL,
+                name TEXT
+            )",
+        ))
+        .unwrap();
+
+    let f = client
+        .batch_execute("INSERT INTO foo (name) VALUES ('steven')")
+        .map_err(|e| Box::new(e) as Box<Error>)
+        .and_then(|_| Err::<(), _>(Box::<Error>::from("")));
+    runtime.block_on(client.transaction(f)).unwrap_err();
+
+    let rows = runtime
+        .block_on(
+            client
+                .prepare("SELECT name FROM foo")
+                .and_then(|s| client.query(&s, &[]).collect()),
+        )
+        .unwrap();
+
+    assert_eq!(rows.len(), 0);
+}