MultiSearchTemplateSpec.hs

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

module Test.MultiSearchTemplateSpec where

import Data.Aeson qualified as Aeson
import Data.Aeson.Key qualified as K
import Data.Aeson.KeyMap qualified as KM
import Data.ByteString.Lazy.Char8 qualified as LC
import Data.Either (rights)
import Data.List (sort)
import Database.Bloodhound.Common.Types
import TestsUtils.Common
import TestsUtils.Import

-- | Smart constructor that mirrors the 2-field form of
-- 'MultiSearchTemplateItem'. Keeps the integration-test fixtures readable
-- while the record itself carries the full per-header field set.
mkItem :: Maybe IndexName -> SearchTemplate -> MultiSearchTemplateItem
mkItem idx t =
  MultiSearchTemplateItem
    { multiSearchTemplateItemIndex = idx,
      multiSearchTemplateItemRouting = Nothing,
      multiSearchTemplateItemSearchType = Nothing,
      multiSearchTemplateItemPreference = Nothing,
      multiSearchTemplateItemAllowPartialSearchResults = Nothing,
      multiSearchTemplateItemSearchTemplate = t
    }

-- | Decode a ToJSON-encoded value back to a sorted list of key names so the
-- tests can assert on the header line as a set rather than relying on
-- aeson's key ordering.
headerKeys :: LC.ByteString -> [Text]
headerKeys bs = case Aeson.decode bs :: Maybe Aeson.Object of
  Just o -> sort (map K.toText (KM.keys o))
  Nothing -> error "headerKeys: decode failed"

-- | Split the NDJSON body into its non-empty lines so tests can assert on
-- line count and per-line shape without depending on trailing-newline
-- placement.
bodyLines :: LC.ByteString -> [LC.ByteString]
bodyLines = filter (not . LC.null) . LC.lines

-- | A canned inline-source template that matches the seeded @Tweet@ in
-- @insertData@ (user = "bitemyapp") when its params are filled in.
matchingInlineTemplate :: SearchTemplate
matchingInlineTemplate =
  mkSearchTemplate
    (Right (SearchTemplateSource "{\"query\": { \"match\" : { \"{{my_field}}\" : \"{{my_value}}\" } } }"))
    templateParams

-- | Params for 'matchingInlineTemplate', factored out so the same values
-- can be reused for a stored-id variant in an integration test.
templateParams :: TemplateQueryKeyValuePairs
templateParams =
  TemplateQueryKeyValuePairs $
    KM.fromList
      [ ("my_field", "user"),
        ("my_value", "bitemyapp")
      ]

