ResizeSpec.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE QuasiQuotes #-}

module Test.ResizeSpec where

import Control.Exception (finally)
import Data.Aeson.KeyMap qualified as KM
import Data.List qualified as L
import Database.Bloodhound.Common.Client qualified as Client
import Database.Bloodhound.Common.Requests as Requests
import TestsUtils.Common
import TestsUtils.Import

-- Test fixtures: a source index plus one target per operation. Each
-- integration test seeds the source, runs the operation, then asserts
-- the target was created with the requested shard count. Cleanup is
-- always run before AND after so partial state from a previous failed
-- run does not poison the next one.
resizeSrcIndex :: IndexName
resizeSrcIndex = [qqIndexName|bh-resize-src|]

shrinkTargetIndex :: IndexName
shrinkTargetIndex = [qqIndexName|bh-resize-shrunk|]

splitTargetIndex :: IndexName
splitTargetIndex = [qqIndexName|bh-resize-split|]

cloneTargetIndex :: IndexName
cloneTargetIndex = [qqIndexName|bh-resize-cloned|]

-- | Delete every index used by this spec, ignoring any that do not
-- exist. Called both before seeding (defensive cleanup of leftover
-- state) and after the assertions (exception-safe teardown).
cleanupResizeIndices :: BH IO ()
cleanupResizeIndices = do
  _ <- tryEsError $ performBHRequest $ deleteIndex resizeSrcIndex
  _ <- tryEsError $ performBHRequest $ deleteIndex shrinkTargetIndex
  _ <- tryEsError $ performBHRequest $ deleteIndex splitTargetIndex
  _ <- tryEsError $ performBHRequest $ deleteIndex cloneTargetIndex
  pure ()

-- | Create @bh-resize-src@ with the requested primary shard count (and
-- zero replicas, to keep the operation local to a single node), then
-- mark it read-only. Shrink, split and clone all refuse to operate on
-- a source index that is still writable.
--
-- Replica count is forced to zero because the docker-compose test
-- cluster is single-node; assigning replicas that can never be
-- allocated would leave the index YELLOW and could trip the
-- @wait_for_active_shards@ check on the resize call.
seedReadOnlySource :: ShardCount -> BH IO ()
seedReadOnlySource shardCount = do
  _ <- cleanupResizeIndices
  let settings = IndexSettings shardCount (ReplicaCount 0) defaultIndexMappingsLimits
  _ <- performBHRequest $ createIndex settings resizeSrcIndex
  -- Source must be marked read-only before the resize call: shrink,
  -- split and clone all refuse to operate on a still-writable index.
  _ <- performBHRequest $ updateIndexSettings (BlocksWrite True :| []) resizeSrcIndex
  pure ()

-- | Wait for the cluster to settle on the resize result and confirm
-- the target index exists with the expected shard count. ES returns
-- Acknowledged=True before the target is fully allocated, so we ask
-- the cluster to reach yellow health before reading it back.
assertTargetIndex :: IndexName -> ShardCount -> BH IO ()
assertTargetIndex target expectedShards = do
  _ <- performBHRequest $ Requests.waitForYellowIndex target
  info <- performBHRequest $ getIndex target
  liftIO $ do
    iiIndexName info `shouldBe` target
    indexShards (iiFixedSettings info) `shouldBe` expectedShards

-- | Body for the three 'defaultResizeSettings' values: an empty object
-- because no @settings@ \/ @mappings@ \/ @aliases@ are supplied.

-- Single-alias payload reused by the "carries aliases" assertion.
sampleAliases :: KM.KeyMap Value
sampleAliases = KM.singleton "bh-resize-alias" (object [])

-- | Sort query-string pairs so the assertion is order-insensitive —
-- 'withQueries' is documented as order-insensitive, so the test
-- shouldn't depend on how the params happen to be sequenced.
normalizeQueryString :: (Ord a, Ord b) => [(a, b)] -> [(a, b)]
normalizeQueryString = L.sort

