content: {-# LANGUAGE CPP #-}
  {-# LANGUAGE OverloadedStrings #-}
  
  module Test.BulkAPISpec (spec) where
  
  import qualified Data.Aeson.KeyMap as X
  import qualified Data.Aeson.Optics as LMA
  import qualified Data.Vector as V
  import Optics
  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
      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)]
                  }
  
          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)]
                  }
  
          -- 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)]
                  }
  
          -- 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)
          -- 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` HitsTotal 6 HTR_EQ
          let nameList :: [Text]
              nameList =
                hits (searchHits sr)
                  ^.. traversed
                  % to hitSource
                  % _Just
                  % LMA.key "name"
                  % LMA._String
  
          liftIO $
            nameList
              `shouldBe` ["blah", "bloo", "graffle", "garabadoo", "serenity"]
  