From 9a2dbbd3cceaa93513d09e1a116628d9a32342f0 Mon Sep 17 00:00:00 2001
From: Paul Brinkmeier <hallo@pbrinkmeier.de>
Date: Tue, 15 Jul 2025 23:22:52 +0200
Subject: [PATCH] Avoid double-closing the connection

---
 lib/Database/PostgreSQL/Opium.hs            | 12 +++++----
 lib/Database/PostgreSQL/Opium/Connection.hs | 27 +++++++++++++++------
 test/SpecHook.hs                            |  2 +-
 3 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/lib/Database/PostgreSQL/Opium.hs b/lib/Database/PostgreSQL/Opium.hs
index b4c4244..aaf84b3 100644
--- a/lib/Database/PostgreSQL/Opium.hs
+++ b/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
diff --git a/lib/Database/PostgreSQL/Opium/Connection.hs b/lib/Database/PostgreSQL/Opium/Connection.hs
index dc6c2c3..ff8c091 100644
--- a/lib/Database/PostgreSQL/Opium/Connection.hs
+++ b/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
diff --git a/test/SpecHook.hs b/test/SpecHook.hs
index 102a7e6..507c920 100644
--- a/test/SpecHook.hs
+++ b/test/SpecHook.hs
@@ -22,6 +22,7 @@ setupConnection = do
   let dsn = printf "host=localhost user=%s password=%s dbname=%s port=%s" dbUser dbPass dbName dbPort
   conn <- Opium.connect $ Encoding.encodeUtf8 $ Text.pack dsn
 
+  Right _ <- Opium.execute_ "DROP TABLE IF EXISTS person" conn
   Right _ <- Opium.execute_ "CREATE TABLE person (name TEXT NOT NULL, age INT NOT NULL, score DOUBLE PRECISION NOT NULL, motto TEXT)" conn
   Right _ <- Opium.execute_ "INSERT INTO person VALUES ('paul', 25, 30), ('albus', 103, 50.42)" conn
 
@@ -29,7 +30,6 @@ setupConnection = do
 
 teardownConnection :: Opium.Connection -> IO ()
 teardownConnection conn = do
-  Right _ <- Opium.execute_ "DROP TABLE person" conn
   Opium.close conn
 
 hook :: SpecWith Opium.Connection -> Spec