packages feed

bloodhound-1.0.0.0: tests/Test/BulkAPISpec.hs

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

module Test.BulkAPISpec (spec) where

import Data.Aeson.KeyMap qualified as X
import Data.Function ((&))
import Data.Functor ((<&>))
import Data.Vector qualified as V
import Data.Word (Word64)
import TestsUtils.Common
import TestsUtils.Import

newtype BulkTest
  = BulkTest Text
  deriving stock (Eq, Show)

instance ToJSON BulkTest where
  toJSON (BulkTest name') =
    object ["name" .= name']

instance FromJSON BulkTest where
  parseJSON = withObject "BulkTest" parse
    where
      parse o = do
        t <- o .: "name"
        BulkTest <$> parseJSON t

data BulkScriptTest = BulkScriptTest
  { bstName :: Text,
    bstCounter :: Int
  }
  deriving stock (Eq, Show)

instance ToJSON BulkScriptTest where
  toJSON (BulkScriptTest name' count) =
    object ["name" .= name', "counter" .= count]

instance FromJSON BulkScriptTest where
  parseJSON = withObject "BulkScriptTest" $ \v ->
    BulkScriptTest
      <$> v
        .: "name"
      <*> v
        .: "counter"

assertDocs :: (FromJSON a, Show a, Eq a) => [(DocId, a)] -> BH IO ()
assertDocs as = do
  let (ids, docs) = unzip as
  res <- traverse (fmap getSource . performBHRequest . getDocument testIndex) ids
  liftIO $ res `shouldBe` map Just docs

upsertDocs ::
  (ToJSON a, Show a, Eq a) =>
  (Value -> UpsertPayload) ->
  [(DocId, a)] ->
  BH IO ()
upsertDocs f as = do
  let batch = as <&> (\(id_, doc) -> BulkUpsert testIndex id_ (f $ toJSON doc) []) & V.fromList
  _ <- performBHRequest (bulk @StatusDependant batch)
  _ <- performBHRequest (refreshIndex testIndex)
  pure ()

spec :: Spec
spec =
  describe "Bulk API" $ do
    describe "BulkItem response fields" $ do
      let minimalItemJson :: Value
          minimalItemJson =
            object
              [ "_index" .= ("test-idx" :: Text),
                "_id" .= ("1" :: Text),
                "status" .= (404 :: Int),
                "error"
                  .= object
                    [ "type" .= ("document_missing_exception" :: Text),
                      "reason" .= ("doc not found" :: Text)
                    ]
              ]

          fullItemJson :: Value
          fullItemJson =
            object
              [ "_index" .= ("test-idx" :: Text),
                "_id" .= ("1" :: Text),
                "status" .= (201 :: Int),
                "_seq_no" .= (0 :: Int),
                "_primary_term" .= (1 :: Int),
                "_version" .= (1 :: Int),
                "_shards"
                  .= object
                    [ "total" .= (1 :: Int),
                      "successful" .= (1 :: Int),
                      "skipped" .= (0 :: Int),
                      "failed" .= (0 :: Int)
                    ],
                "result" .= ("created" :: Text)
              ]

          expectedMinimal :: BulkItem
          expectedMinimal =
            BulkItem
              { biIndex = "test-idx",
                biId = Just "1",
                biStatus = Just 404,
                biError =
                  Just $
                    BulkError
                      { beType = "document_missing_exception",
                        beReason = "doc not found"
                      },
                biSeqNo = Nothing,
                biPrimaryTerm = Nothing,
                biVersion = Nothing,
                biShards = Nothing,
                biResult = Nothing
              }

          expectedFull :: BulkItem
          expectedFull =
            BulkItem
              { biIndex = "test-idx",
                biId = Just "1",
                biStatus = Just 201,
                biError = Nothing,
                biSeqNo = Just 0,
                biPrimaryTerm = Just 1,
                biVersion = Just 1,
                biShards =
                  Just $
                    ShardResult
                      { shardTotal = 1,
                        shardsSuccessful = 1,
                        shardsSkipped = 0,
                        shardsFailed = 0
                      },
                biResult = Just "created"
              }

      it "decodes a fully-populated per-item result" $
        decode (encode fullItemJson) `shouldBe` Just expectedFull

      it "decodes a minimal per-item result (new fields default to Nothing)" $
        decode (encode minimalItemJson) `shouldBe` Just expectedMinimal

      it "parses _shards into ShardResult with skipped/failed present" $ do
        let jsonWithSkippedFailed :: Value
            jsonWithSkippedFailed =
              object
                [ "_index" .= ("test-idx" :: Text),
                  "_id" .= ("1" :: Text),
                  "_shards"
                    .= object
                      [ "total" .= (3 :: Int),
                        "successful" .= (2 :: Int),
                        "skipped" .= (1 :: Int),
                        "failed" .= (0 :: Int)
                      ]
                ]
        case decode (encode jsonWithSkippedFailed) of
          Just (bi :: BulkItem) ->
            biShards bi
              `shouldBe` Just
                ShardResult
                  { shardTotal = 3,
                    shardsSuccessful = 2,
                    shardsSkipped = 1,
                    shardsFailed = 0
                  }
          _ -> expectationFailure "expected BulkItem decode to succeed"

      it "propagates per-item shard failures via biShards (failed > 0)" $ do
        -- This is the motivating bug for the bead: callers previously could
        -- not see per-item shard failures. Verify shardsFailed round-trips.
        let jsonWithFailedShard :: Value
            jsonWithFailedShard =
              object
                [ "_index" .= ("test-idx" :: Text),
                  "_id" .= ("1" :: Text),
                  "status" .= (200 :: Int),
                  "_version" .= (3 :: Int),
                  "result" .= ("updated" :: Text),
                  "_shards"
                    .= object
                      [ "total" .= (2 :: Int),
                        "successful" .= (1 :: Int),
                        "skipped" .= (0 :: Int),
                        "failed" .= (1 :: Int)
                      ]
                ]
        case decode (encode jsonWithFailedShard) of
          Just (bi :: BulkItem) ->
            (biShards bi >>= Just . shardsFailed) `shouldBe` Just 1
          _ -> expectationFailure "expected BulkItem decode to succeed"

      it "decodes a delete result (no _shards, no _seq_no, no _primary_term)" $ do
        -- ES bulk delete items return _version and result="deleted" but
        -- typically omit _shards, _seq_no, _primary_term. The new fields
        -- must decode safely as Nothing.
        let deleteItemJson :: Value
            deleteItemJson =
              object
                [ "_index" .= ("test-idx" :: Text),
                  "_id" .= ("1" :: Text),
                  "status" .= (200 :: Int),
                  "_version" .= (7 :: Int),
                  "result" .= ("deleted" :: Text)
                ]
        decode (encode deleteItemJson)
          `shouldBe` Just
            BulkItem
              { biIndex = "test-idx",
                biId = Just "1",
                biStatus = Just 200,
                biError = Nothing,
                biSeqNo = Nothing,
                biPrimaryTerm = Nothing,
                biVersion = Just 7,
                biShards = Nothing,
                biResult = Just "deleted"
              }

    it "upsert operations" $
      withTestEnv $ do
        _ <- insertData

        -- Upserting in a to a fresh index should insert
        let toInsert = [(DocId "3", BulkTest "stringer"), (DocId "5", BulkTest "sobotka"), (DocId "7", BulkTest "snoop")]
        upsertDocs UpsertDoc toInsert
        assertDocs toInsert

        -- Upserting existing documents should update
        let toUpsert = [(DocId "3", BulkTest "bell"), (DocId "5", BulkTest "frank"), (DocId "7", BulkTest "snoop")]
        upsertDocs UpsertDoc toUpsert
        assertDocs toUpsert

    it "upsert with a script" $
      withTestEnv $ do
        _ <- insertData

        -- first insert the batch
        let batch = [(DocId "3", BulkScriptTest "stringer" 0), (DocId "5", BulkScriptTest "sobotka" 3)]
        upsertDocs UpsertDoc batch

        -- then upsert with the script

        let script =
              Script
                { scriptLanguage = Just $ ScriptLanguage "painless",
                  scriptSource = ScriptInline "ctx._source.counter += params.count",
                  scriptParams = Just $ ScriptParams $ X.fromList [("count", Number 2)],
                  scriptOptions = Nothing
                }

        upsertDocs (UpsertScript False script) batch
        assertDocs (batch <&> (\(i, v) -> (i, v {bstCounter = bstCounter v + 2})))

    it "script upsert without scripted_upsert" $
      withTestEnv $ do
        _ <- insertData

        let batch = [(DocId "3", BulkScriptTest "stringer" 0), (DocId "5", BulkScriptTest "sobotka" 3)]

        let script =
              Script
                { scriptLanguage = Just $ ScriptLanguage "painless",
                  scriptSource = ScriptInline "ctx._source.counter += params.count",
                  scriptParams = Just $ ScriptParams $ X.fromList [("count", Number 2)],
                  scriptOptions = Nothing
                }

        -- Without "script_upsert" flag new documents are simply inserted and are not handled by the script
        upsertDocs (UpsertScript False script) batch
        assertDocs batch

    it "script upsert with scripted_upsert -- will fail if a bug on elasticsearch is fix, delete patch line" $
      withTestEnv $ do
        _ <- insertData

        let batch = [(DocId "3", BulkScriptTest "stringer" 0), (DocId "5", BulkScriptTest "sobotka" 3)]

        let script =
              Script
                { scriptLanguage = Just $ ScriptLanguage "painless",
                  scriptSource = ScriptInline "ctx._source.counter += params.count",
                  scriptParams = Just $ ScriptParams $ X.fromList [("count", Number 2)],
                  scriptOptions = Nothing
                }

        -- Without "script_upsert" flag new documents are simply inserted and are not handled by the script
        upsertDocs (UpsertScript True script) batch

        -- if this test fails due to a bug in ES7: https://github.com/elastic/elasticsearch/issues/48670, delete next line when it is solved.
        assertDocs (batch <&> (\(i, v) -> (i, v {bstCounter = bstCounter v + 2})))

    it "inserts all documents we request" $
      withTestEnv $ do
        _ <- insertData
        let firstTest = BulkTest "blah"
        let secondTest = BulkTest "bloo"
        let thirdTest = BulkTest "graffle"
        let fourthTest = BulkTest "garabadoo"
        let fifthTest = BulkTest "serenity"
        let firstDoc = BulkIndex testIndex (DocId "2") (toJSON firstTest) []
        let secondDoc = BulkCreate testIndex (DocId "3") (toJSON secondTest) []
        let thirdDoc = BulkCreateEncoding testIndex (DocId "4") (toEncoding thirdTest) []
        let fourthDoc = BulkIndexAuto testIndex (toJSON fourthTest) []
        let fifthDoc = BulkIndexEncodingAuto testIndex (toEncoding fifthTest) []
        let stream = V.fromList [firstDoc, secondDoc, thirdDoc, fourthDoc, fifthDoc]
        bulkResponse <- performBHRequest $ bulk @StatusDependant stream
        -- liftIO $ pPrint bulkResp
        _ <- performBHRequest $ refreshIndex testIndex
        -- liftIO $ pPrint refreshResp
        maybeFirst <- performBHRequest $ getDocument @BulkTest testIndex (DocId "2")
        maybeSecond <- performBHRequest $ getDocument @BulkTest testIndex (DocId "3")
        maybeThird <- performBHRequest $ getDocument @BulkTest testIndex (DocId "4")
        -- note that we cannot query for fourthDoc and fifthDoc since we
        -- do not know their autogenerated ids.
        -- liftIO $ pPrint [maybeFirst, maybeSecond, maybeThird]
        liftIO $ do
          getSource maybeFirst `shouldBe` Just firstTest
          getSource maybeSecond `shouldBe` Just secondTest
          getSource maybeThird `shouldBe` Just thirdTest
          bulkErrors bulkResponse `shouldBe` False
          length (bulkActionItems bulkResponse) `shouldBe` 5
          baiAction <$> bulkActionItems bulkResponse `shouldBe` [Index, Create, Create, Index, Index]
          bulkActionItems bulkResponse `shouldSatisfy` all ((== "bloodhound-tests-twitter-1") . biIndex . baiItem)
          bulkActionItems bulkResponse `shouldSatisfy` all ((== Just 201) . biStatus . baiItem)
          bulkActionItems bulkResponse `shouldSatisfy` all ((== Nothing) . biError . baiItem)
          -- Per-item enrichment: a successful create/update carries
          -- _seq_no, _primary_term, _version, _shards, and result on every backend.
          bulkActionItems bulkResponse `shouldSatisfy` all (isJust . biSeqNo . baiItem)
          bulkActionItems bulkResponse `shouldSatisfy` all (isJust . biPrimaryTerm . baiItem)
          bulkActionItems bulkResponse `shouldSatisfy` all (isJust . biVersion . baiItem)
          bulkActionItems bulkResponse `shouldSatisfy` all (isJust . biShards . baiItem)
          bulkActionItems bulkResponse `shouldSatisfy` all (isJust . biResult . baiItem)
          -- The first op is BulkIndex (upsert) on a fresh doc id, the next two
          -- are creates, the last two are index-auto (also upserts).
          (biResult . baiItem <$> bulkActionItems bulkResponse)
            `shouldSatisfy` all (`elem` [Just "created", Just "updated"])
          (biVersion . baiItem <$> bulkActionItems bulkResponse)
            `shouldSatisfy` all (>= Just 1)
        -- Since we can't get the docs by doc id, we check for their existence in
        -- a match all query.
        let query = MatchAllQuery Nothing
        let search = mkSearch (Just query) Nothing
        sr <- performBHRequest $ searchByIndex @Value testIndex search
        liftIO $
          hitsTotal (searchHits sr) `shouldBe` Just (HitsTotal 6 HTR_EQ)
        let nameList :: [Text]
            nameList =
              [ name'
              | h <- hits (searchHits sr),
                Just v <- [hitSource h],
                Object o <- [v],
                Just (String name') <- [X.lookup "name" o]
              ]

        liftIO $
          nameList
            `shouldBe` ["blah", "bloo", "graffle", "garabadoo", "serenity"]

    it "bulkWith passes URI params (refresh=wait_for) to the server" $
      withTestEnv $ do
        _ <- insertData
        let doc = BulkTest "bulkwith-refresh"
        -- bsRefresh = RefreshWaitFor forces the server to acknowledge
        -- a refresh before responding, so the doc is immediately
        -- searchable without an explicit refreshIndex call.
        let op = BulkIndex testIndex (DocId "bw-1") (toJSON doc) []
        _ <-
          performBHRequest $
            bulkWith @StatusDependant defaultBulkSettings {bsRefresh = Just RefreshWaitFor} (V.singleton op)
        -- No refreshIndex here — bsRefresh did the work.
        maybeDoc <- performBHRequest $ getDocument @BulkTest testIndex (DocId "bw-1")
        liftIO $ getSource maybeDoc `shouldBe` Just doc

    it "bulkWith respects per-op _routing metadata end-to-end" $
      withTestEnv $ do
        _ <- insertData
        let doc = BulkTest "routing-end-to-end"
        -- Custom routing places the doc on a specific shard; a follow-up
        -- getDocument *without* routing would miss it on a multi-shard
        -- index. testIndex is single-shard (see TestsUtils.Common), so we
        -- verify the round-trip via the bulk response rather than shard
        -- placement: the response should report no errors and a created
        -- result.
        let op =
              BulkIndex
                testIndex
                (DocId "rt-1")
                (toJSON doc)
                [UA_Routing "custom-route-1"]
        bulkResponse <-
          performBHRequest $
            bulkWith @StatusDependant defaultBulkSettings (V.singleton op)
        _ <- performBHRequest $ refreshIndex testIndex
        liftIO $ do
          bulkErrors bulkResponse `shouldBe` False
          length (bulkActionItems bulkResponse) `shouldBe` 1
          (biResult . baiItem <$> bulkActionItems bulkResponse)
            `shouldSatisfy` all (`elem` [Just "created", Just "updated"])

    describe "Per-op metadata (UpsertActionMetadata)" $ do
      it "renders every constructor to the expected wire pair" $ do
        buildUpsertActionMetadata (UA_Routing "user-1")
          `shouldBe` "routing"
          .= ("user-1" :: Text)
        buildUpsertActionMetadata (UA_Version 7)
          `shouldBe` "version"
          .= (7 :: Int)
        buildUpsertActionMetadata (UA_VersionType VersionTypeExternalGTE)
          `shouldBe` "version_type"
          .= ("external_gte" :: Text)
        buildUpsertActionMetadata (UA_IfSeqNo 12)
          `shouldBe` "if_seq_no"
          .= (12 :: Word64)
        buildUpsertActionMetadata (UA_IfPrimaryTerm 3)
          `shouldBe` "if_primary_term"
          .= (3 :: Word64)
        buildUpsertActionMetadata (UA_RequireAlias True)
          `shouldBe` "require_alias"
          .= True
        buildUpsertActionMetadata (UA_RetryOnConflict 5)
          `shouldBe` "retry_on_conflict"
          .= (5 :: Int)

      it "encodes a BulkIndex op with metadata into the correct NDJSON line" $ do
        let op =
              BulkIndex
                [qqIndexName|twitter|]
                (DocId "1")
                (toJSON (BulkTest "blah"))
                [ UA_Routing "user-1",
                  UA_IfSeqNo 5,
                  UA_IfPrimaryTerm 2,
                  UA_VersionType VersionTypeExternal
                ]
        -- aeson sorts object keys alphabetically, so the rendered line is stable.
        encodeBulkOperation op
          `shouldBe` "{\"index\":{\"_id\":\"1\",\"_index\":\"twitter\",\"if_primary_term\":2,\"if_seq_no\":5,\"routing\":\"user-1\",\"version_type\":\"external\"}}\n{\"name\":\"blah\"}"

      it "encodes a BulkDelete op with metadata into the correct NDJSON line" $ do
        let op =
              BulkDelete
                [qqIndexName|twitter|]
                (DocId "1")
                [UA_Routing "user-1"]
        encodeBulkOperation op
          `shouldBe` "{\"delete\":{\"_id\":\"1\",\"_index\":\"twitter\",\"routing\":\"user-1\"}}"

      it "encodes a BulkUpdate op with metadata into the correct NDJSON line" $ do
        let op =
              BulkUpdate
                [qqIndexName|twitter|]
                (DocId "1")
                (toJSON (BulkTest "blah"))
                [UA_RequireAlias True]
        encodeBulkOperation op
          `shouldBe` "{\"update\":{\"_id\":\"1\",\"_index\":\"twitter\",\"require_alias\":true}}\n{\"doc\":{\"name\":\"blah\"}}"

      it "encodes a BulkIndexAuto op with metadata into the correct NDJSON line" $ do
        let op =
              BulkIndexAuto
                [qqIndexName|twitter|]
                (toJSON (BulkTest "blah"))
                [UA_Routing "user-1"]
        encodeBulkOperation op
          `shouldBe` "{\"index\":{\"_index\":\"twitter\",\"routing\":\"user-1\"}}\n{\"name\":\"blah\"}"

      it "encodes a BulkCreateEncoding op with metadata into the correct NDJSON line" $ do
        let op =
              BulkCreateEncoding
                [qqIndexName|twitter|]
                (DocId "1")
                (toEncoding (BulkTest "blah"))
                [UA_IfPrimaryTerm 1, UA_IfSeqNo 2]
        encodeBulkOperation op
          `shouldBe` "{\"create\":{\"_id\":\"1\",\"_index\":\"twitter\",\"if_primary_term\":1,\"if_seq_no\":2}}\n{\"name\":\"blah\"}"

      it "is byte-for-byte identical whether metadata is [] or omitted-style (empty list)" $ do
        -- Regression: pre-refactor BulkIndex ix id val produced exactly
        -- this string; the refactored BulkIndex ix id val [] must match.
        let opOldShape = BulkIndex [qqIndexName|twitter|] (DocId "2") (toJSON (BulkTest "blah")) []
        encodeBulkOperation opOldShape
          `shouldBe` "{\"index\":{\"_id\":\"2\",\"_index\":\"twitter\"}}\n{\"name\":\"blah\"}"

    describe "BulkSettings (URI params)" $ do
      it "defaultBulkSettings emits no query parameters" $
        bulkSettingsParams defaultBulkSettings `shouldBe` []

      it "renders refresh / wait_for_active_shards as bare values" $ do
        let opts =
              defaultBulkSettings
                { bsRefresh = Just RefreshWaitFor,
                  bsWaitForActiveShards = Just AllActiveShards
                }
        bulkSettingsParams opts
          `shouldBe` [("refresh", Just "wait_for"), ("wait_for_active_shards", Just "all")]

      it "renders booleans as true/false strings" $ do
        let opts =
              defaultBulkSettings
                { bsSource = Just False,
                  bsRequireAlias = Just True
                }
        bulkSettingsParams opts
          `shouldBe` [("_source", Just "false"), ("require_alias", Just "true")]

      it "renders text params verbatim" $ do
        let opts =
              defaultBulkSettings
                { bsTimeout = Just "5s",
                  bsRouting = Just "user-1",
                  bsPipeline = Just "enrich",
                  bsSourceIncludes = Just "metadata.*,name",
                  bsSourceExcludes = Just "internal.*"
                }
            params = bulkSettingsParams opts
        -- Order of the rendered list is unspecified (callers should treat
        -- it as a set), so verify per-key rather than as a sublist.
        lookup "timeout" params `shouldBe` Just (Just "5s")
        lookup "routing" params `shouldBe` Just (Just "user-1")
        lookup "pipeline" params `shouldBe` Just (Just "enrich")
        lookup "_source_includes" params `shouldBe` Just (Just "metadata.*,name")
        lookup "_source_excludes" params `shouldBe` Just (Just "internal.*")