agent: option to add SQLite functions to DB connection

simplexmq.cabal

@@ -216,6 +216,7 @@ library
           Simplex.Messaging.Agent.Store.SQLite.Migrations.M20250702_conn_invitations_remove_cascade_delete
           Simplex.Messaging.Agent.Store.SQLite.Migrations.M20251009_queue_to_subscribe
           Simplex.Messaging.Agent.Store.SQLite.Migrations.M20251010_client_notices
+          Simplex.Messaging.Agent.Store.SQLite.Util
   if flag(client_postgres) || flag(server_postgres)
       exposed-modules:
           Simplex.Messaging.Agent.Store.Postgres

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

@@ -42,6 +42,9 @@ module Simplex.Messaging.Agent.Store.SQLite
   )
 where
 
+import Control.Concurrent.MVar
+import Control.Concurrent.STM
+import Control.Exception (bracketOnError, onException, throwIO)
 import Control.Monad
 import Data.ByteArray (ScrubbedBytes)
 import qualified Data.ByteArray as BA
@@ -58,21 +61,19 @@ import Simplex.Messaging.Agent.Store.Migrations (DBMigrate (..), sharedMigrateSc
 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.SQLite.Util
 import Simplex.Messaging.Agent.Store.Shared (Migration (..), MigrationConfig (..), MigrationError (..))
 import Simplex.Messaging.Util (ifM, safeDecodeUtf8)
 import System.Directory (copyFile, createDirectoryIfMissing, doesFileExist)
 import System.FilePath (takeDirectory, takeFileName, (</>))
-import UnliftIO.Exception (bracketOnError, onException)
-import UnliftIO.MVar
-import UnliftIO.STM
 
 -- * SQLite Store implementation
 
 createDBStore :: DBOpts -> [Migration] -> MigrationConfig -> IO (Either MigrationError DBStore)
-createDBStore opts@DBOpts {dbFilePath, dbKey, keepKey, track} migrations migrationConfig = do
+createDBStore opts@DBOpts {dbFilePath} migrations migrationConfig = do
   let dbDir = takeDirectory dbFilePath
   createDirectoryIfMissing True dbDir
-  st <- connectSQLiteStore dbFilePath dbKey keepKey track
+  st <- connectSQLiteStore opts
   r <- migrateDBSchema st opts Nothing migrations migrationConfig `onException` closeDBStore st
   case r of
     Right () -> pure $ Right st
@@ -91,34 +92,37 @@ migrateDBSchema st DBOpts {dbFilePath, vacuum} migrationsTable migrations Migrat
       dbm = DBMigrate {initialize, getCurrent, run, backup}
    in sharedMigrateSchema dbm (dbNew st) migrations confirm
 
-connectSQLiteStore :: FilePath -> ScrubbedBytes -> Bool -> DB.TrackQueries -> IO DBStore
-connectSQLiteStore dbFilePath key keepKey track = do
+connectSQLiteStore :: DBOpts -> IO DBStore
+connectSQLiteStore DBOpts {dbFilePath, dbFunctions, dbKey = key, keepKey, track} = do
   dbNew <- not <$> doesFileExist dbFilePath
-  dbConn <- dbBusyLoop (connectDB dbFilePath key track)
+  dbConn <- dbBusyLoop $ connectDB dbFilePath dbFunctions key track
   dbConnection <- newMVar dbConn
   dbKey <- newTVarIO $! storeKey key keepKey
   dbClosed <- newTVarIO False
   dbSem <- newTVarIO 0
-  pure DBStore {dbFilePath, dbKey, dbSem, dbConnection, dbNew, dbClosed}
+  pure DBStore {dbFilePath, dbFunctions, dbKey, dbSem, dbConnection, dbNew, dbClosed}
 
-connectDB :: FilePath -> ScrubbedBytes -> DB.TrackQueries -> IO DB.Connection
-connectDB path key track = do
+connectDB :: FilePath -> [SQLiteFuncDef] -> ScrubbedBytes -> DB.TrackQueries -> IO DB.Connection
+connectDB path functions key track = do
   db <- DB.open path track
   prepare db `onException` DB.close db
   -- _printPragmas db path
   pure db
   where
     prepare db = do
-      let exec = SQLite3.exec $ SQL.connectionHandle $ DB.conn db
-      unless (BA.null key) . exec $ "PRAGMA key = " <> keyString key <> ";"
-      exec . fromQuery $
+      let db' = SQL.connectionHandle $ DB.conn db
+      unless (BA.null key) . SQLite3.exec db' $ "PRAGMA key = " <> keyString key <> ";"
+      SQLite3.exec db' . fromQuery $
         [sql|
           PRAGMA busy_timeout = 100;
           PRAGMA foreign_keys = ON;
           -- PRAGMA trusted_schema = OFF;
           PRAGMA secure_delete = ON;
           PRAGMA auto_vacuum = FULL;
         |]
+      forM_ functions $ \SQLiteFuncDef {funcName, argCount, deterministic, funcPtr} ->
+        createStaticFunction db' funcName argCount deterministic funcPtr
+          >>= either (throwIO . userError . show) pure
 
 closeDBStore :: DBStore -> IO ()
 closeDBStore st@DBStore {dbClosed} =
@@ -132,12 +136,12 @@ openSQLiteStore st@DBStore {dbClosed} key keepKey =
   ifM (readTVarIO dbClosed) (openSQLiteStore_ st key keepKey) (putStrLn "openSQLiteStore: already opened")
 
 openSQLiteStore_ :: DBStore -> ScrubbedBytes -> Bool -> IO ()
-openSQLiteStore_ DBStore {dbConnection, dbFilePath, dbKey, dbClosed} key keepKey =
+openSQLiteStore_ DBStore {dbConnection, dbFilePath, dbFunctions, dbKey, dbClosed} key keepKey =
   bracketOnError
     (takeMVar dbConnection)
     (tryPutMVar dbConnection)
     $ \DB.Connection {slow, track} -> do
-      DB.Connection {conn} <- connectDB dbFilePath key track
+      DB.Connection {conn} <- connectDB dbFilePath dbFunctions key track
       atomically $ do
         writeTVar dbClosed False
         writeTVar dbKey $! storeKey key keepKey

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

@@ -6,6 +6,7 @@
 module Simplex.Messaging.Agent.Store.SQLite.Common
   ( DBStore (..),
     DBOpts (..),
+    SQLiteFuncDef (..),
     withConnection,
     withConnection',
     withTransaction,
@@ -20,9 +21,13 @@ import Control.Concurrent (threadDelay)
 import Control.Concurrent.STM (retry)
 import Data.ByteArray (ScrubbedBytes)
 import qualified Data.ByteArray as BA
+import Data.ByteString (ByteString)
 import Database.SQLite.Simple (SQLError)
 import qualified Database.SQLite.Simple as SQL
+import Database.SQLite3.Bindings
+import Foreign.Ptr
 import qualified Simplex.Messaging.Agent.Store.SQLite.DB as DB
+import Simplex.Messaging.Agent.Store.SQLite.Util
 import Simplex.Messaging.Util (ifM, unlessM)
 import qualified UnliftIO.Exception as E
 import UnliftIO.MVar
@@ -33,6 +38,7 @@ storeKey key keepKey = if keepKey || BA.null key then Just key else Nothing
 
 data DBStore = DBStore
   { dbFilePath :: FilePath,
+    dbFunctions :: [SQLiteFuncDef],
     dbKey :: TVar (Maybe ScrubbedBytes),
     dbSem :: TVar Int,
     dbConnection :: MVar DB.Connection,
@@ -42,12 +48,21 @@ data DBStore = DBStore
 
 data DBOpts = DBOpts
   { dbFilePath :: FilePath,
+    dbFunctions :: [SQLiteFuncDef],
     dbKey :: ScrubbedBytes,
     keepKey :: Bool,
     vacuum :: Bool,
     track :: DB.TrackQueries
   }
 
+-- e.g. `SQLiteFuncDef "name" 2 True f`
+data SQLiteFuncDef = SQLiteFuncDef
+  { funcName :: ByteString,
+    argCount :: CArgCount,
+    deterministic :: Bool,
+    funcPtr :: FunPtr SQLiteFunc
+  }
+
 withConnectionPriority :: DBStore -> Bool -> (DB.Connection -> IO a) -> IO a
 withConnectionPriority DBStore {dbSem, dbConnection} priority action
   | priority = E.bracket_ signal release $ withMVar dbConnection action

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

@@ -0,0 +1,41 @@
+module Simplex.Messaging.Agent.Store.SQLite.Util where
+
+import Control.Exception (SomeException, catch, mask_)
+import Data.ByteString (ByteString)
+import qualified Data.ByteString as B
+import Database.SQLite3.Direct (Database (..), FuncArgs (..), FuncContext (..))
+import Database.SQLite3.Bindings
+import Foreign.C.String
+import Foreign.Ptr
+import Foreign.StablePtr
+
+data CFuncPtrs = CFuncPtrs (FunPtr CFunc) (FunPtr CFunc) (FunPtr CFuncFinal)
+
+type SQLiteFunc = Ptr CContext -> CArgCount -> Ptr (Ptr CValue) -> IO ()
+
+mkSQLiteFunc :: (FuncContext -> FuncArgs -> IO ()) -> SQLiteFunc
+mkSQLiteFunc f cxt nArgs cvals = catchAsResultError cxt $ f (FuncContext cxt) (FuncArgs nArgs cvals)
+{-# INLINE mkSQLiteFunc #-}
+
+-- Based on createFunction from Database.SQLite3.Direct, but uses static function pointer to avoid dynamic wrapper that triggers DCL.
+createStaticFunction :: Database -> ByteString -> CArgCount -> Bool -> FunPtr SQLiteFunc -> IO (Either Error ())
+createStaticFunction (Database db) name nArgs isDet funPtr = mask_ $ do
+  u <- newStablePtr $ CFuncPtrs funPtr nullFunPtr nullFunPtr
+  let flags = if isDet then c_SQLITE_DETERMINISTIC else 0
+  B.useAsCString name $ \namePtr ->
+    toResult () <$> c_sqlite3_create_function_v2 db namePtr nArgs flags (castStablePtrToPtr u) funPtr nullFunPtr nullFunPtr nullFunPtr
+
+-- Convert a 'CError' to a 'Either Error', in the common case where
+-- SQLITE_OK signals success and anything else signals an error.
+--
+-- Note that SQLITE_OK == 0.
+toResult :: a -> CError -> Either Error a
+toResult a (CError 0) = Right a
+toResult _ code = Left $ decodeError code
+
+-- call c_sqlite3_result_error in the event of an error
+catchAsResultError :: Ptr CContext -> IO () -> IO ()
+catchAsResultError ctx action = catch action $ \exn -> do
+  let msg = show (exn :: SomeException)
+  withCAStringLen msg $ \(ptr, len) ->
+    c_sqlite3_result_error ctx ptr (fromIntegral len)

tests/AgentTests/FunctionalAPITests.hs

@@ -3685,7 +3685,7 @@ insertUser :: DBStore -> IO ()
 insertUser st = withTransaction st (`DB.execute_` "INSERT INTO users DEFAULT VALUES")
 #else
 createStore :: String -> IO (Either MigrationError DBStore)
-createStore dbPath = createAgentStore (DBOpts dbPath "" False True DB.TQOff) (MigrationConfig MCError Nothing)
+createStore dbPath = createAgentStore (DBOpts dbPath [] "" False True DB.TQOff) (MigrationConfig MCError Nothing)
 
 insertUser :: DBStore -> IO ()
 insertUser st = withTransaction st (`DB.execute_` "INSERT INTO users (user_id) VALUES (1)")

tests/AgentTests/MigrationTests.hs

@@ -230,6 +230,7 @@ createStore randSuffix migrations confirmMigrations = do
   let dbOpts =
         DBOpts {
           dbFilePath = testDB randSuffix,
+          dbFunctions = [],
           dbKey = "",
           keepKey = False,
           vacuum = True,

tests/AgentTests/SQLiteTests.hs

@@ -72,7 +72,7 @@ withStore2 = before connect2 . after (removeStore . fst)
     connect2 :: IO (DBStore, DBStore)
     connect2 = do
       s1@DBStore {dbFilePath} <- createStore'
-      s2 <- connectSQLiteStore dbFilePath "" False DB.TQOff
+      s2 <- connectSQLiteStore $ DBOpts dbFilePath [] "" False False DB.TQOff
       pure (s1, s2)
 
 createStore' :: IO DBStore
@@ -83,7 +83,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) appMigrations (MigrationConfig MCError Nothing)
+  Right st <- createDBStore (DBOpts (testDB <> show r) [] key keepKey True DB.TQOff) appMigrations (MigrationConfig MCError Nothing)
   withTransaction' st (`SQL.execute_` "INSERT INTO users (user_id) VALUES (1);")
   pure st
 

tests/AgentTests/SchemaDump.hs

@@ -51,22 +51,22 @@ testVerifySchemaDump :: IO ()
 testVerifySchemaDump = do
   savedSchema <- ifM (doesFileExist appSchema) (readFile appSchema) (pure "")
   savedSchema `deepseq` pure ()
-  void $ createDBStore (DBOpts testDB "" False True TQOff) appMigrations (MigrationConfig MCConsole Nothing)
+  void $ createDBStore (DBOpts testDB [] "" False True TQOff) appMigrations (MigrationConfig MCConsole Nothing)
   getSchema testDB appSchema `shouldReturn` savedSchema
   removeFile testDB
 
 testVerifyLintFKeyIndexes :: IO ()
 testVerifyLintFKeyIndexes = do
   savedLint <- ifM (doesFileExist appLint) (readFile appLint) (pure "")
   savedLint `deepseq` pure ()
-  void $ createDBStore (DBOpts testDB "" False True TQOff) appMigrations (MigrationConfig MCConsole Nothing)
+  void $ createDBStore (DBOpts testDB [] "" False True TQOff) appMigrations (MigrationConfig MCConsole Nothing)
   getLintFKeyIndexes testDB "tests/tmp/agent_lint.sql" `shouldReturn` savedLint
   removeFile testDB
 
 testSchemaMigrations :: IO ()
 testSchemaMigrations = do
   let noDownMigrations = dropWhileEnd (\Migration {down} -> isJust down) appMigrations
-  Right st <- createDBStore (DBOpts testDB "" False True TQOff) noDownMigrations (MigrationConfig MCError Nothing)
+  Right st <- createDBStore (DBOpts testDB [] "" False True TQOff) noDownMigrations (MigrationConfig MCError Nothing)
   mapM_ (testDownMigration st) $ drop (length noDownMigrations) appMigrations
   closeDBStore st
   removeFile testDB
@@ -89,19 +89,19 @@ testSchemaMigrations = do
 
 testUsersMigrationNew :: IO ()
 testUsersMigrationNew = do
-  Right st <- createDBStore (DBOpts testDB "" False True TQOff) appMigrations (MigrationConfig MCError Nothing)
+  Right st <- createDBStore (DBOpts testDB [] "" False True TQOff) appMigrations (MigrationConfig MCError Nothing)
   withTransaction' st (`SQL.query_` "SELECT user_id FROM users;")
     `shouldReturn` ([] :: [Only Int])
   closeDBStore st
 
 testUsersMigrationOld :: IO ()
 testUsersMigrationOld = do
   let beforeUsers = takeWhile (("m20230110_users" /=) . name) appMigrations
-  Right st <- createDBStore (DBOpts testDB "" False True TQOff) beforeUsers (MigrationConfig MCError Nothing)
+  Right st <- createDBStore (DBOpts testDB [] "" False True TQOff) beforeUsers (MigrationConfig MCError Nothing)
   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) appMigrations (MigrationConfig MCYesUp Nothing)
+  Right st' <- createDBStore (DBOpts testDB [] "" False True TQOff) appMigrations (MigrationConfig MCYesUp Nothing)
   withTransaction' st' (`SQL.query_` "SELECT user_id FROM users;")
     `shouldReturn` ([Only (1 :: Int)])
   closeDBStore st'