From 09c6adeabc533537dcc039e2195123c6f7167ebe Mon Sep 17 00:00:00 2001
From: Evgeny Poberezkin <2769109+epoberezkin@users.noreply.github.com>
Date: Thu, 24 Jun 2021 18:39:59 +0100
Subject: [PATCH] pass migration as parameter (to use in simplex-chat) (#161)

* pass migration as parameter (to use in simplex-chat)

* add connId parameter to create/joinConnection
---
 package.yaml                                  |  1 +
 src/Simplex/Messaging/Agent.hs                | 37 ++++++++++---------
 src/Simplex/Messaging/Agent/Client.hs         |  6 ++-
 src/Simplex/Messaging/Agent/Env/SQLite.hs     |  3 +-
 src/Simplex/Messaging/Agent/Store/SQLite.hs   | 13 ++++---
 .../Agent/Store/SQLite/Migrations.hs          |  3 +-
 tests/AgentTests.hs                           |  9 ++---
 tests/AgentTests/SQLiteTests.hs               |  3 +-
 8 files changed, 42 insertions(+), 33 deletions(-)

diff --git a/package.yaml b/package.yaml
index a12a0333a..ce4f42a9c 100644
--- a/package.yaml
+++ b/package.yaml
@@ -30,6 +30,7 @@ dependencies:
   - base >= 4.7 && < 5
   - base64-bytestring >= 1.0 && < 1.3
   - bytestring == 0.10.*
+  - composition == 1.0.*
   - constraints == 0.12.*
   - containers == 0.6.*
   - cryptonite == 0.27.*
diff --git a/src/Simplex/Messaging/Agent.hs b/src/Simplex/Messaging/Agent.hs
index 84b5f69bf..0e71cdc02 100644
--- a/src/Simplex/Messaging/Agent.hs
+++ b/src/Simplex/Messaging/Agent.hs
@@ -33,6 +33,7 @@ module Simplex.Messaging.Agent
     runAgentClient,
 
     -- * SMP agent functional API
+    AgentClient (..),
     AgentMonad,
     AgentErrorMonad,
     getSMPAgentClient,
@@ -64,9 +65,11 @@ import Crypto.Random (MonadRandom)
 import Data.Bifunctor (second)
 import Data.ByteString.Char8 (ByteString)
 import qualified Data.ByteString.Char8 as B
+import Data.Composition ((.:), (.:.))
 import Data.Functor (($>))
 import Data.List.NonEmpty (NonEmpty (..))
 import qualified Data.List.NonEmpty as L
+import Data.Maybe (fromMaybe)
 import qualified Data.Text as T
 import Data.Text.Encoding (decodeUtf8)
 import Data.Time.Clock
@@ -83,7 +86,7 @@ import qualified Simplex.Messaging.Protocol as SMP
 import Simplex.Messaging.Transport (ATransport (..), TProxy, Transport (..), runTransportServer)
 import Simplex.Messaging.Util (bshow)
 import System.Random (randomR)
-import UnliftIO.Async (Async, async, race_)
+import UnliftIO.Async (async, race_)
 import qualified UnliftIO.Exception as E
 import UnliftIO.STM
 
@@ -111,14 +114,14 @@ runSMPAgentBlocking (ATransport t) started cfg@AgentConfig {tcpPort} = runReader
         `E.finally` disconnectServers c
 
 -- | Creates an SMP agent client instance
-getSMPAgentClient :: (MonadRandom m, MonadUnliftIO m) => AgentConfig -> m (Async (), AgentClient)
+getSMPAgentClient :: (MonadRandom m, MonadUnliftIO m) => AgentConfig -> m AgentClient
 getSMPAgentClient cfg = newSMPAgentEnv cfg >>= runReaderT runAgent
   where
     runAgent = do
       c <- getAgentClient
       st <- agentDB
       action <- async $ subscriber c st `E.finally` disconnectServers c
-      pure (action, c)
+      pure c {smpSubscriber = action}
 
 disconnectServers :: MonadUnliftIO m => AgentClient -> m ()
 disconnectServers c = closeSMPServerClients c >> logConnection c False
@@ -127,20 +130,20 @@ disconnectServers c = closeSMPServerClients c >> logConnection c False
 type AgentErrorMonad m = (MonadUnliftIO m, MonadError AgentErrorType m)
 
 -- | Create SMP agent connection (NEW command) in Reader monad
-createConnection' :: AgentMonad m => AgentClient -> m (ConnId, SMPQueueInfo)
-createConnection' c = newConn c "" Nothing 0
+createConnection' :: AgentMonad m => AgentClient -> Maybe ConnId -> m (ConnId, SMPQueueInfo)
+createConnection' c connId = newConn c (fromMaybe "" connId) Nothing 0
 
 -- | Create SMP agent connection (NEW command)
-createConnection :: AgentErrorMonad m => AgentClient -> m (ConnId, SMPQueueInfo)
-createConnection c = createConnection' c `runReaderT` agentEnv c
+createConnection :: AgentErrorMonad m => AgentClient -> Maybe ConnId -> m (ConnId, SMPQueueInfo)
+createConnection c = (`runReaderT` agentEnv c) . createConnection' c
 
 -- | Join SMP agent connection (JOIN command) in Reader monad
-joinConnection' :: AgentMonad m => AgentClient -> SMPQueueInfo -> m ConnId
-joinConnection' c qInfo = joinConn c "" qInfo (ReplyMode On) Nothing 0
+joinConnection' :: AgentMonad m => AgentClient -> Maybe ConnId -> SMPQueueInfo -> m ConnId
+joinConnection' c connId qInfo = joinConn c (fromMaybe "" connId) qInfo (ReplyMode On) Nothing 0
 
 -- | Join SMP agent connection (JOIN command)
-joinConnection :: AgentErrorMonad m => AgentClient -> SMPQueueInfo -> m ConnId
-joinConnection c qInfo = joinConnection' c qInfo `runReaderT` agentEnv c
+joinConnection :: AgentErrorMonad m => AgentClient -> Maybe ConnId -> SMPQueueInfo -> m ConnId
+joinConnection c = (`runReaderT` agentEnv c) .: joinConnection' c
 
 -- | Accept invitation (ACPT command) in Reader monad
 acceptInvitation' :: AgentMonad m => AgentClient -> InvitationId -> ConnInfo -> m ConnId
@@ -148,27 +151,27 @@ acceptInvitation' c = acceptInv c ""
 
 -- | Accept invitation (ACPT command)
 acceptInvitation :: AgentErrorMonad m => AgentClient -> InvitationId -> ConnInfo -> m ConnId
-acceptInvitation c invId cInfo = acceptInvitation c invId cInfo `runReaderT` agentEnv c
+acceptInvitation c = (`runReaderT` agentEnv c) .: acceptInvitation c
 
 -- | Send introduction of the second connection the first (INTRO command)
 sendIntroduction :: AgentErrorMonad m => AgentClient -> ConnId -> ConnId -> ConnInfo -> m ()
-sendIntroduction c toConn reConn reInfo = sendIntroduction' c toConn reConn reInfo `runReaderT` agentEnv c
+sendIntroduction c = (`runReaderT` agentEnv c) .:. sendIntroduction' c
 
 -- | Subscribe to receive connection messages (SUB command)
 subscribeConnection :: AgentErrorMonad m => AgentClient -> ConnId -> m ()
-subscribeConnection c connId = subscribeConnection' c connId `runReaderT` agentEnv c
+subscribeConnection c = (`runReaderT` agentEnv c) . subscribeConnection' c
 
 -- | Send message to the connection (SEND command)
 sendMessage :: AgentErrorMonad m => AgentClient -> ConnId -> MsgBody -> m InternalId
-sendMessage c connId msgBody = sendMessage' c connId msgBody `runReaderT` agentEnv c
+sendMessage c = (`runReaderT` agentEnv c) .: sendMessage' c
 
 -- | Suspend SMP agent connection (OFF command)
 suspendConnection :: AgentErrorMonad m => AgentClient -> ConnId -> m ()
-suspendConnection c connId = suspendConnection' c connId `runReaderT` agentEnv c
+suspendConnection c = (`runReaderT` agentEnv c) . suspendConnection' c
 
 -- | Delete SMP agent connection (DEL command)
 deleteConnection :: AgentErrorMonad m => AgentClient -> ConnId -> m ()
-deleteConnection c connId = deleteConnection' c connId `runReaderT` agentEnv c
+deleteConnection c = (`runReaderT` agentEnv c) . deleteConnection' c
 
 -- | Creates an SMP agent client instance that receives commands and sends responses via 'TBQueue's.
 getAgentClient :: (MonadUnliftIO m, MonadReader Env m) => m AgentClient
diff --git a/src/Simplex/Messaging/Agent/Client.hs b/src/Simplex/Messaging/Agent/Client.hs
index efb9d6699..abc35e47a 100644
--- a/src/Simplex/Messaging/Agent/Client.hs
+++ b/src/Simplex/Messaging/Agent/Client.hs
@@ -31,6 +31,7 @@ module Simplex.Messaging.Agent.Client
   )
 where
 
