From 2ac7379a58a6b5ebb7cd8acbcc55886c71bc11e2 Mon Sep 17 00:00:00 2001
From: Steven Fackler <sfackler@gmail.com>
Date: Fri, 26 Jul 2013 01:34:02 -0400
Subject: [PATCH] A bit of transaction support

---
 src/sqlite3/lib.rs  | 49 +++++++++++++++++++++++++++++++++------------
 src/sqlite3/test.rs | 46 ++++++++++++++++++++++++++++++++++++------
 2 files changed, 76 insertions(+), 19 deletions(-)

diff --git a/src/sqlite3/lib.rs b/src/sqlite3/lib.rs
index 4910273c..fa9d5d70 100644
--- a/src/sqlite3/lib.rs
+++ b/src/sqlite3/lib.rs
@@ -1,5 +1,3 @@
-extern mod sql;
-
 use std::libc::c_int;
 use std::ptr;
 use std::str;
@@ -62,6 +60,24 @@ impl Connection {
     }
 }
 
+macro_rules! ret_err(
+    ($mat:expr { $($p:pat => $blk:expr),+ }) => (
+        match $mat {
+            Err(err) => return Err(err),
+            $(
+                $p => $blk,
+            )+
+        }
+    );
+
+    ($mat:expr) => (
+        match $mat {
+            Err(err) => return Err(err),
+            _ => ()
+        }
+    )
+)
+
 impl Connection {
     pub fn prepare<'a>(&'a self, query: &str)
             -> Result<~PreparedStatement<'a>, ~str> {
@@ -85,18 +101,25 @@ impl Connection {
 
     pub fn query<T>(&self, query: &str, blk: &fn (&mut ResultIterator) -> T)
             -> Result<T, ~str> {
-        let stmt = match self.prepare(query) {
-            Ok(stmt) => stmt,
-            Err(err) => return Err(err)
-        };
-
-        let mut it = match stmt.query() {
-            Ok(it) => it,
-            Err(err) => return Err(err)
-        };
-
+        let stmt = ret_err!(self.prepare(query) { Ok(stmt) => stmt });
+        let mut it = ret_err!(stmt.query() { Ok(it) => it });
         Ok(blk(&mut it))
     }
+
+    pub fn in_transaction<T>(&self, blk: &fn(&Connection) -> Result<T, ~str>)
+            -> Result<T, ~str> {
+        ret_err!(self.update("BEGIN"));
+
+        let ret = blk(self);
+
+        // TODO: What to do with errors here?
+        match ret {
+            Ok(_) => self.update("COMMIT"),
+            Err(_) => self.update("ROLLBACK")
+        };
+
+        ret
+    }
 }
 
 pub struct PreparedStatement<'self> {
@@ -124,7 +147,7 @@ impl<'self> PreparedStatement<'self> {
         }
     }
 
-    pub fn query<'a>(&'a self) -> Result<ResultIterator<'a>, ~str> {
+    pub fn query(&'self self) -> Result<ResultIterator<'self>, ~str> {
         unsafe { ffi::sqlite3_reset(self.stmt); }
         Ok(ResultIterator {stmt: self})
     }
diff --git a/src/sqlite3/test.rs b/src/sqlite3/test.rs
index e6b45f33..2856f9bd 100644
--- a/src/sqlite3/test.rs
+++ b/src/sqlite3/test.rs
@@ -1,13 +1,22 @@
 extern mod sqlite3;
 
+macro_rules! chk (
+    ($result:expr) => (
+        match $result {
+            Err(err) => fail!(err),
+            Ok(val) => val
+        }
+    )
+)
+
 #[test]
-fn test() {
-    let conn = sqlite3::open("db").unwrap();
-    conn.update("DROP TABLE IF EXISTS foo");
-    conn.update("CREATE TABLE foo (
+fn test_basic() {
+    let conn = chk!(sqlite3::open("db1"));
+    chk!(conn.update("DROP TABLE IF EXISTS foo"));
+    chk!(conn.update("CREATE TABLE foo (
                     id  BIGINT PRIMARY KEY
-                 )");
-    conn.update("INSERT INTO foo (id) VALUES (101), (102)");
+                 )"));
+    chk!(conn.update("INSERT INTO foo (id) VALUES (101), (102)"));
 
     do conn.query("SELECT id FROM foo") |it| {
         for it.advance |row| {
@@ -15,3 +24,28 @@ fn test() {
         }
     };
 }
+
+#[test]
+fn test_trans() {
+    let conn = chk!(sqlite3::open("db2"));
+    chk!(conn.update("DROP TABLE IF EXISTS bar"));
+    chk!(conn.update("CREATE TABLE bar (
+                        id BIGINT PRIMARY KEY
+                     )"));
+    do conn.in_transaction |conn| {
+        chk!(conn.update("INSERT INTO bar (id) VALUES (100)"));
+        Err::<(), ~str>(~"")
+    };
+    assert_eq!(0, chk!(conn.query("SELECT COUNT(*) FROM bar", |it| {
+        it.next().get().get(0).get()
+    })));
+
+    do conn.in_transaction |conn| {
+        chk!(conn.update("INSERT INTO bar (id) VALUES (100)"));
+        Ok(())
+    };
+
+    assert_eq!(1, chk!(conn.query("SELECT COUNT(*) FROM bar", |it| {
+        it.next().get().get(0).get()
+    })));
+}