transaction support

Author: sfackler

postgres-tokio/src/error.rs

@@ -8,10 +8,10 @@ use Connection;
 pub use postgres_shared::error::*;
 
 #[derive(Debug)]
-pub enum Error {
+pub enum Error<C = Connection> {
     Io(io::Error),
-    Db(Box<DbError>, Connection),
-    Conversion(Box<error::Error + Sync + Send>, Connection),
+    Db(Box<DbError>, C),
+    Conversion(Box<error::Error + Sync + Send>, C),
 }
 
 impl fmt::Display for Error {

postgres-tokio/src/lib.rs

@@ -726,9 +726,44 @@ impl Row {
     }
 }
 
+#[derive(Debug)]
 pub struct Transaction(Connection);
 
 impl Transaction {
+    pub fn batch_execute(self, query: &str) -> BoxFuture<Transaction, Error<Transaction>> {
+        self.0.batch_execute(query)
+            .map(Transaction)
+            .map_err(transaction_err)
+            .boxed()
+    }
+
+    pub fn prepare(self, query: &str) -> BoxFuture<(Statement, Transaction), Error<Transaction>> {
+        self.0.prepare(query)
+            .map(|(s, c)| (s, Transaction(c)))
+            .map_err(transaction_err)
+            .boxed()
+    }
+
+    pub fn execute(self,
+                   statement: &Statement,
+                   params: &[&ToSql])
+                   -> BoxFuture<(u64, Transaction), Error<Transaction>> {
+        self.0.execute(statement, params)
+            .map(|(n, c)| (n, Transaction(c)))
+            .map_err(transaction_err)
+            .boxed()
+    }
+
+    pub fn query(self,
+                 statement: &Statement,
+                 params: &[&ToSql])
+                 -> BoxStateStream<Row, Transaction, Error<Transaction>> {
+        self.0.query(statement, params)
+            .map_state(Transaction)
+            .map_err(transaction_err)
+            .boxed()
+    }
+
     pub fn commit(self) -> BoxFuture<Connection, Error> {
         self.finish("COMMIT")
     }
@@ -756,3 +791,11 @@ fn bad_message<T>() -> T
 {
     io::Error::new(io::ErrorKind::InvalidInput, "unexpected message").into()
 }
+
+fn transaction_err(e: Error) -> Error<Transaction> {
+    match e {
+        Error::Io(e) => Error::Io(e),
+        Error::Db(e, c) => Error::Db(e, Transaction(c)),
+        Error::Conversion(e, c) => Error::Conversion(e, Transaction(c))
+    }
+}

postgres-tokio/src/test.rs

@@ -134,11 +134,33 @@ fn query() {
         .and_then(|c| c.prepare("SELECT id, name FROM foo ORDER BY id"))
         .and_then(|(s, c)| c.query(&s, &[]).collect())
         .and_then(|(r, c)| {
+            assert_eq!(r[0].get::<i32, _>("id"), 1);
             assert_eq!(r[0].get::<String, _>("name"), "joe");
+            assert_eq!(r[1].get::<i32, _>("id"), 2);
             assert_eq!(r[1].get::<String, _>("name"), "bob");
             c.prepare("")
         })
         .and_then(|(s, c)| c.query(&s, &[]).collect())
         .map(|(r, _)| assert!(r.is_empty()));
     l.run(done).unwrap();
 }
+
+#[test]
+fn transaction() {
+    let mut l = Core::new().unwrap();
+    let done = Connection::connect("postgres://postgres@localhost", &l.handle())
+        .then(|c| c.unwrap().batch_execute("CREATE TEMPORARY TABLE foo (id SERIAL, name VARCHAR);"))
+        .then(|c| c.unwrap().transaction())
+        .then(|t| t.unwrap().batch_execute("INSERT INTO foo (name) VALUES ('joe');"))
+        .then(|t| t.unwrap().rollback())
+        .then(|c| c.unwrap().transaction())
+        .then(|t| t.unwrap().batch_execute("INSERT INTO foo (name) VALUES ('bob');"))
+        .then(|t| t.unwrap().commit())
+        .then(|c| c.unwrap().prepare("SELECT name FROM foo"))
+        .and_then(|(s, c)| c.query(&s, &[]).collect())
+        .map(|(r, _)| {
+            assert_eq!(r.len(), 1);
+            assert_eq!(r[0].get::<String, _>("name"), "bob");
+        });
+    l.run(done).unwrap();
+}