diff --git a/src/transaction.rs b/src/transaction.rs index d606b06f..b12b1662 100644 --- a/src/transaction.rs +++ b/src/transaction.rs @@ -151,6 +151,7 @@ impl Config { pub struct Transaction<'conn> { conn: &'conn Connection, depth: u32, + savepoint_name: Option, commit: Cell, 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> { + 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> { 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, }) } diff --git a/tests/test.rs b/tests/test.rs index 8aee2788..bbf49e0c 100644 --- a/tests/test.rs +++ b/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));