spec :: Spec
spec = do
  describe "shrinkIndex / splitIndex / cloneIndex (request shape)" $ do
    -- The body builder delegates to 'createIndexOptionsBody', so an
    -- empty 'ShrinkSettings' must produce an empty body (not 'Nothing',
    -- not "{}") — the server treats an absent body as "copy source
    -- verbatim", which is exactly the documented behaviour.
    it "shrinkIndex with defaultShrinkSettings sends an empty body" $ do
      let req = Requests.shrinkIndex resizeSrcIndex shrinkTargetIndex defaultShrinkSettings
      bhRequestBody req `shouldBe` Just ""

    it "splitIndex with defaultSplitSettings sends an empty body" $ do
      let req = Requests.splitIndex resizeSrcIndex splitTargetIndex defaultSplitSettings
      bhRequestBody req `shouldBe` Just ""

    it "cloneIndex with defaultCloneSettings sends an empty body" $ do
      let req = Requests.cloneIndex resizeSrcIndex cloneTargetIndex defaultCloneSettings
      bhRequestBody req `shouldBe` Just ""

    -- When 'cioSettings' is populated, the body must contain a
    -- top-level @settings.index.number_of_shards@ field — this is how
    -- the caller tells the server the target shard count for shrink
    -- (must be a divisor of the source) or split (must be a multiple).
    it "shrinkIndex forwards cioSettings as {settings:{index:{...}}}" $ do
      let settings =
            ShrinkSettings
              defaultCreateIndexOptions
                { cioSettings =
                    Just $
                      IndexSettings (ShardCount 1) (ReplicaCount 0) defaultIndexMappingsLimits
                }
      let req = Requests.shrinkIndex resizeSrcIndex shrinkTargetIndex settings
      let expected =
            object
              [ "settings"
                  .= object
                    [ "index"
                        .= object
                          [ "number_of_shards" .= (1 :: Int),
                            "number_of_replicas" .= (0 :: Int),
                            "mapping" .= defaultIndexMappingsLimits
                          ]
                    ]
              ]
      (decode =<< bhRequestBody req) `shouldBe` Just expected

    -- Likewise, aliases supplied via 'cioAliases' must land under the
    -- top-level @aliases@ key. Verified on cloneIndex so we exercise
    -- the third newtype's path through 'resizeEndpoint'.
    it "cloneIndex forwards cioAliases as {aliases:{...}}" $ do
      let settings =
            CloneSettings
              defaultCreateIndexOptions
                { cioAliases = Just sampleAliases
                }
      let req = Requests.cloneIndex resizeSrcIndex cloneTargetIndex settings
      let expected = object ["aliases" .= object ["bh-resize-alias" .= object []]]
      (decode =<< bhRequestBody req) `shouldBe` Just expected

    -- The newtypes must also forward URI parameters through
    -- 'createIndexOptionsParams'. This is the whole point of wrapping
    -- the full 'CreateIndexOptions' rather than rolling a narrower
    -- record — the typed resize settings preserve the entire
    -- body+param surface of create-index.
    it "splitIndex forwards cioWaitForActiveShards / master_timeout / timeout as query params" $ do
      let settings =
            SplitSettings
              defaultCreateIndexOptions
                { cioWaitForActiveShards = Just AllActiveShards,
                  cioMasterTimeout = Just (TimeUnitSeconds, 30),
                  cioTimeout = Just (TimeUnitMilliseconds, 500)
                }
      let req = Requests.splitIndex resizeSrcIndex splitTargetIndex settings
      normalizeQueryString (getRawEndpointQueries (bhRequestEndpoint req))
        `shouldBe` [ ("master_timeout", Just "30s"),
                     ("timeout", Just "500ms"),
                     ("wait_for_active_shards", Just "all")
                   ]

  describe "shrinkIndex (integration)" $ do
    -- Source @bh-resize-src@ has 2 primary shards; the shrink target
    -- @bh-resize-shrunk@ requests 1 shard (1 is a divisor of 2). The
    -- single-node docker-compose cluster satisfies the "all shards on
    -- one node" precondition trivially.
    it "shrinks a 2-shard source into a 1-shard target" $
      withTestEnv
        ( do
            _ <- seedReadOnlySource (ShardCount 2)
            resp <-
              Client.shrinkIndex
                resizeSrcIndex
                shrinkTargetIndex
                ( ShrinkSettings
                    defaultCreateIndexOptions
                      { cioSettings =
                          Just $
                            IndexSettings (ShardCount 1) (ReplicaCount 0) defaultIndexMappingsLimits
                      }
                )
            liftIO $ resp `shouldBe` Acknowledged True
            assertTargetIndex shrinkTargetIndex (ShardCount 1)
        )
        `finally` withTestEnv cleanupResizeIndices

  describe "splitIndex (integration)" $ do
    -- Source @bh-resize-src@ has 1 primary shard; the split target
    -- @bh-resize-split@ requests 2 shards (2 is a multiple of 1).
    it "splits a 1-shard source into a 2-shard target" $
      withTestEnv
        ( do
            _ <- seedReadOnlySource (ShardCount 1)
            resp <-
              Client.splitIndex
                resizeSrcIndex
                splitTargetIndex
                ( SplitSettings
                    defaultCreateIndexOptions
                      { cioSettings =
                          Just $
                            IndexSettings (ShardCount 2) (ReplicaCount 0) defaultIndexMappingsLimits
                      }
                )
            liftIO $ resp `shouldBe` Acknowledged True
            assertTargetIndex splitTargetIndex (ShardCount 2)
        )
        `finally` withTestEnv cleanupResizeIndices

  describe "cloneIndex (integration)" $ do
    -- Source @bh-resize-src@ has 1 primary shard; the clone target
    -- @bh-resize-cloned@ requests no overriding settings, so it
    -- inherits the source's shard count.
    it "clones a 1-shard source into a 1-shard target" $
      withTestEnv
        ( do
            _ <- seedReadOnlySource (ShardCount 1)
            resp <- Client.cloneIndex resizeSrcIndex cloneTargetIndex defaultCloneSettings
            liftIO $ resp `shouldBe` Acknowledged True
            assertTargetIndex cloneTargetIndex (ShardCount 1)
        )
        `finally` withTestEnv cleanupResizeIndices