+import Control.Concurrent.Async (Async)
 import Control.Concurrent.STM (stateTVar)
 import Control.Logger.Simple
 import Control.Monad.Except
@@ -68,7 +69,8 @@ data AgentClient = AgentClient
     subscrConns :: TVar (Map ConnId SMPServer),
     clientId :: Int,
     store :: SQLiteStore,
-    agentEnv :: Env
+    agentEnv :: Env,
+    smpSubscriber :: Async ()
   }
 
 newAgentClient :: SQLiteStore -> Env -> STM AgentClient
@@ -81,7 +83,7 @@ newAgentClient store agentEnv = do
   subscrSrvrs <- newTVar M.empty
   subscrConns <- newTVar M.empty
   clientId <- stateTVar (clientCounter agentEnv) $ \i -> (i + 1, i + 1)
-  return AgentClient {rcvQ, subQ, msgQ, smpClients, subscrSrvrs, subscrConns, clientId, store, agentEnv}
+  return AgentClient {rcvQ, subQ, msgQ, smpClients, subscrSrvrs, subscrConns, clientId, store, agentEnv, smpSubscriber = undefined}
 
 -- | Agent monad with MonadReader Env and MonadError AgentErrorType
 type AgentMonad m = (MonadUnliftIO m, MonadReader Env m, MonadError AgentErrorType m)
