Avoid double-closing the connection

lib/Database/PostgreSQL/Opium.hs

@@ -99,12 +99,8 @@ execParams
   -> a
   -> ExceptT Error IO Result
 execParams conn query params = do
-  let queryBytes = Encoding.encodeUtf8 query
   -- Actually run the query while locking the connection.
-  mbResult <-
-    flip withRawConnection conn $ \mbRawConn -> do
-      rawConn <- mbRawConn `orThrow` ErrorConnectionClosed
-      liftIO $ LibPQ.execParams rawConn queryBytes (toParamList params) LibPQ.Binary
+  mbResult <- ExceptT runQuery
   -- Check whether the result is valid or nah.
   result <- mbResult `orThrow` ErrorNoResult
   status <- liftIO $ LibPQ.resultStatus result
@@ -117,3 +113,9 @@ execParams conn query params = do
   where
     orThrow (Just x) _ = pure x
     orThrow Nothing  e = throwE e
+
+    queryBytes = Encoding.encodeUtf8 query
+    runQuery =
+      flip withRawConnection conn $ \mbRawConn -> runExceptT $ do
+        rawConn <- mbRawConn `orThrow` ErrorConnectionClosed
+        liftIO $ LibPQ.execParams rawConn queryBytes (toParamList params) LibPQ.Binary

lib/Database/PostgreSQL/Opium/Connection.hs

@@ -1,4 +1,5 @@
 {-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE OverloadedRecordDot #-}
 
 module Database.PostgreSQL.Opium.Connection
   ( Connection
@@ -8,24 +9,28 @@ module Database.PostgreSQL.Opium.Connection
   , close
   ) where
 
+import Control.Concurrent.MVar (MVar, newMVar, modifyMVar_, withMVar)
 import Data.ByteString (ByteString)
 import GHC.Stack (HasCallStack)
 
 import qualified Database.PostgreSQL.LibPQ as LibPQ
 
-newtype Connection = Connection LibPQ.Connection
+newtype Connection = Connection
+  { rawConnection :: MVar (Maybe LibPQ.Connection)
+  } 
 
 withRawConnection
-  :: (Maybe LibPQ.Connection -> m a)
+  :: HasCallStack
+  => (Maybe LibPQ.Connection -> IO a)
   -> Connection
-  -> m a
-withRawConnection f (Connection rawConn) = f (Just rawConn)
+  -> IO a
+withRawConnection f conn = withMVar conn.rawConnection f
 
 unsafeWithRawConnection
   :: HasCallStack
-  => (LibPQ.Connection -> m a)
+  => (LibPQ.Connection -> IO a)
   -> Connection
-  -> m a
+  -> IO a
 unsafeWithRawConnection f = withRawConnection $ \case
   Nothing -> error "raw connection is missing! perhaps the connection was already closed."
   Just rawConn -> f rawConn
@@ -33,7 +38,13 @@ unsafeWithRawConnection f = withRawConnection $ \case
 connect :: ByteString -> IO Connection
 connect connectionString = do
   rawConn <- LibPQ.connectdb connectionString
-  pure $ Connection rawConn
+
+  Connection <$> newMVar (Just rawConn)
 
 close :: Connection -> IO ()
-close (Connection rawConn) = LibPQ.finish rawConn
+close conn = modifyMVar_ conn.rawConnection $ \case
+  Just rawConn -> do
+    LibPQ.finish rawConn
+    pure Nothing
+  Nothing ->
+    pure Nothing