From f4d00960e71624902bb9ea8fb9ff101923827f29 Mon Sep 17 00:00:00 2001
From: Evgeny Poberezkin <evgeny@poberezkin.com>
Date: Tue, 9 Sep 2025 20:00:37 +0100
Subject: [PATCH] refactor

---
 src/Simplex/Messaging/Server.hs               |  4 +--
 src/Simplex/Messaging/Server/Main.hs          |  2 +-
 .../Messaging/Server/MsgStore/Journal.hs      |  8 ++---
 src/Simplex/Messaging/Server/MsgStore/STM.hs  |  2 +-
 .../Messaging/Server/MsgStore/Types.hs        |  4 +--
 .../Messaging/Server/QueueStore/Postgres.hs   | 36 ++++++++++++-------
 tests/CoreTests/MsgStoreTests.hs              |  2 +-
 7 files changed, 34 insertions(+), 24 deletions(-)

diff --git a/src/Simplex/Messaging/Server.hs b/src/Simplex/Messaging/Server.hs
index 5814d546b..61748dd5b 100644
--- a/src/Simplex/Messaging/Server.hs
+++ b/src/Simplex/Messaging/Server.hs
@@ -2117,7 +2117,7 @@ exportMessages tty st f drainMsgs = do
     StoreMemory ms -> exportMessages_ ms $ getMsgs ms
     StoreJournal ms -> exportMessages_ ms $ getJournalMsgs ms
   where
-    exportMessages_ ms get = fmap (\(Sum n) -> n) . unsafeWithAllMsgQueues tty True ms . saveQueueMsgs get
+    exportMessages_ ms get = fmap (\(Sum n) -> n) . unsafeWithAllMsgQueues tty ms . saveQueueMsgs get
     run :: (Handle -> IO Int) -> IO ()
     run a = liftIO $ withFile f WriteMode $ tryAny . a >=> \case
       Right n -> logNote $ "messages saved: " <> tshow n
@@ -2160,7 +2160,7 @@ processServerMessages StartOptions {skipWarnings} = do
               run processValidateQueue
         | otherwise = logWarn "skipping message expiration" $> Nothing
         where
-          run a = unsafeWithAllMsgQueues False False ms a `catchAny` \_ -> exitFailure
+          run a = unsafeWithAllMsgQueues False ms a `catchAny` \_ -> exitFailure
           processExpireQueue :: Int64 -> JournalQueue s -> IO MessageStats
           processExpireQueue old q = unsafeRunStore q "processExpireQueue" $ do
             mq <- getMsgQueue ms q False
diff --git a/src/Simplex/Messaging/Server/Main.hs b/src/Simplex/Messaging/Server/Main.hs
index c095f7160..59aacc2c4 100644
--- a/src/Simplex/Messaging/Server/Main.hs
+++ b/src/Simplex/Messaging/Server/Main.hs
@@ -591,7 +591,7 @@ exportDatabaseToStoreLog logPath dbOpts storeLogFilePath = do
   ps <- newJournalMsgStore logPath $ PQStoreCfg storeCfg
   sl <- openWriteStoreLog False storeLogFilePath
   Sum sCnt <- foldServiceRecs (postgresQueueStore ps) $ \sr -> logNewService sl sr $> Sum (1 :: Int)
-  Sum qCnt <- foldQueueRecs True True (postgresQueueStore ps) Nothing $ \(rId, qr) -> logCreateQueue sl rId qr $> Sum (1 :: Int)
+  Sum qCnt <- foldQueueRecs True True (postgresQueueStore ps) $ \(rId, qr) -> logCreateQueue sl rId qr $> Sum (1 :: Int)
   closeStoreLog sl
   pure (sCnt, qCnt)
 #endif
diff --git a/src/Simplex/Messaging/Server/MsgStore/Journal.hs b/src/Simplex/Messaging/Server/MsgStore/Journal.hs
index 89b181748..be73ab681 100644
--- a/src/Simplex/Messaging/Server/MsgStore/Journal.hs
+++ b/src/Simplex/Messaging/Server/MsgStore/Journal.hs
@@ -406,11 +406,11 @@ instance MsgStoreClass (JournalMsgStore s) where
 
   -- This function can only be used in server CLI commands or before server is started.
   -- It does not cache queues and is NOT concurrency safe.
