Merge pull request rust-postgres#184 from nvanbenschoten/nvanbenschoten/savepoints

Add savepoint method to Transaction

Author: sfackler

src/transaction.rs

@@ -151,6 +151,7 @@ impl Config {
 pub struct Transaction<'conn> {
     conn: &'conn Connection,
     depth: u32,
+    savepoint_name: Option<String>,
     commit: Cell<bool>,
     finished: bool,
 }
@@ -177,6 +178,7 @@ impl<'conn> TransactionInternals<'conn> for Transaction<'conn> {
         Transaction {
             conn: conn,
             depth: depth,
+            savepoint_name: None,
             commit: Cell::new(false),
             finished: false,
         }
@@ -195,14 +197,13 @@ impl<'conn> Transaction<'conn> {
     fn finish_inner(&mut self) -> Result<()> {
         let mut conn = self.conn.conn.borrow_mut();
         debug_assert!(self.depth == conn.trans_depth);
-        let query = match (self.commit.get(), self.depth != 1) {
-            (false, true) => "ROLLBACK TO sp",
-            (false, false) => "ROLLBACK",
-            (true, true) => "RELEASE sp",
-            (true, false) => "COMMIT",
-        };
         conn.trans_depth -= 1;
-        conn.quick_query(query).map(|_| ())
+        match (self.commit.get(), &self.savepoint_name) {
+            (false, &Some(ref savepoint_name)) => conn.quick_query(&format!("ROLLBACK TO {}", savepoint_name)),
+            (false, &None) => conn.quick_query("ROLLBACK"),
+            (true, &Some(ref savepoint_name)) => conn.quick_query(&format!("RELEASE {}", savepoint_name)),
+            (true, &None) => conn.quick_query("COMMIT"),
+        }.map(|_| ())
     }
 
     /// Like `Connection::prepare`.
@@ -233,22 +234,33 @@ impl<'conn> Transaction<'conn> {
         self.conn.batch_execute(query)
     }
 
-    /// Like `Connection::transaction`.
+    /// Like `Connection::transaction`, but creates a nested transaction.
     ///
     /// # Panics
     ///
     /// Panics if there is an active nested transaction.
     pub fn transaction<'a>(&'a self) -> Result<Transaction<'a>> {
+        self.savepoint("sp")
+    }
+    
+    /// Like `Connection::transaction`, but creates a nested transaction
+    /// with the provided name.
+    ///
+    /// # Panics
+    ///
+    /// Panics if there is an active nested transaction.
+    pub fn savepoint<'a>(&'a self, name: &str) -> Result<Transaction<'a>> {
         let mut conn = self.conn.conn.borrow_mut();
         check_desync!(conn);
         assert!(conn.trans_depth == self.depth,
-                "`transaction` may only be called on the active transaction");
-        try!(conn.quick_query("SAVEPOINT sp"));
+                "`savepoint` may only be called on the active transaction");
+        try!(conn.quick_query(&format!("SAVEPOINT {}", name)));
         conn.trans_depth += 1;
         Ok(Transaction {
             conn: self.conn,
-            commit: Cell::new(false),
             depth: self.depth + 1,
+            savepoint_name: Some(name.to_owned()),
+            commit: Cell::new(false),
             finished: false,
         })
     }

tests/test.rs

@@ -202,9 +202,9 @@ fn test_nested_transactions() {
             }
 
             {
-                let trans3 = or_panic!(trans2.transaction());
-                or_panic!(trans3.execute("INSERT INTO foo (id) VALUES (6)", &[]));
-                assert!(trans3.commit().is_ok());
+                let sp = or_panic!(trans2.savepoint("custom"));
+                or_panic!(sp.execute("INSERT INTO foo (id) VALUES (6)", &[]));
+                assert!(sp.commit().is_ok());
             }
 
             assert!(trans2.commit().is_ok());
@@ -250,10 +250,10 @@ fn test_nested_transactions_finish() {
             }
 
             {
-                let trans3 = or_panic!(trans2.transaction());
-                or_panic!(trans3.execute("INSERT INTO foo (id) VALUES (6)", &[]));
-                trans3.set_commit();
-                assert!(trans3.finish().is_ok());
+                let sp = or_panic!(trans2.savepoint("custom"));
+                or_panic!(sp.execute("INSERT INTO foo (id) VALUES (6)", &[]));
+                sp.set_commit();
+                assert!(sp.finish().is_ok());
             }
 
             trans2.set_commit();
@@ -294,6 +294,15 @@ fn test_trans_with_nested_trans() {
     trans.transaction().unwrap();
 }
 
+#[test]
+#[should_panic(expected = "active transaction")]
+fn test_trans_with_savepoints() {
+    let conn = or_panic!(Connection::connect("postgres://postgres@localhost", SslMode::None));
+    let trans = or_panic!(conn.transaction());
+    let _sp = or_panic!(trans.savepoint("custom"));
+    trans.savepoint("custom2").unwrap();
+}
+
 #[test]
 fn test_stmt_execute_after_transaction() {
     let conn = or_panic!(Connection::connect("postgres://postgres@localhost", SslMode::None));