Improve ES migration test coverage

services/brig/src/Brig/Index/Options.hs

@@ -39,7 +39,7 @@ module Brig.Index.Options
     commandParser,
     mkCreateIndexSettings,
     toESServer,
-    ReindexFromAnotherIndexSettings,
+    ReindexFromAnotherIndexSettings (..),
     reindexDestIndex,
     reindexTimeoutSeconds,
     reindexEsConnection,

services/brig/test/integration/API/Search.hs

@@ -1,6 +1,7 @@
 {-# LANGUAGE OverloadedRecordDot #-}
 {-# LANGUAGE PartialTypeSignatures #-}
 {-# LANGUAGE QuasiQuotes #-}
+{-# LANGUAGE RecordWildCards #-}
 {-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}
 {-# OPTIONS_GHC -Wno-partial-type-signatures #-}
 {-# OPTIONS_GHC -Wno-redundant-constraints #-}
@@ -33,12 +34,17 @@ import API.User.Util
 import Bilge
 import Bilge.Assert
 import Brig.App (initHttpManagerWithTLSConfig)
+import Brig.Index.Eval (runCommand)
+import Brig.Index.Options
+import Brig.Index.Options qualified as IndexOpts
+import Brig.Options (ElasticSearchOpts)
 import Brig.Options qualified as Opt
 import Brig.Options qualified as Opts
+import Cassandra qualified as C
+import Cassandra.Options qualified as CassOpts
 import Control.Lens ((.~), (?~), (^.))
-import Control.Monad.Catch (MonadCatch, MonadThrow)
-import Control.Retry
-import Data.Aeson (FromJSON, Value, decode)
+import Control.Monad.Catch (MonadCatch)
+import Data.Aeson (Value, decode)
 import Data.Aeson qualified as Aeson
 import Data.Domain (Domain (Domain))
 import Data.Handle (fromHandle)
@@ -58,13 +64,15 @@ import Network.Wai qualified as Wai
 import Network.Wai.Handler.Warp qualified as Warp
 import Network.Wai.Test qualified as WaiTest
 import Safe (headMay)
+import System.Logger qualified as Log
 import Test.QuickCheck (Arbitrary (arbitrary), generate)
 import Test.Tasty
 import Test.Tasty.HUnit
 import Text.RawString.QQ (r)
 import URI.ByteString qualified as URI
 import UnliftIO (Concurrently (..), async, bracket, cancel, runConcurrently)
 import Util
+import Util.Options (Endpoint)
 import Wire.API.Federation.API.Brig (SearchResponse (SearchResponse))
 import Wire.API.Team.Feature
 import Wire.API.Team.SearchVisibility
@@ -90,7 +98,11 @@ tests opts mgr galley brig = do
         testWithBothIndices opts mgr "Non ascii names" $ testSearchNonAsciiNames brig,
         testWithBothIndices opts mgr "user with umlaut" $ testSearchWithUmlaut brig,
         testWithBothIndices opts mgr "user with japanese name" $ testSearchCJK brig,
-        test mgr "migration to new index" $ testMigrationToNewIndex opts brig,
+        testGroup "index migration" $
+          [ test mgr "migration to new index from existing index" $ testMigrationToNewIndex opts brig runReindexFromAnotherIndex,
+            test mgr "migration to new index from database" $ testMigrationToNewIndex opts brig (runReindexFromDatabase Reindex),
+            test mgr "migration to new index from database (force sync)" $ testMigrationToNewIndex opts brig (runReindexFromDatabase ReindexSameOrNewer)
+          ],
         testGroup "team A: SearchVisibilityStandard (= unrestricted outbound search)" $
           [ testGroup "team A: SearchableByOwnTeam (= restricted inbound search)" $
               [ testWithBothIndices opts mgr "  I. non-team user cannot find team A member by display name" $ testSearchTeamMemberAsNonMemberDisplayName mgr brig galley FeatureStatusDisabled,
@@ -608,8 +620,13 @@ testSearchOtherDomain opts brig = do
 -- cluster. This test spins up a proxy server to pass requests to our only ES
 -- server. The proxy server ensures that only requests to the 'old' index go
 -- through.
-testMigrationToNewIndex :: (TestConstraints m, MonadUnliftIO m) => Opt.Opts -> Brig -> m ()
-testMigrationToNewIndex opts brig = do
+testMigrationToNewIndex ::
+  (TestConstraints m, MonadUnliftIO m) =>
+  Opt.Opts ->
+  Brig ->
+  (Log.Logger -> Opt.Opts -> ES.IndexName -> IO ()) ->
+  m ()
+testMigrationToNewIndex opts brig reindexCommand = do
   withOldESProxy opts $ \oldESUrl oldESIndex -> do
     let optsOldIndex =
           opts
@@ -652,9 +669,11 @@ testMigrationToNewIndex opts brig = do
     assertCanFindByName brig phase1TeamUser1 phase2TeamUser
 
     -- Run Migrations
-    let newIndexName = opts ^. Opt.elasticsearchLens . Opt.indexLens
-    taskNodeId <- assertRight =<< runBH opts (ES.reindexAsync $ ES.mkReindexRequest (ES.IndexName oldESIndex) newIndexName)
-    runBH opts $ waitForTaskToComplete @ES.ReindexResponse taskNodeId
+    let oldIndexName = ES.IndexName oldESIndex
+    logger <- Log.create Log.StdOut
+    liftIO $ do
+      createCommand logger opts oldIndexName
+      reindexCommand logger opts oldIndexName
 
     -- Phase 3: Using old index for search, writing to both indices, migrations have run
     refreshIndex brig
@@ -688,6 +707,71 @@ testMigrationToNewIndex opts brig = do
     assertCanFindByName brig phase1TeamUser1 phase3NonTeamUser
     assertCanFindByName brig phase1TeamUser1 phase3TeamUser
 
+createCommand :: Log.Logger -> Opt.Opts -> ES.IndexName -> IO ()
+createCommand logger opts oldIndexName =
+  let newIndexName = opts ^. Opt.elasticsearchLens . Opt.indexLens
+      esOldOpts :: Opt.ElasticSearchOpts = (opts ^. Opt.elasticsearchLens) & (Opt.indexLens .~ oldIndexName)
+      esOldConnectionSettings :: ESConnectionSettings = toESConnectionSettings esOldOpts
+      esNewConnectionSettings = esOldConnectionSettings {esIndex = newIndexName}
+      replicas = 2
+      shards = 2
+      refreshInterval = 5
+      esSettings =
+        IndexOpts.localElasticSettings
+          & IndexOpts.esConnection .~ esNewConnectionSettings
+          & IndexOpts.esIndexReplicas .~ ES.ReplicaCount replicas
+          & IndexOpts.esIndexShardCount .~ shards
+          & IndexOpts.esIndexRefreshInterval .~ refreshInterval
+   in runCommand logger $ Create esSettings opts.galley
+
+runReindexFromAnotherIndex :: Log.Logger -> Opt.Opts -> ES.IndexName -> IO ()
+runReindexFromAnotherIndex logger opts oldIndexName =
+  let newIndexName = opts ^. Opt.elasticsearchLens . Opt.indexLens
+      esOldOpts :: Opt.ElasticSearchOpts = (opts ^. Opt.elasticsearchLens) & (Opt.indexLens .~ oldIndexName)
+      esOldConnectionSettings :: ESConnectionSettings = toESConnectionSettings esOldOpts
+      reindexSettings = ReindexFromAnotherIndexSettings esOldConnectionSettings newIndexName 5
+   in runCommand logger $ ReindexFromAnotherIndex reindexSettings
+
+runReindexFromDatabase ::
+  (ElasticSettings -> CassandraSettings -> Endpoint -> Command) ->
+  Log.Logger ->
+  Opt.Opts ->
+  ES.IndexName ->
+  IO ()
+runReindexFromDatabase syncCommand logger opts oldIndexName =
+  let newIndexName = opts ^. Opt.elasticsearchLens . Opt.indexLens
+      esOldOpts :: Opt.ElasticSearchOpts = (opts ^. Opt.elasticsearchLens) & (Opt.indexLens .~ oldIndexName)
+      esOldConnectionSettings :: ESConnectionSettings = toESConnectionSettings esOldOpts
+      esNewConnectionSettings = esOldConnectionSettings {esIndex = newIndexName}
+      replicas = 2
+      shards = 2
+      refreshInterval = 5
+      elasticSettings :: ElasticSettings =
+        IndexOpts.localElasticSettings
+          & IndexOpts.esConnection .~ esNewConnectionSettings
+          & IndexOpts.esIndexReplicas .~ ES.ReplicaCount replicas
+          & IndexOpts.esIndexShardCount .~ shards
+          & IndexOpts.esIndexRefreshInterval .~ refreshInterval
+      cassandraSettings :: CassandraSettings =
+        ( localCassandraSettings
+            & IndexOpts.cHost .~ (Text.unpack opts.cassandra.endpoint.host)
+            & IndexOpts.cPort .~ (opts.cassandra.endpoint.port)
+            & IndexOpts.cKeyspace .~ (C.Keyspace opts.cassandra.keyspace)
+        )
+
+      endpoint :: Endpoint = opts.galley
+   in runCommand logger $ syncCommand elasticSettings cassandraSettings endpoint
+
+toESConnectionSettings :: ElasticSearchOpts -> ESConnectionSettings
+toESConnectionSettings opts = ESConnectionSettings {..}
+  where
+    toText (ES.Server url) = url
+    esServer = (fromRight undefined . URI.parseURI URI.strictURIParserOptions . Text.encodeUtf8 . toText) opts.url
+    esIndex = opts.index
+    esCaCert = opts.caCert
+    esInsecureSkipVerifyTls = opts.insecureSkipVerifyTls
+    esCredentials = opts.credentials
+
 withOldESProxy :: (TestConstraints m, MonadUnliftIO m, HasCallStack) => Opt.Opts -> (Text -> Text -> m a) -> m a
 withOldESProxy opts f = do
   indexName <- randomHandle
@@ -712,18 +796,6 @@ indexProxyServer idx opts mgr =
             else Wai.WPRResponse (Wai.responseLBS HTTP.status400 [] $ "Refusing to proxy to path=" <> cs (Wai.rawPathInfo req))
    in waiProxyTo proxyApp Wai.defaultOnExc mgr
 
-waitForTaskToComplete :: forall a m. (ES.MonadBH m, MonadThrow m, FromJSON a) => ES.TaskNodeId -> m ()
-waitForTaskToComplete taskNodeId = do
-  let policy = constantDelay 100000 <> limitRetries 30
-  let retryCondition _ = fmap not . isTaskComplete
-  task <- retrying policy retryCondition (const $ ES.getTask @m @a taskNodeId)
-  taskCompleted <- isTaskComplete task
-  liftIO $ assertBool "Timed out waiting for task" taskCompleted
-  where
-    isTaskComplete :: Either ES.EsError (ES.TaskResponse a) -> m Bool
-    isTaskComplete (Left e) = liftIO $ assertFailure $ "Expected Right, got Left: " <> show e
-    isTaskComplete (Right taskRes) = pure $ ES.taskResponseCompleted taskRes
-
 testWithBothIndices :: Opt.Opts -> Manager -> TestName -> WaiTest.Session a -> TestTree
 testWithBothIndices opts mgr name f = do
   testGroup