From 2028500da7b968d919662fdc87568007fb10c5e3 Mon Sep 17 00:00:00 2001
From: "Adam C. Baker" <adam.durandal@gmail.com>
Date: Mon, 13 Oct 2014 13:04:25 -0700
Subject: [PATCH] withSavepoint middleware

---
 src/Middleware.hs       |  8 +++++++-
 test/Unit/ErrorsSpec.hs | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)
 create mode 100644 test/Unit/ErrorsSpec.hs

diff --git a/src/Middleware.hs b/src/Middleware.hs
index 14cfd8bdf..b9c730d8e 100644
--- a/src/Middleware.hs
+++ b/src/Middleware.hs
@@ -11,13 +11,19 @@ import Network.HTTP.Types.Header (hContentType)
 import Network.HTTP.Types.Status (status400)
 import Database.HDBC.Types (SqlError(..))
 import Network.Wai (Application, responseLBS)
-import Control.Exception (finally, catchJust)
+import Control.Exception (finally, throw, catchJust, catch,  SomeException)
 
 
 inTransaction :: (Connection -> Application) -> (Connection -> Application)
 inTransaction app conn req respond =
   finally (runRaw conn "begin" >> app conn req respond) (runRaw conn "commit")
 
+withSavepoint :: (Connection -> Application) -> Connection -> Application
+withSavepoint app conn req respond = do
+  runRaw conn "savepoint req_sp"
+  catch (app conn req respond) (\e -> let _ = (e::SomeException) in
+    runRaw conn "rollback to savepoint req_sp" >> throw e)
+
 instance ToJSON SqlError where
   toJSON t = object [
       "error" .= object [
diff --git a/test/Unit/ErrorsSpec.hs b/test/Unit/ErrorsSpec.hs
new file mode 100644
index 000000000..c6f8f84d3
--- /dev/null
+++ b/test/Unit/ErrorsSpec.hs
@@ -0,0 +1,36 @@
+{-# LANGUAGE OverloadedStrings #-}
+module Unit.ErrorsSpec where
+
+import Test.Hspec
+
+import Database.HDBC (runRaw, quickQuery, fromSql, SqlError)
+import SpecHelper (dbWithSchema)
+import Middleware (withSavepoint)
+import PgQuery (insert)
+import Types(SqlRow(..))
+import Control.Exception(catch)
+import Control.Monad(void)
+import Network.Wai (defaultRequest, responseLBS)
+import Network.HTTP.Types.Status (ok200)
+
+spec :: Spec
+spec = let
+  dbErrApp conn _ res = do
+    putStrLn "In fake app"
+    _ <- insert "1" "items" (SqlRow []) conn
+    runRaw conn "select 1/0"
+    _ <- insert "1" "items" (SqlRow []) conn
+    res $ responseLBS ok200 [("Content-Type", "application/json")] "{}"
+  in around dbWithSchema $ do
+
+    describe "withSavepoint" $ do
+      it "allows partial rollback of request" $ \c -> do
+        let app = withSavepoint dbErrApp c
+        [[beforeCount]] <- quickQuery c "select count(*) from \"1\".items" []
+        runRaw c "set role dbapi_anonymous"
+        _ <- insert "1" "items" (SqlRow []) c
+        catch (void $ app defaultRequest (const undefined) ) $
+          \e -> let _ = (e::SqlError) in do
+          _ <- insert "1" "items" (SqlRow []) c
+          [[afterCount]] <- quickQuery c "select count(*) from \"1\".items" []
+          fromSql afterCount `shouldBe` (fromSql beforeCount::Int) + 2