From 6bb52d057489220f5410c8756c011ade6b926477 Mon Sep 17 00:00:00 2001
From: Nikita Volkov <nikita.y.volkov@mail.ru>
Date: Fri, 13 Sep 2024 22:42:14 +0300
Subject: [PATCH] Add the resetting of the registry in case of errors

---
 library/Hasql/PreparedStatementRegistry.hs | 12 +++++++++++-
 library/Hasql/Session/Core.hs              |  7 +++++--
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/library/Hasql/PreparedStatementRegistry.hs b/library/Hasql/PreparedStatementRegistry.hs
index 37fae6a..d876f52 100644
--- a/library/Hasql/PreparedStatementRegistry.hs
+++ b/library/Hasql/PreparedStatementRegistry.hs
@@ -2,6 +2,7 @@ module Hasql.PreparedStatementRegistry
   ( PreparedStatementRegistry,
     new,
     update,
+    reset,
     LocalKey (..),
   )
 where
@@ -9,7 +10,7 @@ where
 import ByteString.StrictBuilder qualified as B
 import Data.HashTable.IO qualified as A
 import Hasql.LibPq14 qualified as Pq
-import Hasql.Prelude hiding (lookup)
+import Hasql.Prelude hiding (lookup, reset)
 
 data PreparedStatementRegistry
   = PreparedStatementRegistry !(A.BasicHashTable LocalKey ByteString) !(IORef Word)
@@ -42,6 +43,15 @@ update localKey onNewRemoteKey onOldRemoteKey (PreparedStatementRegistry table c
     old =
       onOldRemoteKey
 
+reset :: PreparedStatementRegistry -> IO ()
+reset (PreparedStatementRegistry table counter) = do
+  -- TODO: This is a temporary measure.
+  -- We should just move to a pure implementation.
+  do
+    entries <- A.toList table
+    forM_ entries \(k, _) -> A.delete table k
+  writeIORef counter 0
+
 -- |
 -- Local statement key.
 data LocalKey
diff --git a/library/Hasql/Session/Core.hs b/library/Hasql/Session/Core.hs
index d0e18c9..f544632 100644
--- a/library/Hasql/Session/Core.hs
+++ b/library/Hasql/Session/Core.hs
@@ -11,6 +11,7 @@ import Hasql.IO qualified as IO
 import Hasql.LibPq14 qualified as Pq
 import Hasql.Pipeline.Core qualified as Pipeline
 import Hasql.Prelude
+import Hasql.PreparedStatementRegistry qualified as PreparedStatementRegistry
 import Hasql.Statement qualified as Statement
 
 -- |
@@ -29,11 +30,13 @@ run (Session impl) connection =
       runExceptT $ runReaderT impl connection
     handler =
       case connection of
-        Connection.Connection pqConnVar _ _ ->
+        Connection.Connection pqConnVar _ registry ->
           withMVar pqConnVar \pqConn ->
             Pq.transactionStatus pqConn >>= \case
               Pq.TransIdle -> pure ()
-              _ -> Pq.reset pqConn
+              _ -> do
+                PreparedStatementRegistry.reset registry
+                Pq.reset pqConn
 
 -- |
 -- Possibly a multi-statement query,