diff --git a/postgres/src/test.rs b/postgres/src/test.rs index db7cc1d3..a55ce8e7 100644 --- a/postgres/src/test.rs +++ b/postgres/src/test.rs @@ -93,3 +93,54 @@ fn transaction_drop() { let rows = client.query("SELECT * FROM foo", &[]).unwrap(); assert_eq!(rows.len(), 0); } + +#[test] +fn nested_transactions() { + let mut client = Client::connect("host=localhost port=5433 user=postgres", NoTls).unwrap(); + + client + .batch_execute("CREATE TEMPORARY TABLE foo (id INT PRIMARY KEY)") + .unwrap(); + + let mut transaction = client.transaction().unwrap(); + + transaction + .execute("INSERT INTO foo (id) VALUES (1)", &[]) + .unwrap(); + + let mut transaction2 = transaction.transaction().unwrap(); + + transaction2 + .execute("INSERT INTO foo (id) VALUES (2)", &[]) + .unwrap(); + + transaction2.rollback().unwrap(); + + let rows = transaction + .query("SELECT id FROM foo ORDER BY id", &[]) + .unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].get::<_, i32>(0), 1); + + let mut transaction3 = transaction.transaction().unwrap(); + + transaction3 + .execute("INSERT INTO foo (id) VALUES(3)", &[]) + .unwrap(); + + let mut transaction4 = transaction3.transaction().unwrap(); + + transaction4 + .execute("INSERT INTO foo (id) VALUES(4)", &[]) + .unwrap(); + + transaction4.commit().unwrap(); + transaction3.commit().unwrap(); + transaction.commit().unwrap(); + + let rows = client.query("SELECT id FROM foo ORDER BY id", &[]).unwrap(); + assert_eq!(rows.len(), 3); + assert_eq!(rows[0].get::<_, i32>(0), 1); + assert_eq!(rows[1].get::<_, i32>(0), 3); + assert_eq!(rows[2].get::<_, i32>(0), 4); +} diff --git a/postgres/src/transaction.rs b/postgres/src/transaction.rs index b61b5d58..0a252177 100644 --- a/postgres/src/transaction.rs +++ b/postgres/src/transaction.rs @@ -5,6 +5,7 @@ use crate::{Client, Query, Statement}; pub struct Transaction<'a> { client: &'a mut Client, + depth: u32, done: bool, } @@ -20,13 +21,19 @@ impl<'a> Transaction<'a> { pub(crate) fn new(client: &'a mut Client) -> Transaction<'a> { Transaction { client, + depth: 0, done: false, } } pub fn commit(mut self) -> Result<(), Error> { self.done = true; - self.client.batch_execute("COMMIT") + if self.depth == 0 { + self.client.batch_execute("COMMIT") + } else { + self.client + .batch_execute(&format!("RELEASE sp{}", self.depth)) + } } pub fn rollback(mut self) -> Result<(), Error> { @@ -35,7 +42,12 @@ impl<'a> Transaction<'a> { } fn rollback_inner(&mut self) -> Result<(), Error> { - self.client.batch_execute("ROLLBACK") + if self.depth == 0 { + self.client.batch_execute("ROLLBACK") + } else { + self.client + .batch_execute(&format!("ROLLBACK TO sp{}", self.depth)) + } } pub fn prepare(&mut self, query: &str) -> Result { @@ -63,4 +75,15 @@ impl<'a> Transaction<'a> { pub fn batch_execute(&mut self, query: &str) -> Result<(), Error> { self.client.batch_execute(query) } + + pub fn transaction(&mut self) -> Result, Error> { + let depth = self.depth + 1; + self.client + .batch_execute(&format!("SAVEPOINT sp{}", depth))?; + Ok(Transaction { + client: self.client, + depth, + done: false, + }) + } }