From cf5bd66218fa68b480fb34771b34468e1cfcd867 Mon Sep 17 00:00:00 2001
From: Steven Fackler <sfackler@gmail.com>
Date: Wed, 4 Sep 2013 23:28:44 -0700
Subject: [PATCH] Add nested transaction support

---
 README.md   |  3 ++-
 src/lib.rs  | 26 ++++++++++++++++++++++++++
 src/test.rs | 42 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index f8e54396..8dcc071f 100644
--- a/README.md
+++ b/README.md
@@ -138,7 +138,8 @@ do conn.in_transaction |trans| {
     }
 }
 ```
-A transaction will commit by default. Transactions cannot currently be nested.
+A transaction will commit by default. Nested transactions are supported via
+savepoints.
 
 Lazy Queries
 ------------
diff --git a/src/lib.rs b/src/lib.rs
index dcfcda4e..73f580c1 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -328,6 +328,7 @@ impl PostgresConnection {
 
         let trans = PostgresTransaction {
             conn: self,
+            next_savepoint_id: Cell::new(0),
             commit: Cell::new(true)
         };
         // If this fails, Postgres will rollback when the connection closes
@@ -378,6 +379,7 @@ impl PostgresConnection {
 
 pub struct PostgresTransaction<'self> {
     priv conn: &'self PostgresConnection,
+    priv next_savepoint_id: Cell<uint>,
     priv commit: Cell<bool>
 }
 
@@ -403,6 +405,30 @@ impl<'self> PostgresTransaction<'self> {
         self.conn.try_update(query, params)
     }
 
+    pub fn in_transaction<T>(&self, blk: &fn(&PostgresTransaction) -> T) -> T {
+        let id = self.next_savepoint_id.take();
+        let savepoint = fmt!("savepoint_%u", id);
+        self.next_savepoint_id.put_back(id + 1);
+
+        self.conn.quick_query(fmt!("SAVEPOINT %s", savepoint));
+
+        let nested_trans = PostgresTransaction {
+            conn: self.conn,
+            next_savepoint_id: Cell::new(id + 1),
+            commit: Cell::new(true)
+        };
+
+        let ret = blk(&nested_trans);
+
+        if nested_trans.commit.take() {
+            self.conn.quick_query(fmt!("RELEASE %s", savepoint));
+        } else {
+            self.conn.quick_query(fmt!("ROLLBACK TO %s", savepoint));
+        }
+
+        ret
+    }
+
     pub fn will_commit(&self) -> bool {
         let commit = self.commit.take();
         self.commit.put_back(commit);
diff --git a/src/test.rs b/src/test.rs
index a98ae7e5..beb1b4de 100644
--- a/src/test.rs
+++ b/src/test.rs
@@ -56,6 +56,48 @@ fn test_transaction_rollback() {
     assert_eq!(~[1i32], result.map(|row| { row[0] }).collect());
 }
 
+#[test]
+fn test_nested_transactions() {
+    let conn = PostgresConnection::connect("postgres://postgres@127.0.0.1:5432");
+    conn.update("CREATE TEMPORARY TABLE foo (id INT PRIMARY KEY)", []);
+
+    conn.update("INSERT INTO foo (id) VALUES (1)", []);
+
+    do conn.in_transaction |trans1| {
+        trans1.update("INSERT INTO foo (id) VALUES (2)", []);
+
+        do trans1.in_transaction |trans2| {
+            trans2.update("INSERT INTO foo (id) VALUES (3)", []);
+            trans2.set_rollback();
+        }
+
+        do trans1.in_transaction |trans2| {
+            trans2.update("INSERT INTO foo (id) VALUES (4)", []);
+
+            do trans2.in_transaction |trans3| {
+                trans3.update("INSERT INTO foo (id) VALUES (5)", []);
+                trans3.set_rollback();
+            }
+
+            do trans2.in_transaction |trans3| {
+                trans3.update("INSERT INTO foo (id) VALUES (6)", []);
+            }
+        }
+
+        let stmt = conn.prepare("SELECT * FROM foo ORDER BY id");
+        let result = stmt.query([]);
+
+        assert_eq!(~[1i32, 2, 4, 6], result.map(|row| { row[0] }).collect());
+
+        trans1.set_rollback();
+    }
+
+    let stmt = conn.prepare("SELECT * FROM foo ORDER BY id");
+    let result = stmt.query([]);
+
+    assert_eq!(~[1i32], result.map(|row| { row[0] }).collect());
+}
+
 #[test]
 fn test_query() {
     let conn = PostgresConnection::connect("postgres://postgres@127.0.0.1:5432");