spec :: Spec
spec = do
  describe "MultiSearchTemplateItem ToJSON header rendering" $ do
    let bareTemplate = matchingInlineTemplate
        logs = [qqIndexName|logs|]

    it "renders only index when no per-header fields are set" $ do
      let item = mkItem (Just logs) bareTemplate
      Aeson.encode item `shouldBe` "{\"index\":\"logs\"}"

    it "renders an empty object when no fields (including index) are set" $ do
      let item = mkItem Nothing bareTemplate
      Aeson.encode item `shouldBe` "{}"

    it "renders routing alongside index" $ do
      let item = (mkItem (Just logs) bareTemplate) {multiSearchTemplateItemRouting = Just "r1"}
      headerKeys (Aeson.encode item) `shouldBe` ["index", "routing"]

    it "renders search_type using the SearchType Text encoding" $ do
      let item = (mkItem (Just logs) bareTemplate) {multiSearchTemplateItemSearchType = Just SearchTypeDfsQueryThenFetch}
      Aeson.encode item
        `shouldBe` "{\"index\":\"logs\",\"search_type\":\"dfs_query_then_fetch\"}"

    it "renders preference and allow_partial_search_results" $ do
      let item =
            (mkItem (Just logs) bareTemplate)
              { multiSearchTemplateItemPreference = Just "_local",
                multiSearchTemplateItemAllowPartialSearchResults = Just False
              }
      headerKeys (Aeson.encode item)
        `shouldBe` ["allow_partial_search_results", "index", "preference"]

    it "renders every per-header field simultaneously" $ do
      let item =
            (mkItem (Just logs) bareTemplate)
              { multiSearchTemplateItemRouting = Just "r1",
                multiSearchTemplateItemSearchType = Just SearchTypeDfsQueryThenFetch,
                multiSearchTemplateItemPreference = Just "_local",
                multiSearchTemplateItemAllowPartialSearchResults = Just True
              }
      Aeson.encode item
        `shouldBe` "{\"allow_partial_search_results\":true,\"index\":\"logs\",\"preference\":\"_local\",\"routing\":\"r1\",\"search_type\":\"dfs_query_then_fetch\"}"

  describe "encodeMultiSearchTemplateItems NDJSON shape" $ do
    let logs = [qqIndexName|logs|]
        t1 = matchingInlineTemplate
        t2 =
          mkSearchTemplate
            (Right (SearchTemplateSource "{\"query\":{\"match_all\":{}}}"))
            (TemplateQueryKeyValuePairs KM.empty)

    it "produces one header + one body line per item, plus a trailing newline" $ do
      let items = mkItem (Just logs) t1 :| [mkItem Nothing t2]
          encoded = encodeMultiSearchTemplateItems items
      -- 2 items * 2 lines = 4 content newlines, plus the outer
      -- wrapper's trailing newline (mirrors 'encodeMultiSearchItems').
      LC.count '\n' encoded `shouldBe` 5
      length (bodyLines encoded) `shouldBe` 4

    it "every body line decodes to the JSON of one of the input templates" $ do
      -- Order-agnostic: the encoder builds the NDJSON stream with a
      -- 'foldr', which reverses iteration order, but each (header, body)
      -- pair still corresponds to one input item. Asserting on the *set*
      -- of body JSONs keeps this test honest without pinning an order
      -- that the existing @_msearch@ encoder does not guarantee either.
      let items = mkItem (Just logs) t1 :| [mkItem Nothing t2]
          encoded = encodeMultiSearchTemplateItems items
          lines' = bodyLines encoded
          -- Lines 0 and 2 are headers (objects with @index@ at most);
          -- lines 1 and 3 are bodies (objects with @source@ or @id@).
          body1 = Aeson.decode (lines' !! 1) :: Maybe Aeson.Value
          body2 = Aeson.decode (lines' !! 3) :: Maybe Aeson.Value
          decodedSet = sort [show body1, show body2]
          expectedSet =
            sort
              [ show (Aeson.decode (Aeson.encode t1) :: Maybe Aeson.Value),
                show (Aeson.decode (Aeson.encode t2) :: Maybe Aeson.Value)
              ]
      decodedSet `shouldBe` expectedSet

    it "the header line matches encode <item>" $ do
      -- The header line is the ToJSON of the MultiSearchTemplateItem
      -- (header-only, body dropped). Lock the wire shape so a regression
      -- in the encoder (e.g. accidentally serialising the SearchTemplate
      -- body in the header line) surfaces here.
      let item = mkItem (Just logs) t1
          items = item :| []
          encoded = encodeMultiSearchTemplateItems items
          lines' = bodyLines encoded
      Aeson.decode (head lines') `shouldBe` (Aeson.decode (Aeson.encode item) :: Maybe Aeson.Value)

  describe "multi-search template API" $ do
    it "executes a single inline-source template by index" $
      withTestEnv $ do
        _ <- insertData
        result <-
          tryPerformBHRequest $
            multiSearchTemplateByIndex @Tweet testIndex (matchingInlineTemplate :| [])
        case result of
          Left _ -> liftIO $ expectationFailure "multi-search template by index failed"
          Right resp -> do
            let responses = multiSearchResponses resp
            liftIO $ responses `shouldSatisfy` all (either (const False) (const True))
            liftIO $ length responses `shouldBe` 1
            let firstResult = head (rights responses)
            let myTweet = grabFirst firstResult
            liftIO $ myTweet `shouldBe` Right exampleTweet

    it "executes multiple templates in a single request" $
      withTestEnv $ do
        _ <- insertData
        _ <- insertOther
        let otherTemplate =
              mkSearchTemplate
                (Right (SearchTemplateSource "{\"query\":{\"match_all\":{}}}"))
                (TemplateQueryKeyValuePairs KM.empty)
            items = mkItem (Just testIndex) matchingInlineTemplate :| [mkItem (Just testIndex) otherTemplate]
        result <- tryPerformBHRequest $ multiSearchTemplate @Tweet items
        case result of
          Left _ -> liftIO $ expectationFailure "multi-search template failed"
          Right resp -> do
            let responses = multiSearchResponses resp
            liftIO $ responses `shouldSatisfy` all (either (const False) (const True))
            liftIO $ length responses `shouldBe` 2
            let myTweet = grabFirst (head (rights responses))
            liftIO $ myTweet `shouldBe` Right exampleTweet

    it "resolves a stored template id alongside an inline source" $
      withTestEnv $ do
        _ <- insertData
        let tid = SearchTemplateId "bh-msearch-template-spec"
        _ <- performBHRequest $ storeSearchTemplate tid (SearchTemplateSource "{\"query\": { \"match\" : { \"{{my_field}}\" : \"{{my_value}}\" } } }")
        let storedTemplate = mkSearchTemplate (Left tid) templateParams
            items = mkItem (Just testIndex) matchingInlineTemplate :| [mkItem (Just testIndex) storedTemplate]
        result <- tryPerformBHRequest $ multiSearchTemplate @Tweet items
        -- cleanup before assertions so a later re-run doesn't see a stale script
        _ <- performBHRequest $ deleteSearchTemplate tid
        case result of
          Left _ -> liftIO $ expectationFailure "multi-search template with stored id failed"
          Right resp -> do
            let responses = multiSearchResponses resp
            liftIO $ responses `shouldSatisfy` all (either (const False) (const True))
            liftIO $ length responses `shouldBe` 2
            let myTweet = grabFirst (head (rights responses))
            liftIO $ myTweet `shouldBe` Right exampleTweet

    it "multiSearchTemplateWith forwards URI params (max_concurrent_searches, typed_keys)" $
      withTestEnv $ do
        _ <- insertData
        _ <- insertOther
        let otherTemplate =
              mkSearchTemplate
                (Right (SearchTemplateSource "{\"query\":{\"match_all\":{}}}"))
                (TemplateQueryKeyValuePairs KM.empty)
            -- Index-pinned items: the cluster has system indices (e.g.
            -- @.tasks@) whose documents do not decode as 'Tweet', so a
            -- cluster-wide @match_all@ template would trip a parse
            -- exception in the response decoder. Pinning to 'testIndex'
            -- keeps the @match_all@ template scoped to twitter docs.
            items = mkItem (Just testIndex) matchingInlineTemplate :| [mkItem (Just testIndex) otherTemplate]
            opts =
              defaultSearchOptions
                { soMaxConcurrentSearches = Just 1,
                  soTypedKeys = Just True
                }
        result <- tryPerformBHRequest $ multiSearchTemplateWith @Tweet opts items
        case result of
          Left _ -> liftIO $ expectationFailure "multi-search template with options failed"
          Right resp -> do
            let responses = multiSearchResponses resp
            liftIO $ responses `shouldSatisfy` all (either (const False) (const True))
            liftIO $ length responses `shouldBe` 2

    it "per-header fields land on the header line and are accepted by ES" $
      withTestEnv $ do
        _ <- insertData
        let item =
              (mkItem (Just testIndex) matchingInlineTemplate)
                { multiSearchTemplateItemPreference = Just "_local",
                  multiSearchTemplateItemSearchType = Just SearchTypeQueryThenFetch,
                  multiSearchTemplateItemAllowPartialSearchResults = Just True
                }
        result <- tryPerformBHRequest $ multiSearchTemplate @Tweet (item :| [])
        case result of
          Left _ -> liftIO $ expectationFailure "multi-search template with header fields failed"
          Right resp -> do
            let responses = multiSearchResponses resp
            liftIO $ responses `shouldSatisfy` all (either (const False) (const True))
            liftIO $ length responses `shouldBe` 1