diff --git a/src/Simplex/Messaging/Agent/Env/SQLite.hs b/src/Simplex/Messaging/Agent/Env/SQLite.hs
index 13445643a..77e1e1d6f 100644
--- a/src/Simplex/Messaging/Agent/Env/SQLite.hs
+++ b/src/Simplex/Messaging/Agent/Env/SQLite.hs
@@ -12,6 +12,7 @@ import Network.Socket
 import Numeric.Natural
 import Simplex.Messaging.Agent.Protocol (SMPServer)
 import Simplex.Messaging.Agent.Store.SQLite
+import qualified Simplex.Messaging.Agent.Store.SQLite.Migrations as Migrations
 import Simplex.Messaging.Client
 import System.Random (StdGen, newStdGen)
 import UnliftIO.STM
@@ -37,7 +38,7 @@ data Env = Env
 newSMPAgentEnv :: (MonadUnliftIO m, MonadRandom m) => AgentConfig -> m Env
 newSMPAgentEnv config = do
   idsDrg <- newTVarIO =<< drgNew
-  _ <- liftIO $ createSQLiteStore $ dbFile config
+  _ <- liftIO $ createSQLiteStore (dbFile config) Migrations.app
   clientCounter <- newTVarIO 0
   randomServer <- newTVarIO =<< liftIO newStdGen
   return Env {config, idsDrg, clientCounter, reservedMsgSize, randomServer}
diff --git a/src/Simplex/Messaging/Agent/Store/SQLite.hs b/src/Simplex/Messaging/Agent/Store/SQLite.hs
index e37096813..2de757207 100644
--- a/src/Simplex/Messaging/Agent/Store/SQLite.hs
+++ b/src/Simplex/Messaging/Agent/Store/SQLite.hs
@@ -47,6 +47,7 @@ import Database.SQLite.Simple.ToField (ToField (..))
 import Network.Socket (ServiceName)
 import Simplex.Messaging.Agent.Protocol
 import Simplex.Messaging.Agent.Store
+import Simplex.Messaging.Agent.Store.SQLite.Migrations (Migration)
 import qualified Simplex.Messaging.Agent.Store.SQLite.Migrations as Migrations
 import Simplex.Messaging.Parsers (blobFieldParser)
 import qualified Simplex.Messaging.Protocol as SMP
@@ -66,8 +67,8 @@ data SQLiteStore = SQLiteStore
     dbNew :: Bool
   }
 
-createSQLiteStore :: FilePath -> IO SQLiteStore
-createSQLiteStore dbFilePath = do
+createSQLiteStore :: FilePath -> [Migration] -> IO SQLiteStore
+createSQLiteStore dbFilePath migrations = do
   let dbDir = takeDirectory dbFilePath
   createDirectoryIfMissing False dbDir
   store <- connectSQLiteStore dbFilePath
