test: track agent query plans (#5571)

tests/ChatClient.hs

@@ -70,6 +70,7 @@ import Database.PostgreSQL.Simple (ConnectInfo (..), defaultConnectInfo)
 #else
 import Data.ByteArray (ScrubbedBytes)
 import qualified Data.Map.Strict as M
+import Simplex.Messaging.Agent.Client (agentClientStore)
 import Simplex.Messaging.Agent.Store.Common (withConnection)
 import System.FilePath ((</>))
 #endif
@@ -324,8 +325,10 @@ stopTestChat ps TestCC {chatController = cc@ChatController {smpAgent, chatStore}
   uninterruptibleCancel chatAsync
   liftIO $ disposeAgentClient smpAgent
 #if !defined(dbPostgres)
-  stats <- withConnection chatStore $ readTVarIO . DB.slow
-  atomically $ modifyTVar' (queryStats ps) $ M.unionWith combineStats stats
+  chatStats <- withConnection chatStore $ readTVarIO . DB.slow
+  atomically $ modifyTVar' (chatQueryStats ps) $ M.unionWith combineStats chatStats
+  agentStats <- withConnection (agentClientStore smpAgent) $ readTVarIO . DB.slow
+  atomically $ modifyTVar' (agentQueryStats ps) $ M.unionWith combineStats agentStats
 #endif
   closeDBStore chatStore
   threadDelay 200000

tests/ChatTests/DBUtils/SQLite.hs

@@ -6,5 +6,6 @@ import Simplex.Messaging.TMap (TMap)
 
 data TestParams = TestParams
   { tmpPath :: FilePath,
-    queryStats :: TMap Query SlowQueryStats
+    chatQueryStats :: TMap Query SlowQueryStats,
+    agentQueryStats :: TMap Query SlowQueryStats
   }

tests/SchemaDump.hs

@@ -20,12 +20,13 @@ import qualified Data.Text.IO as T
 import Database.SQLite.Simple (Query (..))
 import Simplex.Chat.Store (createChatStore)
 import qualified Simplex.Chat.Store as Store
+import Simplex.Messaging.Agent.Env.SQLite (createAgentStore)
 import Simplex.Messaging.Agent.Store.Common (withConnection)
-import Simplex.Messaging.Agent.Store.Interface
-import Simplex.Messaging.Agent.Store.Shared (Migration (..), MigrationConfirmation (..), MigrationsToRun (..), toDownMigration)
 import Simplex.Messaging.Agent.Store.DB (TrackQueries (..))
 import qualified Simplex.Messaging.Agent.Store.DB as DB
+import Simplex.Messaging.Agent.Store.Interface
 import qualified Simplex.Messaging.Agent.Store.SQLite.Migrations as Migrations
+import Simplex.Messaging.Agent.Store.Shared (Migration (..), MigrationConfirmation (..), MigrationsToRun (..), toDownMigration)
 import Simplex.Messaging.Util (ifM, tshow, whenM)
 import System.Directory (doesFileExist, removeFile)
 import System.Process (readCreateProcess, shell)
@@ -34,6 +35,9 @@ import Test.Hspec
 testDB :: FilePath
 testDB = "tests/tmp/test_chat.db"
 
+testAgentDB :: FilePath
+testAgentDB = "tests/tmp/test_agent.db"
+
 appSchema :: FilePath
 appSchema = "src/Simplex/Chat/Store/SQLite/Migrations/chat_schema.sql"
 
@@ -53,8 +57,11 @@ appSchema = "src/Simplex/Chat/Store/SQLite/Migrations/chat_schema.sql"
 appLint :: FilePath
 appLint = "src/Simplex/Chat/Store/SQLite/Migrations/chat_lint.sql"
 
-appQueryPlans :: FilePath
-appQueryPlans = "src/Simplex/Chat/Store/SQLite/Migrations/chat_query_plans.txt"
+appChatQueryPlans :: FilePath
+appChatQueryPlans = "src/Simplex/Chat/Store/SQLite/Migrations/chat_query_plans.txt"
+
+appAgentQueryPlans :: FilePath
+appAgentQueryPlans = "src/Simplex/Chat/Store/SQLite/Migrations/agent_query_plans.txt"
 
 testSchema :: FilePath
 testSchema = "tests/tmp/test_agent_schema.sql"
@@ -138,18 +145,35 @@ getLintFKeyIndexes dbPath lintPath = do
   lint `deepseq` pure lint
 
 saveQueryPlans :: SpecWith TestParams
-saveQueryPlans = it "verify and overwrite query plans" $ \TestParams {queryStats} -> do
-  savedPlans <- ifM (doesFileExist appQueryPlans) (T.readFile appQueryPlans) (pure "")
-  savedPlans `deepseq` pure ()
-  queries <- sort . M.keys <$> readTVarIO queryStats
-  Right st <- createChatStore (DBOpts testDB "" False True TQOff) MCError
-  plans' <- withConnection st $ \db -> do
-    DB.execute_ db "CREATE TABLE IF NOT EXISTS temp_conn_ids (conn_id BLOB)"
-    mapM (getQueryPlan db) queries
-  let savedPlans' = T.unlines plans'
-  T.writeFile appQueryPlans savedPlans'
-  savedPlans' `shouldBe` savedPlans
+saveQueryPlans = it "verify and overwrite query plans" $ \TestParams {chatQueryStats, agentQueryStats} -> do
+  (chatSavedPlans, chatSavedPlans') <-
+    updatePlans
+      appChatQueryPlans
+      chatQueryStats
+      (createChatStore (DBOpts testDB "" False True TQOff) MCError)
+      (`DB.execute_` "CREATE TABLE IF NOT EXISTS temp_conn_ids (conn_id BLOB)")
+  (agentSavedPlans, agentSavedPlans') <-
+    updatePlans
+      appAgentQueryPlans
+      agentQueryStats
+      (createAgentStore (DBOpts testAgentDB "" False True TQOff) MCError)
+      (const $ pure ())
+  chatSavedPlans' `shouldBe` chatSavedPlans
+  agentSavedPlans' `shouldBe` agentSavedPlans
+  removeFile testDB
+  removeFile testAgentDB
   where
+    updatePlans plansFile statsSel createStore prepareStore = do
+      savedPlans <- ifM (doesFileExist plansFile) (T.readFile plansFile) (pure "")
+      savedPlans `deepseq` pure ()
+      queries <- sort . M.keys <$> readTVarIO statsSel
+      Right st <- createStore
+      plans' <- withConnection st $ \db -> do
+        void $ prepareStore db
+        mapM (getQueryPlan db) queries
+      let savedPlans' = T.unlines plans'
+      T.writeFile plansFile savedPlans'
+      pure (savedPlans, savedPlans')
     getQueryPlan :: DB.Connection -> Query -> IO Text
     getQueryPlan db q =
       (("Query: " <> fromQuery q) <>) . result <$> E.try (DB.query_ db $ "explain query plan " <> q)

tests/Test.hs

@@ -34,7 +34,8 @@ main :: IO ()
 main = do
   setLogLevel LogError
 #if !defined(dbPostgres)
-  queryStats <- TM.emptyIO
+  chatQueryStats <- TM.emptyIO
+  agentQueryStats <- TM.emptyIO
 #endif
   withGlobalLogging logCfg . hspec
 #if defined(dbPostgres)
@@ -59,7 +60,7 @@ main = do
       around testBracket
         . after_ (dropAllSchemasExceptSystem testDBConnectInfo)
 #else
-      around (testBracket queryStats)
+      around (testBracket chatQueryStats agentQueryStats)
 #endif
         $ do
 #if !defined(dbPostgres)
@@ -73,10 +74,11 @@ main = do
           xdescribe'' "Save query plans" saveQueryPlans
 #endif
   where
-#if defined(dbPostgres)    
+#if defined(dbPostgres)
     testBracket test = withSmpServer $ tmpBracket $ test . TestParams
 #else
-    testBracket queryStats test = withSmpServer $ tmpBracket $ \tmpPath -> test TestParams {tmpPath, queryStats}
+    testBracket chatQueryStats agentQueryStats test =
+      withSmpServer $ tmpBracket $ \tmpPath -> test TestParams {tmpPath, chatQueryStats, agentQueryStats}
 #endif
     tmpBracket test = do
       t <- getSystemTime