From afd7d8800d62b3720703bc4a656391e15e048d21 Mon Sep 17 00:00:00 2001
From: Steven Fackler <sfackler@gmail.com>
Date: Tue, 27 Aug 2013 01:40:23 -0400
Subject: [PATCH] Better transaction syntax

Closes #9
---
 src/lib.rs  | 45 ++++++++++++++++++++++++++++++++++++++++-----
 src/test.rs | 20 ++++++++++----------
 2 files changed, 50 insertions(+), 15 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 7c53fdf5..99259c8a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -202,15 +202,18 @@ impl PostgresConnection {
         })
     }
 
-    pub fn in_transaction<T, E: ToStr>(&self, blk: &fn(&PostgresConnection)
-                                                       -> Result<T, E>)
-                                       -> Result<T, E> {
+    pub fn in_transaction<T>(&self, blk: &fn(&PostgresTransaction) -> T)
+                             -> T {
         self.quick_query("BEGIN");
 
+        let trans = PostgresTransaction {
+            conn: self,
+            commit: Cell::new(true)
+        };
         // If this fails, Postgres will rollback when the connection closes
-        let ret = blk(self);
+        let ret = blk(&trans);
 
-        if ret.is_ok() {
+        if trans.commit.take() {
             self.quick_query("COMMIT");
         } else {
             self.quick_query("ROLLBACK");
@@ -241,6 +244,38 @@ impl PostgresConnection {
     }
 }
 
+pub struct PostgresTransaction<'self> {
+    priv conn: &'self PostgresConnection,
+    priv commit: Cell<bool>
+}
+
+impl<'self> PostgresTransaction<'self> {
+    pub fn prepare<'a>(&'a self, query: &str) -> PostgresStatement<'a> {
+        self.conn.prepare(query)
+    }
+
+    pub fn try_prepare<'a>(&'a self, query: &str)
+                           -> Result<PostgresStatement<'a>, PostgresDbError> {
+        self.conn.try_prepare(query)
+    }
+
+    pub fn will_commit(&self) -> bool {
+        let commit = self.commit.take();
+        self.commit.put_back(commit);
+        commit
+    }
+
+    pub fn set_commit(&self) {
+        self.commit.take();
+        self.commit.put_back(true);
+    }
+
+    pub fn set_rollback(&self) {
+        self.commit.take();
+        self.commit.put_back(false);
+    }
+}
+
 pub struct PostgresStatement<'self> {
     priv conn: &'self PostgresConnection,
     priv name: ~str,
diff --git a/src/test.rs b/src/test.rs
index caba369f..6d861c03 100644
--- a/src/test.rs
+++ b/src/test.rs
@@ -6,10 +6,10 @@ use postgres::{PostgresConnection, ToSql};
 fn test_basic() {
     let conn = PostgresConnection::connect("postgres://postgres@127.0.0.1:5432");
 
-    do conn.in_transaction |conn| {
-        conn.prepare("CREATE TABLE foo (id BIGINT PRIMARY KEY)").update([]);
+    do conn.in_transaction |trans| {
+        trans.prepare("CREATE TABLE foo (id BIGINT PRIMARY KEY)").update([]);
 
-        Err::<(), ()>(())
+        trans.set_rollback();
     };
 }
 
@@ -23,16 +23,16 @@ fn test_prepare_err() {
 fn test_query() {
     let conn = PostgresConnection::connect("postgres://postgres@127.0.0.1:5432");
 
-    do conn.in_transaction |conn| {
-        conn.prepare("CREATE TABLE foo (id BIGINT PRIMARY KEY)").update([]);
-        conn.prepare("INSERT INTO foo (id) VALUES ($1), ($2)")
+    do conn.in_transaction |trans| {
+        trans.prepare("CREATE TABLE foo (id BIGINT PRIMARY KEY)").update([]);
+        trans.prepare("INSERT INTO foo (id) VALUES ($1), ($2)")
                 .update([&1 as &ToSql, &2 as &ToSql]);
-        let stmt = conn.prepare("SELECT * from foo ORDER BY id");
+        let stmt = trans.prepare("SELECT * from foo ORDER BY id");
         let result = stmt.query([]);
 
         assert_eq!(~[1, 2], result.iter().map(|row| { row[0] }).collect());
 
-        Err::<(), ()>(())
+        trans.set_rollback();
     };
 }
 
@@ -40,7 +40,7 @@ fn test_query() {
 fn test_nulls() {
     let conn = PostgresConnection::connect("postgres://postgres@127.0.0.1:5432");
 
-    do conn.in_transaction |conn| {
+    do conn.in_transaction |trans| {
         conn.prepare("CREATE TABLE foo (
                         id BIGINT PRIMARY KEY,
                         val VARCHAR
@@ -54,7 +54,7 @@ fn test_nulls() {
         assert_eq!(~[Some(~"foobar"), None],
                    result.iter().map(|row| { row[1] }).collect());
 
-        Err::<(), ()>(())
+        trans.set_rollback();
     };
 }