@@ -77,13 +78,13 @@ createSQLiteStore dbFilePath = do
     Just "THREADSAFE=0" -> confirmOrExit "SQLite compiled with non-threadsafe code."
     Nothing -> putStrLn "Warning: SQLite THREADSAFE compile option not found"
     _ -> return ()
-  migrateSchema store
+  migrateSchema store migrations
   pure store
 
-migrateSchema :: SQLiteStore -> IO ()
-migrateSchema SQLiteStore {dbConn, dbFilePath, dbNew} = do
+migrateSchema :: SQLiteStore -> [Migration] -> IO ()
+migrateSchema SQLiteStore {dbConn, dbFilePath, dbNew} migrations = do
   Migrations.initialize dbConn
-  Migrations.get dbConn Migrations.app >>= \case
+  Migrations.get dbConn migrations >>= \case
     Left e -> confirmOrExit $ "Database error: " <> e
     Right [] -> pure ()
     Right ms -> do
diff --git a/src/Simplex/Messaging/Agent/Store/SQLite/Migrations.hs b/src/Simplex/Messaging/Agent/Store/SQLite/Migrations.hs
index b022ea5bf..4e6128493 100644
--- a/src/Simplex/Messaging/Agent/Store/SQLite/Migrations.hs
+++ b/src/Simplex/Messaging/Agent/Store/SQLite/Migrations.hs
@@ -7,7 +7,8 @@
 {-# LANGUAGE TupleSections #-}
 
 module Simplex.Messaging.Agent.Store.SQLite.Migrations
-  ( app,
+  ( Migration (..),
+    app,
     initialize,
     get,
     run,
diff --git a/tests/AgentTests.hs b/tests/AgentTests.hs
index c6c588f97..34993c8e8 100644
--- a/tests/AgentTests.hs
+++ b/tests/AgentTests.hs
@@ -21,7 +21,6 @@ import qualified Data.ByteString.Char8 as B
 import SMPAgentClient
 import SMPClient (withSmpServer)
 import Simplex.Messaging.Agent
-import Simplex.Messaging.Agent.Client
 import Simplex.Messaging.Agent.Env.SQLite (dbFile)
 import Simplex.Messaging.Agent.Protocol
 import Simplex.Messaging.Agent.Store (InternalId (..))
@@ -126,11 +125,11 @@ testDuplexConnection _ alice bob = do
 
 testAgentClient :: IO ()
 testAgentClient = do
-  (_, alice) <- getSMPAgentClient cfg
-  (_, bob) <- getSMPAgentClient cfg {dbFile = testDB2}
+  alice <- getSMPAgentClient cfg
+  bob <- getSMPAgentClient cfg {dbFile = testDB2}
   Right () <- runExceptT $ do
-    (bobId, qInfo) <- createConnection alice
-    aliceId <- joinConnection bob qInfo
+    (bobId, qInfo) <- createConnection alice Nothing
+    aliceId <- joinConnection bob Nothing qInfo
     get alice ##> ("", bobId, CON)
     get bob ##> ("", aliceId, CON)
     InternalId 1 <- sendMessage alice bobId "hello"
diff --git a/tests/AgentTests/SQLiteTests.hs b/tests/AgentTests/SQLiteTests.hs
index 2f8383a8c..c3c1bcdd2 100644
--- a/tests/AgentTests/SQLiteTests.hs
+++ b/tests/AgentTests/SQLiteTests.hs
@@ -24,6 +24,7 @@ import SMPClient (testKeyHash)
 import Simplex.Messaging.Agent.Protocol
 import Simplex.Messaging.Agent.Store
 import Simplex.Messaging.Agent.Store.SQLite
+import qualified Simplex.Messaging.Agent.Store.SQLite.Migrations as Migrations
 import qualified Simplex.Messaging.Crypto as C
 import System.Random (Random (randomIO))
 import Test.Hspec
@@ -49,7 +50,7 @@ createStore = do
   -- Randomize DB file name to avoid SQLite IO errors supposedly caused by asynchronous
   -- IO operations on multiple similarly named files; error seems to be environment specific
   r <- randomIO :: IO Word32
-  createSQLiteStore $ testDB <> show r
+  createSQLiteStore (testDB <> show r) Migrations.app
 
 removeStore :: SQLiteStore -> IO ()
 removeStore store = do