-  unsafeWithAllMsgQueues :: Monoid a => Bool -> Bool -> JournalMsgStore s -> (JournalQueue s -> IO a) -> IO a
-  unsafeWithAllMsgQueues tty withData ms action = case queueStore_ ms of
+  unsafeWithAllMsgQueues :: Monoid a => Bool -> JournalMsgStore s -> (JournalQueue s -> IO a) -> IO a
+  unsafeWithAllMsgQueues tty ms action = case queueStore_ ms of
     MQStore st -> withLoadedQueues st run
 #if defined(dbServerPostgres)
-    PQStore st -> foldQueueRecs tty withData st Nothing $ uncurry (mkQueue ms False) >=> run
+    PQStore st -> foldQueueRecs False tty st $ uncurry (mkQueue ms False) >=> run
 #endif
     where
       run q = do
@@ -430,7 +430,7 @@ instance MsgStoreClass (JournalMsgStore s) where
 #if defined(dbServerPostgres)
     PQStore st -> do
       let JournalMsgStore {queueLocks, sharedLock} = ms
-      foldQueueRecs tty False st (Just veryOld) $ \(rId, qr) -> do
+      foldRecentQueueRecs veryOld tty st $ \(rId, qr) -> do
         q <- mkQueue ms False rId qr
         withSharedWaitLock rId queueLocks sharedLock $ run $ tryStore' "deleteExpiredMsgs" rId $
           getLoadedQueue q >>= unStoreIO . expireQueueMsgs ms now old
