From e4b47825b56122222e5bf4716285b419acdac83d Mon Sep 17 00:00:00 2001
From: Evgeny Poberezkin <2769109+epoberezkin@users.noreply.github.com>
Date: Fri, 2 Sep 2022 15:42:37 +0100
Subject: [PATCH] functions to create and close store (#512)

---
 src/Simplex/Messaging/Agent/Env/SQLite.hs   |  6 ++-
 src/Simplex/Messaging/Agent/Store/SQLite.hs | 41 ++++++++++-----------
 2 files changed, 24 insertions(+), 23 deletions(-)

diff --git a/src/Simplex/Messaging/Agent/Env/SQLite.hs b/src/Simplex/Messaging/Agent/Env/SQLite.hs
index 03cb3ba7b..9004a2bd9 100644
--- a/src/Simplex/Messaging/Agent/Env/SQLite.hs
+++ b/src/Simplex/Messaging/Agent/Env/SQLite.hs
@@ -18,6 +18,7 @@ module Simplex.Messaging.Agent.Env.SQLite
     defaultReconnectInterval,
     Env (..),
     newSMPAgentEnv,
+    createAgentStore,
     NtfSupervisor (..),
     NtfSupervisorCommand (..),
   )
@@ -143,12 +144,15 @@ data Env = Env
 newSMPAgentEnv :: (MonadUnliftIO m, MonadRandom m) => AgentConfig -> m Env
 newSMPAgentEnv config@AgentConfig {dbFile, dbKey, yesToMigrations} = do
   idsDrg <- newTVarIO =<< drgNew
-  store <- liftIO $ createSQLiteStore dbFile dbKey Migrations.app yesToMigrations
+  store <- liftIO $ createAgentStore dbFile dbKey yesToMigrations
   clientCounter <- newTVarIO 0
   randomServer <- newTVarIO =<< liftIO newStdGen
   ntfSupervisor <- atomically . newNtfSubSupervisor $ tbqSize config
   return Env {config, store, idsDrg, clientCounter, randomServer, ntfSupervisor}
 
+createAgentStore :: FilePath -> String -> Bool -> IO SQLiteStore
+createAgentStore dbFilePath dbKey = createSQLiteStore dbFilePath dbKey Migrations.app
+
 data NtfSupervisor = NtfSupervisor
   { ntfTkn :: TVar (Maybe NtfToken),
     ntfSubQ :: TBQueue (ConnId, NtfSupervisorCommand),
diff --git a/src/Simplex/Messaging/Agent/Store/SQLite.hs b/src/Simplex/Messaging/Agent/Store/SQLite.hs
index d8af7226b..ba62da04d 100644
--- a/src/Simplex/Messaging/Agent/Store/SQLite.hs
+++ b/src/Simplex/Messaging/Agent/Store/SQLite.hs
@@ -22,6 +22,7 @@ module Simplex.Messaging.Agent.Store.SQLite
   ( SQLiteStore (..),
     createSQLiteStore,
     connectSQLiteStore,
+    closeSQLiteStore,
     sqlString,
 
     -- * Queues and connections
@@ -107,7 +108,7 @@ import Data.ByteString (ByteString)
 import qualified Data.ByteString.Base64.URL as U
 import Data.Char (toLower)
 import Data.Functor (($>))
-import Data.List (find, foldl')
+import Data.List (foldl')
 import Data.List.NonEmpty (NonEmpty (..))
 import qualified Data.Map.Strict as M
 import Data.Maybe (fromMaybe, listToMaybe)
@@ -141,7 +142,7 @@ import System.Directory (copyFile, createDirectoryIfMissing, doesFileExist)
 import System.Exit (exitFailure)
 import System.FilePath (takeDirectory)
 import System.IO (hFlush, stdout)
-import UnliftIO.Exception (bracket)
+import UnliftIO.Exception (bracket, onException)
 import qualified UnliftIO.Exception as E
 import UnliftIO.STM
 
@@ -159,19 +160,9 @@ createSQLiteStore dbFilePath dbKey migrations yesToMigrations = do
   let dbDir = takeDirectory dbFilePath
   createDirectoryIfMissing False dbDir
   st <- connectSQLiteStore dbFilePath dbKey
-  checkThreadsafe st
-  migrateSchema st migrations yesToMigrations
+  migrateSchema st migrations yesToMigrations `onException` closeSQLiteStore st
   pure st
 
-checkThreadsafe :: SQLiteStore -> IO ()
-checkThreadsafe st = withConnection st $ \db -> do
-  compileOptions <- DB.query_ db "pragma COMPILE_OPTIONS;" :: IO [[Text]]
-  let threadsafeOption = find (T.isPrefixOf "THREADSAFE=") (concat compileOptions)
-  case threadsafeOption of
-    Just "THREADSAFE=0" -> confirmOrExit "SQLite compiled with non-threadsafe code."
-    Nothing -> putStrLn "Warning: SQLite THREADSAFE compile option not found"
-    _ -> return ()
-
 migrateSchema :: SQLiteStore -> [Migration] -> Bool -> IO ()
 migrateSchema st migrations yesToMigrations = withConnection st $ \db -> do
   Migrations.initialize db
@@ -203,17 +194,23 @@ connectSQLiteStore dbFilePath dbKey = do
 connectDB :: FilePath -> String -> IO DB.Connection
 connectDB path key = do
   db <- DB.open path
-  let exec = SQLite3.exec $ DB.connectionHandle db
-  unless (null key) . exec $ "PRAGMA key = " <> sqlString key <> ";"
-  exec . fromQuery $
-    [sql|
-      PRAGMA foreign_keys = ON;
-      -- PRAGMA trusted_schema = OFF;
-      PRAGMA secure_delete = ON;
-      PRAGMA auto_vacuum = FULL;
-    |]
+  prepare db `onException` DB.close db
   -- _printPragmas db path
   pure db
+  where
+    prepare db = do
+      let exec = SQLite3.exec $ DB.connectionHandle db
+      unless (null key) . exec $ "PRAGMA key = " <> sqlString key <> ";"
+      exec . fromQuery $
+        [sql|
+          PRAGMA foreign_keys = ON;
+          -- PRAGMA trusted_schema = OFF;
+          PRAGMA secure_delete = ON;
+          PRAGMA auto_vacuum = FULL;
+        |]
+
+closeSQLiteStore :: SQLiteStore -> IO ()
+closeSQLiteStore st = atomically (takeTMVar $ dbConnection st) >>= DB.close
 
 sqlString :: String -> Text
 sqlString s = quote <> T.replace quote "''" (T.pack s) <> quote