agent: refactor migrations (#1452)

* agent: refactor migrations

* fix postgres

src/Simplex/Messaging/Agent.hs

@@ -170,8 +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.Interface (closeDBStore, execSQL, getCurrentMigrations)
 import Simplex.Messaging.Agent.Store.Shared (UpMigration (..), upMigration)
 import Simplex.Messaging.Client (SMPClientError, ServerTransmission (..), ServerTransmissionBatch, temporaryClientError, unexpectedResponse)
 import qualified Simplex.Messaging.Crypto as C
@@ -2180,7 +2179,7 @@ execAgentStoreSQL :: AgentClient -> Text -> AE [Text]
 execAgentStoreSQL c sql = withAgentEnv c $ withStore' c (`execSQL` sql)
 
 getAgentMigrations :: AgentClient -> AE [UpMigration]
-getAgentMigrations c = withAgentEnv c $ map upMigration <$> withStore' c Migrations.getCurrent
+getAgentMigrations c = withAgentEnv c $ map upMigration <$> withStore' c getCurrentMigrations
 
 debugAgentLocks :: AgentClient -> IO AgentLocks
 debugAgentLocks AgentClient {connLocks = cs, invLocks = is, deleteLock = d} = do

src/Simplex/Messaging/Agent/Store.hs

@@ -30,8 +30,7 @@ import Data.Type.Equality
 import Simplex.Messaging.Agent.Protocol
 import Simplex.Messaging.Agent.RetryInterval (RI2State)
 import Simplex.Messaging.Agent.Store.Common
-import Simplex.Messaging.Agent.Store.Interface (DBOpts, createDBStore)
-import qualified Simplex.Messaging.Agent.Store.Migrations as Migrations
+import Simplex.Messaging.Agent.Store.Interface (DBOpts, appMigrations, createDBStore)
 import Simplex.Messaging.Agent.Store.Shared (MigrationConfirmation (..), MigrationError (..))
 import qualified Simplex.Messaging.Crypto as C
 import Simplex.Messaging.Crypto.Ratchet (PQEncryption, PQSupport, RatchetX448)
@@ -54,7 +53,7 @@ import Simplex.Messaging.Protocol
 import qualified Simplex.Messaging.Protocol as SMP
 
 createStore :: DBOpts -> MigrationConfirmation -> IO (Either MigrationError DBStore)
-createStore dbOpts = createDBStore dbOpts Migrations.app
+createStore dbOpts = createDBStore dbOpts appMigrations
 
 -- * Queue types
 

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

@@ -1,16 +1,11 @@
-{-# LANGUAGE CPP #-}
 {-# LANGUAGE LambdaCase #-}
 
 module Simplex.Messaging.Agent.Store.Migrations
   ( Migration (..),
     MigrationsToRun (..),
     DownMigration (..),
-    Migrations.app,
-    Migrations.getCurrent,
-    get,
-    Migrations.initialize,
-    Migrations.run,
-    migrateSchema,
+    DBMigrate (..),
+    sharedMigrateSchema,
     -- for tests
     migrationsToRun,
     toDownMigration,
@@ -21,19 +16,9 @@ import Control.Monad
 import Data.Char (toLower)
 import Data.Functor (($>))
 import Data.Maybe (isNothing, mapMaybe)
-import Simplex.Messaging.Agent.Store.Common
 import Simplex.Messaging.Agent.Store.Shared
 import System.Exit (exitFailure)
 import System.IO (hFlush, stdout)
-#if defined(dbPostgres)
-import qualified Simplex.Messaging.Agent.Store.Postgres.Migrations as Migrations
-#else
-import qualified Simplex.Messaging.Agent.Store.SQLite.Migrations as Migrations
-import System.Directory (copyFile)
-#endif
-
-get :: DBStore -> [Migration] -> IO (Either MTRError MigrationsToRun)
-get st migrations = migrationsToRun migrations <$> withTransaction st Migrations.getCurrent
 
 migrationsToRun :: [Migration] -> [Migration] -> Either MTRError MigrationsToRun
 migrationsToRun [] [] = Right MTRNone
@@ -48,44 +33,42 @@ migrationsToRun (a : as) (d : ds)
   | name a == name d = migrationsToRun as ds
   | otherwise = Left $ MTREDifferent (name a) (name d)
 
-migrateSchema :: DBStore -> [Migration] -> MigrationConfirmation -> Bool -> IO (Either MigrationError ())
-migrateSchema st migrations confirmMigrations vacuum = do
-  Migrations.initialize st
-  get st migrations >>= \case
+data DBMigrate = DBMigrate
+  { initialize :: IO (),
+    getCurrent :: IO [Migration],
+    run :: MigrationsToRun -> IO (),
+    backup :: IO ()
+  }
+
+sharedMigrateSchema :: DBMigrate -> Bool -> [Migration] -> MigrationConfirmation -> IO (Either MigrationError ())
+sharedMigrateSchema dbm dbNew' migrations confirmMigrations = do
+  initialize dbm
+  currentMs <- getCurrent dbm
+  case migrationsToRun migrations currentMs of
     Left e -> do
       when (confirmMigrations == MCConsole) $ confirmOrExit ("Database state error: " <> mtrErrorDescription e)
       pure . Left $ MigrationError e
     Right MTRNone -> pure $ Right ()
     Right ms@(MTRUp ums)
-      | dbNew st -> Migrations.run st vacuum ms $> Right ()
+      | dbNew' -> run dbm ms $> Right ()
       | otherwise -> case confirmMigrations of
-          MCYesUp -> runWithBackup st vacuum ms
-          MCYesUpDown -> runWithBackup st vacuum ms
-          MCConsole -> confirm err >> runWithBackup st vacuum ms
+          MCYesUp -> runWithBackup ms
+          MCYesUpDown -> runWithBackup ms
+          MCConsole -> confirm err >> runWithBackup ms
           MCError -> pure $ Left err
       where
         err = MEUpgrade $ map upMigration ums -- "The app has a newer version than the database.\nConfirm to back up and upgrade using these migrations: " <> intercalate ", " (map name ums)
     Right ms@(MTRDown dms) -> case confirmMigrations of
-      MCYesUpDown -> runWithBackup st vacuum ms
-      MCConsole -> confirm err >> runWithBackup st vacuum ms
+      MCYesUpDown -> runWithBackup ms
+      MCConsole -> confirm err >> runWithBackup ms
       MCYesUp -> pure $ Left err
       MCError -> pure $ Left err
       where
         err = MEDowngrade $ map downName dms
   where
+    runWithBackup ms = backup dbm >> run dbm ms $> Right ()
     confirm err = confirmOrExit $ migrationErrorDescription err
 
-runWithBackup :: DBStore -> Bool -> MigrationsToRun -> IO (Either a ())
-#if defined(dbPostgres)
-runWithBackup st vacuum ms = Migrations.run st vacuum ms $> Right ()
-#else
-runWithBackup st vacuum ms = do
-  let f = dbFilePath st
-  copyFile f (f <> ".bak")
-  Migrations.run st vacuum ms
-  pure $ Right ()
-#endif
-
 confirmOrExit :: String -> IO ()
 confirmOrExit s = do
   putStrLn s

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

@@ -6,6 +6,8 @@
 
 module Simplex.Messaging.Agent.Store.Postgres
   ( DBOpts (..),
+    Migrations.appMigrations,
+    Migrations.getCurrentMigrations,
     createDBStore,
     closeDBStore,
     reopenDBStore,
@@ -22,7 +24,8 @@ import Data.Text (Text)
 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.Migrations (DBMigrate (..), sharedMigrateSchema)
+import qualified Simplex.Messaging.Agent.Store.Postgres.Migrations as Migrations
 import Simplex.Messaging.Agent.Store.Postgres.Common
 import qualified Simplex.Messaging.Agent.Store.Postgres.DB as DB
 import Simplex.Messaging.Agent.Store.Shared (Migration (..), MigrationConfirmation (..), MigrationError (..))
@@ -43,10 +46,16 @@ data DBOpts = DBOpts
 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
+  r <- migrateSchema st `onException` closeDBStore st
   case r of
     Right () -> pure $ Right st
     Left e -> closeDBStore st $> Left e
+  where
+    migrateSchema st =
+      let initialize = Migrations.initialize st
+          getCurrent = withTransaction st Migrations.getCurrentMigrations
+          dbm = DBMigrate {initialize, getCurrent, run = Migrations.run st, backup = pure ()}
+       in sharedMigrateSchema dbm (dbNew st) migrations confirmMigrations
 
 connectPostgresStore :: ByteString -> String -> IO DBStore
 connectPostgresStore dbConnstr dbSchema = do

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

@@ -17,7 +17,7 @@ module Simplex.Messaging.Agent.Store.Postgres.DB
 where
 
 import Control.Monad (void)
-import Data.Int (Int32, Int64)
+import Data.Int (Int64)
 import Data.Word (Word16, Word32)
 import Database.PostgreSQL.Simple (ResultError (..))
 import qualified Database.PostgreSQL.Simple as PSQL

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

@@ -5,10 +5,10 @@
 {-# LANGUAGE TupleSections #-}
 
 module Simplex.Messaging.Agent.Store.Postgres.Migrations
-  ( app,
+  ( appMigrations,
     initialize,
     run,
-    getCurrent,
+    getCurrentMigrations,
   )
 where
 
@@ -34,8 +34,8 @@ schemaMigrations =
   ]
 
 -- | The list of migrations in ascending order by date
-app :: [Migration]
-app = sortOn name $ map migration schemaMigrations
+appMigrations :: [Migration]
+appMigrations = sortOn name $ map migration schemaMigrations
   where
     migration (name, up, down) = Migration {name, up, down = down}
 
@@ -53,8 +53,8 @@ initialize st = withTransaction' st $ \db ->
       )
     |]
 
-run :: DBStore -> Bool -> MigrationsToRun -> IO ()
-run st _vacuum = \case
+run :: DBStore -> MigrationsToRun -> IO ()
+run st = \case
   MTRUp [] -> pure ()
   MTRUp ms -> mapM_ runUp ms
   MTRDown ms -> mapM_ runDown $ reverse ms
@@ -72,7 +72,7 @@ run st _vacuum = \case
       withMVar (connectionHandle db) $ \pqConn ->
         void $ LibPQ.exec pqConn (TE.encodeUtf8 query)
 
-getCurrent :: PSQL.Connection -> IO [Migration]
-getCurrent db = map toMigration <$> PSQL.query_ db "SELECT name, down FROM migrations ORDER BY name ASC;"
+getCurrentMigrations :: PSQL.Connection -> IO [Migration]
+getCurrentMigrations db = map toMigration <$> PSQL.query_ db "SELECT name, down FROM migrations ORDER BY name ASC;"
   where
     toMigration (name, down) = Migration {name, up = T.pack "", down}

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

@@ -26,6 +26,8 @@
 
 module Simplex.Messaging.Agent.Store.SQLite
   ( DBOpts (..),
+    Migrations.appMigrations,
+    Migrations.getCurrentMigrations,
     createDBStore,
     closeDBStore,
     reopenDBStore,
@@ -52,12 +54,13 @@ import Database.SQLite.Simple (Query (..))
 import qualified Database.SQLite.Simple as SQL
 import Database.SQLite.Simple.QQ (sql)
 import qualified Database.SQLite3 as SQLite3
-import Simplex.Messaging.Agent.Store.Migrations (migrateSchema)
+import Simplex.Messaging.Agent.Store.Migrations (DBMigrate (..), sharedMigrateSchema)
+import qualified Simplex.Messaging.Agent.Store.SQLite.Migrations as Migrations
 import Simplex.Messaging.Agent.Store.SQLite.Common
 import qualified Simplex.Messaging.Agent.Store.SQLite.DB as DB
 import Simplex.Messaging.Agent.Store.Shared (Migration (..), MigrationConfirmation (..), MigrationError (..))
 import Simplex.Messaging.Util (ifM, safeDecodeUtf8)
-import System.Directory (createDirectoryIfMissing, doesFileExist)
+import System.Directory (copyFile, createDirectoryIfMissing, doesFileExist)
 import System.FilePath (takeDirectory)
 import UnliftIO.Exception (bracketOnError, onException)
 import UnliftIO.MVar
@@ -78,10 +81,18 @@ createDBStore DBOpts {dbFilePath, dbKey, keepKey, track, vacuum} migrations conf
   let dbDir = takeDirectory dbFilePath
   createDirectoryIfMissing True dbDir
   st <- connectSQLiteStore dbFilePath dbKey keepKey track
-  r <- migrateSchema st migrations confirmMigrations vacuum `onException` closeDBStore st
+  r <- migrateSchema st `onException` closeDBStore st
   case r of
     Right () -> pure $ Right st
     Left e -> closeDBStore st $> Left e
+  where
+    migrateSchema st =
+      let initialize = Migrations.initialize st
+          getCurrent = withTransaction st Migrations.getCurrentMigrations
+          run = Migrations.run st vacuum
+          backup = copyFile dbFilePath (dbFilePath <> ".bak")
+          dbm = DBMigrate {initialize, getCurrent, run, backup}
+       in sharedMigrateSchema dbm (dbNew st) migrations confirmMigrations
 
 connectSQLiteStore :: FilePath -> ScrubbedBytes -> Bool -> DB.TrackQueries -> IO DBStore
 connectSQLiteStore dbFilePath key keepKey track = do

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

@@ -8,10 +8,10 @@
 {-# LANGUAGE TupleSections #-}
 
 module Simplex.Messaging.Agent.Store.SQLite.Migrations
-  ( app,
+  ( appMigrations,
     initialize,
     run,
-    getCurrent,
+    getCurrentMigrations,
   )
 where
 
@@ -112,13 +112,13 @@ schemaMigrations =
   ]
 
 -- | The list of migrations in ascending order by date
-app :: [Migration]
-app = sortOn name $ map migration schemaMigrations
+appMigrations :: [Migration]
+appMigrations = sortOn name $ map migration schemaMigrations
   where
     migration (name, up, down) = Migration {name, up = fromQuery up, down = fromQuery <$> down}
 
-getCurrent :: DB.Connection -> IO [Migration]
-getCurrent DB.Connection {DB.conn} = map toMigration <$> SQL.query_ conn "SELECT name, down FROM migrations ORDER BY name ASC;"
+getCurrentMigrations :: DB.Connection -> IO [Migration]
+getCurrentMigrations DB.Connection {DB.conn} = map toMigration <$> SQL.query_ conn "SELECT name, down FROM migrations ORDER BY name ASC;"
   where
     toMigration (name, down) = Migration {name, up = "", down}
 

tests/AgentTests/SQLiteTests.hs

@@ -45,6 +45,7 @@ import Simplex.Messaging.Agent.Store.AgentStore
 import Simplex.Messaging.Agent.Store.SQLite
 import Simplex.Messaging.Agent.Store.SQLite.Common (DBStore (..), withTransaction')
 import qualified Simplex.Messaging.Agent.Store.SQLite.DB as DB
+import Simplex.Messaging.Agent.Store.SQLite.Migrations (appMigrations, getCurrentMigrations)
 import qualified Simplex.Messaging.Agent.Store.SQLite.Migrations as Migrations
 import Simplex.Messaging.Agent.Store.Shared (MigrationConfirmation (..))
 import qualified Simplex.Messaging.Crypto as C
@@ -81,7 +82,7 @@ createEncryptedStore key keepKey = 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
-  Right st <- createDBStore (DBOpts (testDB <> show r) key keepKey True DB.TQOff) Migrations.app MCError
+  Right st <- createDBStore (DBOpts (testDB <> show r) key keepKey True DB.TQOff) appMigrations MCError
   withTransaction' st (`SQL.execute_` "INSERT INTO users (user_id) VALUES (1);")
   pure st
 
@@ -635,7 +636,7 @@ testReopenEncryptedStoreKeepKey = do
   hasMigrations st
 
 getMigrations :: DBStore -> IO Bool
-getMigrations st = not . null <$> withTransaction st Migrations.getCurrent
+getMigrations st = not . null <$> withTransaction st getCurrentMigrations
 
 hasMigrations :: DBStore -> Expectation
 hasMigrations st = getMigrations st `shouldReturn` True

tests/AgentTests/SchemaDump.hs

@@ -50,7 +50,7 @@ testVerifySchemaDump :: IO ()
 testVerifySchemaDump = do
   savedSchema <- ifM (doesFileExist appSchema) (readFile appSchema) (pure "")
   savedSchema `deepseq` pure ()
-  void $ createDBStore (DBOpts testDB "" False True TQOff) Migrations.app MCConsole
+  void $ createDBStore (DBOpts testDB "" False True TQOff) appMigrations MCConsole
   getSchema testDB appSchema `shouldReturn` savedSchema
   removeFile testDB
 
@@ -58,7 +58,7 @@ testVerifyLintFKeyIndexes :: IO ()
 testVerifyLintFKeyIndexes = do
   savedLint <- ifM (doesFileExist appLint) (readFile appLint) (pure "")
   savedLint `deepseq` pure ()
-  void $ createDBStore (DBOpts testDB "" False True TQOff) Migrations.app MCConsole
+  void $ createDBStore (DBOpts testDB "" False True TQOff) appMigrations MCConsole
   getLintFKeyIndexes testDB "tests/tmp/agent_lint.sql" `shouldReturn` savedLint
   removeFile testDB
 
@@ -70,9 +70,9 @@ withTmpFiles =
 
 testSchemaMigrations :: IO ()
 testSchemaMigrations = do
-  let noDownMigrations = dropWhileEnd (\Migration {down} -> isJust down) Migrations.app
+  let noDownMigrations = dropWhileEnd (\Migration {down} -> isJust down) appMigrations
   Right st <- createDBStore (DBOpts testDB "" False True TQOff) noDownMigrations MCError
-  mapM_ (testDownMigration st) $ drop (length noDownMigrations) Migrations.app
+  mapM_ (testDownMigration st) $ drop (length noDownMigrations) appMigrations
   closeDBStore st
   removeFile testDB
   removeFile testSchema
@@ -94,19 +94,19 @@ testSchemaMigrations = do
 
 testUsersMigrationNew :: IO ()
 testUsersMigrationNew = do
-  Right st <- createDBStore (DBOpts testDB "" False True TQOff) Migrations.app MCError
+  Right st <- createDBStore (DBOpts testDB "" False True TQOff) appMigrations MCError
   withTransaction' st (`SQL.query_` "SELECT user_id FROM users;")
     `shouldReturn` ([] :: [Only Int])
   closeDBStore st
 
 testUsersMigrationOld :: IO ()
 testUsersMigrationOld = do
-  let beforeUsers = takeWhile (("m20230110_users" /=) . name) Migrations.app
+  let beforeUsers = takeWhile (("m20230110_users" /=) . name) appMigrations
   Right st <- createDBStore (DBOpts testDB "" False True TQOff) beforeUsers MCError
   withTransaction' st (`SQL.query_` "SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'users';")
     `shouldReturn` ([] :: [Only String])
   closeDBStore st
-  Right st' <- createDBStore (DBOpts testDB "" False True TQOff) Migrations.app MCYesUp
+  Right st' <- createDBStore (DBOpts testDB "" False True TQOff) appMigrations MCYesUp
   withTransaction' st' (`SQL.query_` "SELECT user_id FROM users;")
     `shouldReturn` ([Only (1 :: Int)])
   closeDBStore st'