diff --git a/src/Simplex/Messaging/Server/MsgStore/STM.hs b/src/Simplex/Messaging/Server/MsgStore/STM.hs
index ed24e85a4..837e0481d 100644
--- a/src/Simplex/Messaging/Server/MsgStore/STM.hs
+++ b/src/Simplex/Messaging/Server/MsgStore/STM.hs
@@ -82,7 +82,7 @@ instance MsgStoreClass STMMsgStore where
   {-# INLINE closeMsgStore #-}
   withActiveMsgQueues = withLoadedQueues . queueStore_
   {-# INLINE withActiveMsgQueues #-}
-  unsafeWithAllMsgQueues _ _ = withLoadedQueues . queueStore_
+  unsafeWithAllMsgQueues _ = withLoadedQueues . queueStore_
   {-# INLINE unsafeWithAllMsgQueues #-}
 
   expireOldMessages :: Bool -> STMMsgStore -> Int64 -> Int64 -> IO MessageStats
diff --git a/src/Simplex/Messaging/Server/MsgStore/Types.hs b/src/Simplex/Messaging/Server/MsgStore/Types.hs
index b78b101f5..fbb2e194b 100644
--- a/src/Simplex/Messaging/Server/MsgStore/Types.hs
+++ b/src/Simplex/Messaging/Server/MsgStore/Types.hs
@@ -40,8 +40,8 @@ class (Monad (StoreMonad s), QueueStoreClass (StoreQueue s) (QueueStore s)) => M
   closeMsgStore :: s -> IO ()
   withActiveMsgQueues :: Monoid a => s -> (StoreQueue s -> IO a) -> IO a
   -- This function can only be used in server CLI commands or before server is started.
-  -- tty, withData, store
-  unsafeWithAllMsgQueues :: Monoid a => Bool -> Bool -> s -> (StoreQueue s -> IO a) -> IO a
+  -- tty, store
+  unsafeWithAllMsgQueues :: Monoid a => Bool -> s -> (StoreQueue s -> IO a) -> IO a
   -- tty, store, now, ttl
   expireOldMessages :: Bool -> s -> Int64 -> Int64 -> IO MessageStats
   logQueueStates :: s -> IO ()
diff --git a/src/Simplex/Messaging/Server/QueueStore/Postgres.hs b/src/Simplex/Messaging/Server/QueueStore/Postgres.hs
index e4ab1c65b..ff76759d8 100644
--- a/src/Simplex/Messaging/Server/QueueStore/Postgres.hs
+++ b/src/Simplex/Messaging/Server/QueueStore/Postgres.hs
@@ -25,6 +25,7 @@ module Simplex.Messaging.Server.QueueStore.Postgres
     batchInsertQueues,
     foldServiceRecs,
     foldQueueRecs,
+    foldRecentQueueRecs,
     handleDuplicate,
     withLog_,
     withDB',
@@ -545,8 +546,28 @@ foldServiceRecs st f =
     DB.fold_ db "SELECT service_id, service_role, service_cert, service_cert_hash, created_at FROM services" mempty $
       \ !acc -> fmap (acc <>) . f . rowToServiceRec
 
-foldQueueRecs :: forall a q. Monoid a => Bool -> Bool -> PostgresQueueStore q -> Maybe Int64 -> ((RecipientId, QueueRec) -> IO a) -> IO a
-foldQueueRecs tty withData st skipOld_ f = do
+foldQueueRecs :: Monoid a => Bool -> Bool -> PostgresQueueStore q -> ((RecipientId, QueueRec) -> IO a) -> IO a
+foldQueueRecs withData = foldQueueRecs_ foldRecs
+  where
+    foldRecs db acc f'
+      | withData = DB.fold_ db (queueRecQueryWithData <> cond) acc $ \acc' -> f' acc' . rowToQueueRecWithData
+      | otherwise = DB.fold_ db (queueRecQuery <> cond) acc $ \acc' -> f' acc' . rowToQueueRec
+    cond = " WHERE deleted_at IS NULL ORDER BY recipient_id ASC"
+
+foldRecentQueueRecs :: Monoid a => Int64 -> Bool -> PostgresQueueStore q -> ((RecipientId, QueueRec) -> IO a) -> IO a
+foldRecentQueueRecs old = foldQueueRecs_ foldRecs
+  where
+    foldRecs db acc f' = DB.fold db (queueRecQuery <> cond) (Only old) acc $ \acc' -> f' acc' . rowToQueueRec
+    cond = " WHERE deleted_at IS NULL AND updated_at > ? ORDER BY recipient_id ASC"
+
+foldQueueRecs_ ::
+  Monoid a =>
+  (DB.Connection -> (Int, a) -> ((Int, a) -> (RecipientId, QueueRec) -> IO (Int, a)) -> IO (Int, a)) ->
+  Bool ->
+  PostgresQueueStore q ->
+  ((RecipientId, QueueRec) -> IO a) ->
+  IO a
+foldQueueRecs_ foldRecs tty st f = do
   (n, r) <- withTransaction (dbStore st) $ \db ->
     foldRecs db (0 :: Int, mempty) $ \(i, acc) qr -> do
       r <- f qr
@@ -557,17 +578,6 @@ foldQueueRecs tty withData st skipOld_ f = do
   when tty $ putStrLn $ progress n
   pure r
   where
-    foldRecs db acc f' = case skipOld_ of
-      Nothing
-        | withData -> DB.fold_ db (queueRecQueryWithData <> cond) acc $ \acc' -> f' acc' . rowToQueueRecWithData
-        | otherwise -> DB.fold_ db (queueRecQuery <> cond) acc $ \acc' -> f' acc' . rowToQueueRec
-        where
-          cond = " WHERE deleted_at IS NULL ORDER BY recipient_id ASC"
-      Just old
-        | withData -> DB.fold db (queueRecQueryWithData <> cond) (Only old) acc $ \acc' -> f' acc' . rowToQueueRecWithData
-        | otherwise -> DB.fold db (queueRecQuery <> cond) (Only old) acc $ \acc' -> f' acc' . rowToQueueRec
-        where
-          cond = " WHERE deleted_at IS NULL AND updated_at > ? ORDER BY recipient_id ASC"
     progress i = "Processed: " <> show i <> " records"
 
 queueRecQuery :: Query
diff --git a/tests/CoreTests/MsgStoreTests.hs b/tests/CoreTests/MsgStoreTests.hs
index c5cf3a437..2ed0da330 100644
--- a/tests/CoreTests/MsgStoreTests.hs
+++ b/tests/CoreTests/MsgStoreTests.hs
@@ -54,7 +54,7 @@ msgStoreTests = do
   around (withMsgStore testSMTStoreConfig) $ describe "STM message store" someMsgStoreTests
   around (withMsgStore $ testJournalStoreCfg MQStoreCfg) $ describe "Journal message store" $ do
     someMsgStoreTests
-    fit "should export and import journal store" testExportImportStore
+    it "should export and import journal store" testExportImportStore
     describe "queue state" $ do
       it "should restore queue state from the last line" testQueueState
       it "should recover when message is written and state is not" testMessageState