From 1b5b9c25ca839c0714755877bc81e64ea8e4373c Mon Sep 17 00:00:00 2001
From: Heinrich Apfelmus <apfelmus@quantentunnel.de>
Date: Fri, 5 Apr 2024 16:21:57 +0200
Subject: [PATCH] Add `withDBHandleInMemory`

---
 .../src/Cardano/Wallet/Deposit/IO/DB.hs       | 32 ++++++++++++++++---
 lib/wallet/src/Cardano/DB/Sqlite.hs           |  9 ++++++
 2 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/IO/DB.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/IO/DB.hs
index 2fde257174c..aba0afc80c2 100644
--- a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/IO/DB.hs
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/IO/DB.hs
@@ -4,6 +4,7 @@ module Cardano.Wallet.Deposit.IO.DB
     ( SqlM
     , SqlContext (..)
     , withSqliteFile
+    , withSqlContextInMemory
 
     , DBLog (..)
     ) where
@@ -17,10 +18,12 @@ import Cardano.DB.Sqlite
     ( DBLog (..)
     , dbBackend
     , withDBHandle
+    , withDBHandleInMemory
     )
 import Control.Concurrent.MVar
     ( newMVar
     , withMVar
+    , withMVarMasked
     )
 import Control.Tracer
     ( Tracer
@@ -33,7 +36,7 @@ import Data.Time.Clock
 import qualified Database.Persist.Sql as Persistent
 
 {-----------------------------------------------------------------------------
-    Comment layout
+    SqlContext
 ------------------------------------------------------------------------------}
 -- | Monad to run SQL queries in.
 type SqlM = Persistent.SqlPersistT IO
@@ -41,9 +44,28 @@ type SqlM = Persistent.SqlPersistT IO
 -- | A facility to run 'SqlM' computations.
 -- Importantly, computations are not run in parallel, but sequenced.
 newtype SqlContext = SqlContext
-    { runQuery :: forall a. SqlM a -> IO a
+    { runSqlM :: forall a. SqlM a -> IO a
     }
 
+-- | Acquire and release an 'SqlContext' in memory.
+withSqlContextInMemory
+    :: Tracer IO DBLog
+    -- ^ Logging
+    -> (SqlContext -> IO a)
+    -- ^ Action to run
+    -> IO a
+withSqlContextInMemory tr action = do
+    withDBHandleInMemory tr $ \dbhandle -> do
+        -- Lock ensures that database operations are sequenced.
+        lock <- newMVar (dbBackend dbhandle)
+        let runSqlM :: SqlM a -> IO a
+            runSqlM cmd =
+                withMVarMasked lock (observe . Persistent.runSqlConn cmd)
+        action $ SqlContext{runSqlM}
+  where
+    observe :: IO a -> IO a
+    observe = bracketTracer (contramap MsgRun tr)
+
 -- | Open an .sqlite database file
 -- and provide an 'SqlContext' for running 'SqlM' actions.
 withSqliteFile
@@ -61,8 +83,8 @@ withSqliteFile tr fp action = do
         let
             -- Run a query on the open database,
             -- but retry on busy.
-            runQuery :: SqlM a -> IO a
-            runQuery cmd =
+            runSqlM :: SqlM a -> IO a
+            runSqlM cmd =
                 observe
                     . retryOnBusy tr retryOnBusyTimeout
                     $ withMVar lock
@@ -70,7 +92,7 @@ withSqliteFile tr fp action = do
                     $ Persistent.runSqlConn cmd
                     $ dbBackend dbHandle
         in
-            action $ SqlContext{runQuery}
+            action $ SqlContext{runSqlM}
   where
     observe :: IO a -> IO a
     observe = bracketTracer (contramap MsgRun tr)
diff --git a/lib/wallet/src/Cardano/DB/Sqlite.hs b/lib/wallet/src/Cardano/DB/Sqlite.hs
index bf48edb300e..14b8630fa75 100644
--- a/lib/wallet/src/Cardano/DB/Sqlite.hs
+++ b/lib/wallet/src/Cardano/DB/Sqlite.hs
@@ -32,6 +32,7 @@ module Cardano.DB.Sqlite
       -- * DB Connections
     , DBHandle
     , withDBHandle
+    , withDBHandleInMemory
     , dbConn
     , dbFile
     , dbBackend
@@ -302,6 +303,14 @@ withDBHandle
 withDBHandle tr fp =
     bracket (newDBHandle tr fp) (closeDBHandleRetrying tr)
 
+-- | Acquire and release a 'DBHandle' in memory.
+withDBHandleInMemory
+    :: Tracer IO DBLog
+    -> (DBHandle -> IO a)
+    -> IO a
+withDBHandleInMemory tr =
+    bracket (newDBHandleInMemory tr) (closeDBHandle tr)
+
 -- | Create a new 'DBHandle' from a file.
 -- Needs to be closed explicitly.
 newDBHandle