agent: store interface (#1436)

src/Simplex/Messaging/Agent.hs

@@ -170,6 +170,7 @@ import Simplex.Messaging.Agent.Store
 import Simplex.Messaging.Agent.Store.AgentStore
 import Simplex.Messaging.Agent.Store.Common (DBStore)
 import qualified Simplex.Messaging.Agent.Store.DB as DB
+import Simplex.Messaging.Agent.Store.Interface (closeDBStore, execSQL)
 import qualified Simplex.Messaging.Agent.Store.Migrations as Migrations
 import Simplex.Messaging.Agent.Store.Shared (UpMigration (..), upMigration)
 import Simplex.Messaging.Client (SMPClientError, ServerTransmission (..), ServerTransmissionBatch, temporaryClientError, unexpectedResponse)
@@ -279,7 +280,7 @@ disposeAgentClient c@AgentClient {acThread, agentEnv = Env {store}} = do
   t_ <- atomically (swapTVar acThread Nothing) $>>= (liftIO . deRefWeak)
   disconnectAgentClient c
   mapM_ killThread t_
-  liftIO $ closeStore store
+  liftIO $ closeDBStore store
 
 resumeAgentClient :: AgentClient -> IO ()
 resumeAgentClient c = atomically $ writeTVar (active c) True

src/Simplex/Messaging/Agent/Env/SQLite.hs

@@ -1,4 +1,3 @@
-{-# LANGUAGE CPP #-}
 {-# LANGUAGE ConstraintKinds #-}
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE DuplicateRecordFields #-}
@@ -70,6 +69,7 @@ import Simplex.Messaging.Agent.Protocol
 import Simplex.Messaging.Agent.RetryInterval
 import Simplex.Messaging.Agent.Store (createStore)
 import Simplex.Messaging.Agent.Store.Common (DBStore)
+import Simplex.Messaging.Agent.Store.Interface (DBOpts)
 import Simplex.Messaging.Agent.Store.Shared (MigrationConfirmation (..), MigrationError (..))
 import Simplex.Messaging.Client
 import qualified Simplex.Messaging.Crypto as C
@@ -87,11 +87,6 @@ import Simplex.Messaging.Util (allFinally, catchAllErrors, catchAllErrors', tryA
 import System.Mem.Weak (Weak)
 import System.Random (StdGen, newStdGen)
 import UnliftIO.STM
-#if defined(dbPostgres)
-import Database.PostgreSQL.Simple (ConnectInfo (..))
-#else
-import Data.ByteArray (ScrubbedBytes)
-#endif
 
 type AM' a = ReaderT Env IO a
 
@@ -277,13 +272,8 @@ newSMPAgentEnv config store = do
   multicastSubscribers <- newTMVarIO 0
   pure Env {config, store, random, randomServer, ntfSupervisor, xftpAgent, multicastSubscribers}
 
-#if defined(dbPostgres)
-createAgentStore :: ConnectInfo -> String -> MigrationConfirmation -> IO (Either MigrationError DBStore)
+createAgentStore :: DBOpts -> MigrationConfirmation -> IO (Either MigrationError DBStore)
 createAgentStore = createStore
-#else
-createAgentStore :: FilePath -> ScrubbedBytes -> Bool -> MigrationConfirmation -> Bool -> IO (Either MigrationError DBStore)
-createAgentStore = createStore
-#endif
 
 data NtfSupervisor = NtfSupervisor
   { ntfTkn :: TVar (Maybe NtfToken),

src/Simplex/Messaging/Agent/Store.hs

@@ -1,4 +1,3 @@
-{-# LANGUAGE CPP #-}
 {-# LANGUAGE ConstraintKinds #-}
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE DeriveAnyClass #-}
@@ -26,13 +25,12 @@ import Data.List (find)
 import Data.List.NonEmpty (NonEmpty)
 import qualified Data.List.NonEmpty as L
 import Data.Maybe (isJust)
-import Data.Text (Text)
 import Data.Time (UTCTime)
 import Data.Type.Equality
 import Simplex.Messaging.Agent.Protocol
 import Simplex.Messaging.Agent.RetryInterval (RI2State)
 import Simplex.Messaging.Agent.Store.Common
-import qualified Simplex.Messaging.Agent.Store.DB as DB
+import Simplex.Messaging.Agent.Store.Interface (DBOpts, createDBStore)
 import qualified Simplex.Messaging.Agent.Store.Migrations as Migrations
 import Simplex.Messaging.Agent.Store.Shared (MigrationConfirmation (..), MigrationError (..))
 import qualified Simplex.Messaging.Crypto as C
@@ -54,30 +52,9 @@ import Simplex.Messaging.Protocol
     VersionSMPC,
   )
 import qualified Simplex.Messaging.Protocol as SMP
-#if defined(dbPostgres)
-import Database.PostgreSQL.Simple (ConnectInfo (..))
-import qualified Simplex.Messaging.Agent.Store.Postgres as Store
-#else
-import Data.ByteArray (ScrubbedBytes)
-import qualified Simplex.Messaging.Agent.Store.SQLite as Store
-#endif
 
-#if defined(dbPostgres)
-createStore :: ConnectInfo -> String -> MigrationConfirmation -> IO (Either MigrationError DBStore)
-createStore connectInfo schema = Store.createDBStore connectInfo schema Migrations.app
-#else
-createStore :: FilePath -> ScrubbedBytes -> Bool -> MigrationConfirmation -> Bool -> IO (Either MigrationError DBStore)
-createStore dbFilePath dbKey keepKey = Store.createDBStore dbFilePath dbKey keepKey Migrations.app
-#endif
-
-closeStore :: DBStore -> IO ()
-closeStore = Store.closeDBStore
-
-reopenStore :: DBStore -> IO ()
-reopenStore = Store.reopenDBStore
-
-execSQL :: DB.Connection -> Text -> IO [Text]
-execSQL = Store.execSQL
+createStore :: DBOpts -> MigrationConfirmation -> IO (Either MigrationError DBStore)
+createStore dbOpts = createDBStore dbOpts Migrations.app
 
 -- * Queue types
 

src/Simplex/Messaging/Agent/Store/Interface.hs

@@ -0,0 +1,14 @@
+{-# LANGUAGE CPP #-}
+
+module Simplex.Messaging.Agent.Store.Interface
+#if defined(dbPostgres)
+  ( module Simplex.Messaging.Agent.Store.Postgres,
+  )
+  where
+import Simplex.Messaging.Agent.Store.Postgres
+#else
+  ( module Simplex.Messaging.Agent.Store.SQLite,
+  )
+  where
+import Simplex.Messaging.Agent.Store.SQLite
+#endif

src/Simplex/Messaging/Agent/Store/Postgres.hs

@@ -5,7 +5,8 @@
 {-# LANGUAGE ScopedTypeVariables #-}
 
 module Simplex.Messaging.Agent.Store.Postgres
-  ( createDBStore,
+  ( DBOpts (..),
+    createDBStore,
     closeDBStore,
     reopenDBStore,
     execSQL,
@@ -14,47 +15,49 @@ where
 
 import Control.Exception (throwIO)
 import Control.Monad (unless, void)
+import Data.ByteString (ByteString)
 import Data.Functor (($>))
 import Data.String (fromString)
 import Data.Text (Text)
-import Database.PostgreSQL.Simple (ConnectInfo (..), Only (..))
+import Database.PostgreSQL.Simple (Only (..))
 import qualified Database.PostgreSQL.Simple as PSQL
 import Database.PostgreSQL.Simple.SqlQQ (sql)
 import Simplex.Messaging.Agent.Store.Migrations (migrateSchema)
 import Simplex.Messaging.Agent.Store.Postgres.Common
 import qualified Simplex.Messaging.Agent.Store.Postgres.DB as DB
-import Simplex.Messaging.Agent.Store.Postgres.Util (createDBAndUserIfNotExists)
 import Simplex.Messaging.Agent.Store.Shared (Migration (..), MigrationConfirmation (..), MigrationError (..))
 import Simplex.Messaging.Util (ifM)
 import UnliftIO.Exception (bracketOnError, onException)
 import UnliftIO.MVar
 import UnliftIO.STM
 
--- | Create a new Postgres DBStore with the given connection info, schema name and migrations.
--- This function creates the user and/or database passed in connectInfo if they do not exist
--- (expects the default 'postgres' user and 'postgres' db to exist).
+data DBOpts = DBOpts
+  { connstr :: ByteString,
+    schema :: String
+  }
+
+-- | Create a new Postgres DBStore with the given connection string, schema name and migrations.
 -- If passed schema does not exist in connectInfo database, it will be created.
 -- Applies necessary migrations to schema.
 -- TODO [postgres] authentication / user password, db encryption (?)
-createDBStore :: ConnectInfo -> String -> [Migration] -> MigrationConfirmation -> IO (Either MigrationError DBStore)
-createDBStore connectInfo schema migrations confirmMigrations = do
-  createDBAndUserIfNotExists connectInfo
-  st <- connectPostgresStore connectInfo schema
+createDBStore :: DBOpts -> [Migration] -> MigrationConfirmation -> IO (Either MigrationError DBStore)
+createDBStore DBOpts {connstr, schema} migrations confirmMigrations = do
+  st <- connectPostgresStore connstr schema
   r <- migrateSchema st migrations confirmMigrations True `onException` closeDBStore st
   case r of
     Right () -> pure $ Right st
     Left e -> closeDBStore st $> Left e
 
-connectPostgresStore :: ConnectInfo -> String -> IO DBStore
-connectPostgresStore dbConnectInfo dbSchema = do
-  (dbConn, dbNew) <- connectDB dbConnectInfo dbSchema -- TODO [postgres] analogue for dbBusyLoop?
+connectPostgresStore :: ByteString -> String -> IO DBStore
+connectPostgresStore dbConnstr dbSchema = do
+  (dbConn, dbNew) <- connectDB dbConnstr dbSchema -- TODO [postgres] analogue for dbBusyLoop?
   dbConnection <- newMVar dbConn
   dbClosed <- newTVarIO False
-  pure DBStore {dbConnectInfo, dbSchema, dbConnection, dbNew, dbClosed}
+  pure DBStore {dbConnstr, dbSchema, dbConnection, dbNew, dbClosed}
 
-connectDB :: ConnectInfo -> String -> IO (DB.Connection, Bool)
-connectDB dbConnectInfo schema = do
-  db <- PSQL.connect dbConnectInfo
+connectDB :: ByteString -> String -> IO (DB.Connection, Bool)
+connectDB connstr schema = do
+  db <- PSQL.connectPostgreSQL connstr
   schemaExists <- prepare db `onException` PSQL.close db
   let dbNew = not schemaExists
   pure (db, dbNew)
@@ -84,12 +87,12 @@ closeDBStore st@DBStore {dbClosed} =
       atomically $ writeTVar dbClosed True
 
 openPostgresStore_ :: DBStore -> IO ()
-openPostgresStore_ DBStore {dbConnectInfo, dbSchema, dbConnection, dbClosed} =
+openPostgresStore_ DBStore {dbConnstr, dbSchema, dbConnection, dbClosed} =
   bracketOnError
     (takeMVar dbConnection)
     (tryPutMVar dbConnection)
     $ \_dbConn -> do
-      (dbConn, _dbNew) <- connectDB dbConnectInfo dbSchema
+      (dbConn, _dbNew) <- connectDB dbConnstr dbSchema
       atomically $ writeTVar dbClosed False
       putMVar dbConnection dbConn
 

src/Simplex/Messaging/Agent/Store/Postgres/Common.hs

@@ -10,13 +10,14 @@ module Simplex.Messaging.Agent.Store.Postgres.Common
   )
 where
 
+import Data.ByteString (ByteString)
 import qualified Database.PostgreSQL.Simple as PSQL
 import UnliftIO.MVar
 import UnliftIO.STM
 
 -- TODO [postgres] use log_min_duration_statement instead of custom slow queries (SQLite's Connection type)
 data DBStore = DBStore
-  { dbConnectInfo :: PSQL.ConnectInfo,
+  { dbConnstr :: ByteString,
     dbSchema :: String,
     dbConnection :: MVar PSQL.Connection,
     dbClosed :: TVar Bool,

src/Simplex/Messaging/Agent/Store/Postgres/Util.hs

@@ -5,7 +5,6 @@
 
 module Simplex.Messaging.Agent.Store.Postgres.Util
   ( createDBAndUserIfNotExists,
-    -- for tests
     dropSchema,
     dropAllSchemasExceptSystem,
     dropDatabaseAndUser,

src/Simplex/Messaging/Agent/Store/SQLite.hs

@@ -25,7 +25,8 @@
 {-# OPTIONS_GHC -fno-warn-orphans #-}
 
 module Simplex.Messaging.Agent.Store.SQLite
-  ( createDBStore,
+  ( DBOpts (..),
+    createDBStore,
     closeDBStore,
     reopenDBStore,
     execSQL,
@@ -64,8 +65,15 @@ import UnliftIO.STM
 
 -- * SQLite Store implementation
 
-createDBStore :: FilePath -> ScrubbedBytes -> Bool -> [Migration] -> MigrationConfirmation -> Bool -> IO (Either MigrationError DBStore)
-createDBStore dbFilePath dbKey keepKey migrations confirmMigrations vacuum = do
+data DBOpts = DBOpts
+  { dbFilePath :: FilePath,
+    dbKey :: ScrubbedBytes,
+    keepKey :: Bool,
+    vacuum :: Bool
+  }
+
+createDBStore :: DBOpts -> [Migration] -> MigrationConfirmation -> IO (Either MigrationError DBStore)
+createDBStore DBOpts {dbFilePath, dbKey, keepKey, vacuum} migrations confirmMigrations = do
   let dbDir = takeDirectory dbFilePath
   createDirectoryIfMissing True dbDir
   st <- connectSQLiteStore dbFilePath dbKey keepKey