Common.hs

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

module TestsUtils.Common where

import Control.Exception (SomeException, try)
import Data.Char (toLower)
import Data.List (stripPrefix)
import Data.List qualified as L
import Data.Map qualified as M
import Data.Text qualified as T
import Data.Versions qualified as Versions
import Database.Bloodhound.Client.Cluster (BackendType (..))
import Lens.Micro (toListOf)
import Network.HTTP.Types.Method qualified as NHTM
import Network.HTTP.Types.Status qualified as NHTS
import System.Environment (lookupEnv)
import TestsUtils.Import

testServer :: IO Server
testServer = do
  env <- lookupEnv "ES_TEST_SERVER"
  return $ maybe (Server "http://localhost:9200") (Server . T.pack) env

testIndex :: IndexName
testIndex = [qqIndexName|bloodhound-tests-twitter-1|]

withTestEnv :: BH IO a -> IO a
withTestEnv action = testServer >>= withBH defaultManagerSettings <*> pure action

data Location = Location
  { lat :: Double,
    lon :: Double
  }
  deriving stock (Eq, Show)

data Tweet = Tweet
  { user :: Text,
    postDate :: UTCTime,
    message :: Text,
    age :: Int,
    location :: Location,
    extra :: Maybe Text
  }
  deriving stock (Eq, Show)

$(deriveJSON defaultOptions ''Location)
$(deriveJSON defaultOptions ''Tweet)

data ConversationMapping = ConversationMapping deriving stock (Eq, Show)

instance ToJSON ConversationMapping where
  toJSON ConversationMapping =
    object
      [ "properties"
          .= object
            [ "reply_join"
                .= object
                  [ "type" .= ("join" :: Text),
                    "relations" .= object ["message" .= ("reply" :: Text)]
                  ],
              "user"
                .= object
                  [ "type" .= ("text" :: Text),
                    "fielddata" .= True
                  ],
              -- Serializing the date as a date is breaking other tests, mysteriously.
              -- , "postDate" .= object [ "type"   .= ("date" :: Text)
              --                        , "format" .= ("YYYY-MM-dd`T`HH:mm:ss.SSSZZ" :: Text)]
              "message" .= object ["type" .= ("text" :: Text)],
              "age" .= object ["type" .= ("integer" :: Text)],
              "location" .= object ["type" .= ("geo_point" :: Text)],
              "extra" .= object ["type" .= ("keyword" :: Text)]
            ]
      ]

getServerVersion :: IO Versions.Version
getServerVersion = extractVersion <$> withTestEnv (performBHRequest getStatus)
  where
    extractVersion = versionNumber . number . version

createExampleIndex :: (MonadBH m) => m (BHResponse StatusDependant Acknowledged, Acknowledged)
createExampleIndex = do
  result <- tryPerformBHRequest (keepBHResponse $ createIndex (IndexSettings (ShardCount 1) (ReplicaCount 0) defaultIndexMappingsLimits) testIndex)
  case result of
    Left e
      | T.isSuffixOf "already exists" (errorMessage e) -> return (error "TODO rewrite this part too", Acknowledged False)
      | otherwise -> throwEsError e
    Right ack -> return ack

deleteExampleIndex :: (MonadBH m) => m (BHResponse StatusDependant Acknowledged, Acknowledged)
deleteExampleIndex = do
  result <- tryPerformBHRequest (keepBHResponse $ deleteIndex testIndex)
  case result of
    Left e
      | errorStatus e == Just 404 -> return (error "deleteExampleIndex: 404 has no response body", Acknowledged False)
      | otherwise -> throwEsError e
    Right ack -> return ack

validateStatus :: (Show body) => BHResponse contextualized body -> Int -> Expectation
validateStatus resp expected =
  if actual == expected
    then return ()
    else expectationFailure ("Expected " <> show expected <> " but got " <> show actual <> ": " <> show body)
  where
    actual = NHTS.statusCode (responseStatus $ getResponse resp)
    body = responseBody $ getResponse resp

data TweetMapping = TweetMapping deriving stock (Eq, Show)

instance ToJSON TweetMapping where
  toJSON TweetMapping =
    object
      [ "properties"
          .= object
            [ "user"
                .= object
                  [ "type" .= ("text" :: Text),
                    "fielddata" .= True
                  ],
              -- Serializing the date as a date is breaking other tests, mysteriously.
              -- , "postDate" .= object [ "type"   .= ("date" :: Text)
              --                        , "format" .= ("YYYY-MM-dd`T`HH:mm:ss.SSSZZ" :: Text)]
              "message" .= object ["type" .= ("text" :: Text)],
              "age" .= object ["type" .= ("integer" :: Text)],
              "location" .= object ["type" .= ("geo_point" :: Text)],
              "extra" .= object ["type" .= ("keyword" :: Text)]
            ]
      ]

exampleTweet :: Tweet
exampleTweet =
  Tweet
    { user = "bitemyapp",
      postDate =
        UTCTime
          (ModifiedJulianDay 55000)
          (secondsToDiffTime 10),
      message = "Use haskell!",
      age = 10000,
      location = Location 40.12 (-71.34),
      extra = Nothing
    }

tweetWithExtra :: Tweet
tweetWithExtra =
  Tweet
    { user = "bitemyapp",
      postDate =
        UTCTime
          (ModifiedJulianDay 55000)
          (secondsToDiffTime 10),
      message = "Use haskell!",
      age = 10000,
      location = Location 40.12 (-71.34),
      extra = Just "blah blah"
    }

exampleTweetWithAge :: Int -> Tweet
exampleTweetWithAge age' =
  Tweet
    { user = "bitemyapp",
      postDate =
        UTCTime
          (ModifiedJulianDay 55000)
          (secondsToDiffTime 10),
      message = "Use haskell!",
      age = age',
      location = Location 40.12 (-71.34),
      extra = Nothing
    }

newAge :: Int
newAge = 31337

newUser :: Text
newUser = "someotherapp"

tweetPatch :: Value
tweetPatch =
  object
    [ "age" .= newAge,
      "user" .= newUser
    ]

patchedTweet :: Tweet
patchedTweet = exampleTweet {age = newAge, user = newUser}

otherTweet :: Tweet
otherTweet =
  Tweet
    { user = "notmyapp",
      postDate =
        UTCTime
          (ModifiedJulianDay 55000)
          (secondsToDiffTime 11),
      message = "Use haskell!",
      age = 1000,
      location = Location 40.12 (-71.34),
      extra = Nothing
    }

resetIndex :: BH IO ()
resetIndex = do
  _ <- tryEsError deleteExampleIndex
  _ <- createExampleIndex
  _ <- performBHRequest $ putMapping @Value testIndex TweetMapping
  return ()

insertData :: BH IO (BHResponse StatusDependant IndexedDocument, IndexedDocument)
insertData = do
  _ <- tryEsError resetIndex
  insertData' defaultIndexDocumentSettings

insertData' :: IndexDocumentSettings -> BH IO (BHResponse StatusDependant IndexedDocument, IndexedDocument)
insertData' ids = do
  r <- performBHRequest $ keepBHResponse $ indexDocument testIndex ids exampleTweet (DocId "1")
  _ <- performBHRequest $ refreshIndex testIndex
  return r

-- | Returns the `BHResponse` of `indexDocument` without any parsing
insertData'' :: IndexDocumentSettings -> BH IO (BHResponse StatusDependant IndexedDocument)
insertData'' ids = do
  r <- dispatch $ indexDocument testIndex ids exampleTweet (DocId "1")
  _ <- performBHRequest $ refreshIndex testIndex
  return r

insertTweetWithDocId :: Tweet -> Text -> BH IO IndexedDocument
insertTweetWithDocId tweet docId = do
  let ids = defaultIndexDocumentSettings
  r <- performBHRequest $ indexDocument testIndex ids tweet (DocId docId)
  _ <- performBHRequest $ refreshIndex testIndex
  return r

updateData :: BH IO IndexedDocument
updateData = do
  r <- performBHRequest $ updateDocument testIndex defaultIndexDocumentSettings tweetPatch (DocId "1")
  _ <- performBHRequest $ refreshIndex testIndex
  return r

insertOther :: BH IO ()
insertOther = do
  _ <- performBHRequest $ indexDocument testIndex defaultIndexDocumentSettings otherTweet (DocId "2")
  _ <- performBHRequest $ refreshIndex testIndex
  return ()

insertExtra :: BH IO ()
insertExtra = do
  _ <- performBHRequest $ indexDocument testIndex defaultIndexDocumentSettings tweetWithExtra (DocId "4")
  _ <- performBHRequest $ refreshIndex testIndex
  return ()

insertWithSpaceInId :: BH IO ()
insertWithSpaceInId = do
  _ <- performBHRequest $ indexDocument testIndex defaultIndexDocumentSettings exampleTweet (DocId "Hello World")
  _ <- performBHRequest $ refreshIndex testIndex
  return ()

searchTweet :: Search -> BH IO (Either EsError Tweet)
searchTweet search = (>>= grabFirst) <$> searchTweets search

searchTweets :: Search -> BH IO (Either EsError (SearchResult Tweet))
searchTweets search = tryPerformBHRequest $ searchByIndex testIndex search

searchExpectNoResults :: Search -> BH IO ()
searchExpectNoResults search = do
  result <- searchTweets search
  let emptyHits = fmap (hits . searchHits) result
  liftIO $
    emptyHits `shouldBe` Right []

searchExpectAggs :: Search -> BH IO ()
searchExpectAggs search = do
  result <- performBHRequest $ searchByIndex @Tweet testIndex search
  let isEmpty x = return (M.null x)
  liftIO $
    (aggregations result >>= isEmpty) `shouldBe` Just False

searchValidBucketAgg ::
  (BucketAggregation a, FromJSON a, Show a) =>
  Search ->
  Key ->
  (Key -> AggregationResults -> Maybe (Bucket a)) ->
  BH IO ()
searchValidBucketAgg search aggKey extractor = do
  result <- performBHRequest $ searchByIndex @Tweet testIndex search
  let bucketDocs = docCount . head . buckets
  let count = aggregations result >>= extractor aggKey >>= \x -> return (bucketDocs x)
  liftIO $
    count `shouldBe` Just 1

searchTermsAggHint :: [ExecutionHint] -> BH IO ()
searchTermsAggHint hints = do
  let terms hint = TermsAgg $ (mkTermsAggregation $ FieldName "user.keyword") {termExecutionHint = Just hint}
  let search hint = mkAggregateSearch Nothing $ mkAggregations "users" $ terms hint
  forM_ hints $ searchExpectAggs . search

-- forM_ hints (\x -> searchValidBucketAgg (search x) "users" toTerms)

searchTweetHighlight ::
  Search ->
  BH IO (Either EsError (Maybe HitHighlight))
searchTweetHighlight search = do
  result <- searchTweets search
  let tweetHit :: Either EsError (Maybe (Hit Tweet))
      tweetHit = fmap (headMay . hits . searchHits) result
      myHighlight :: Either EsError (Maybe HitHighlight)
      myHighlight = (join . fmap hitHighlight) <$> tweetHit
  return myHighlight

searchExpectSource :: Source -> Either EsError Value -> BH IO ()
searchExpectSource src expected = do
  _ <- insertData
  let query = QueryMatchQuery $ mkMatchQuery (FieldName "message") (QueryString "haskell")
  let search = (mkSearch (Just query) Nothing) {source = Just src}
  result <- performBHRequest $ searchByIndex testIndex search
  let value_ = grabFirst result
  liftIO $
    value_ `shouldBe` expected

-- ---------------------------------------------------------------------------
-- - k-NN test fixtures
-- ---------------------------------------------------------------------------

-- | Dedicated index for kNN tests. Keeps the dense_vector mapping out of the
-- shared @testIndex@ so the rest of the suite is unaffected.
knnTestIndex :: IndexName
knnTestIndex = [qqIndexName|bloodhound-tests-knn-1|]

-- | 5-dimensional vector document for kNN tests. The vector is small enough
-- to run on any laptop yet exercises the real @dense_vector@ pipeline.
--
-- The vector is wrapped in 'Maybe' because ES9+ strips @dense_vector@ fields
-- from @_source@ on read by default. We only need the vector at index time;
-- assertions only inspect @label@.
data KnnDoc = KnnDoc
  { knnDocLabel :: Text,
    knnDocVector :: Maybe [Double]
  }
  deriving stock (Eq, Show)

-- | Drop the @knnDoc@ prefix and lowercase the next character so that
-- @knnDocLabel@ serializes as @label@ (not @Label@). Inlined into the
-- splice due to Template Haskell stage restrictions.
$(deriveJSON defaultOptions {fieldLabelModifier = \s -> case stripPrefix "knnDoc" s of { Just (c : rest) -> toLower c : rest; _ -> s }} ''KnnDoc)

-- | Mapping declaring @vector@ as an indexed @dense_vector@ with cosine
-- similarity. Required by ES8+\/ES9+ kNN search.
knnDocMapping :: Value
knnDocMapping =
  object
    [ "properties"
        .= object
          [ "label" .= object ["type" .= ("keyword" :: Text)],
            "vector"
              .= object
                [ "type" .= ("dense_vector" :: Text),
                  "dims" .= (5 :: Int),
                  "index" .= True
                ]
          ]
    ]

-- | Sample kNN documents: two vectors that point roughly "toward" the query
-- vector @[0.1, 0.2, 0.3, 0.4, 0.5]@.
exampleKnnDocA :: KnnDoc
exampleKnnDocA = KnnDoc "alpha" (Just [0.1, 0.2, 0.3, 0.4, 0.5])

exampleKnnDocB :: KnnDoc
exampleKnnDocB = KnnDoc "beta" (Just [0.9, 0.8, 0.7, 0.6, 0.5])

-- | Reset (delete + recreate + map) the kNN test index.
resetKnnIndex :: BH IO ()
resetKnnIndex = do
  _ <- tryEsError deleteKnnExampleIndex
  _ <- createKnnExampleIndex
  _ <- performBHRequest $ putMapping @Value knnTestIndex knnDocMapping
  return ()

deleteKnnExampleIndex :: BH IO (BHResponse StatusDependant Acknowledged, Acknowledged)
deleteKnnExampleIndex =
  performBHRequest $ keepBHResponse $ deleteIndex knnTestIndex

createKnnExampleIndex :: BH IO (BHResponse StatusDependant Acknowledged, Acknowledged)
createKnnExampleIndex = do
  result <- tryPerformBHRequest (keepBHResponse $ createIndex (IndexSettings (ShardCount 1) (ReplicaCount 0) defaultIndexMappingsLimits) knnTestIndex)
  case result of
    Left e
      | T.isSuffixOf "already exists" (errorMessage e) -> return (error "TODO rewrite this part too", Acknowledged False)
      | otherwise -> throwEsError e
    Right ack -> return ack

-- | Seed the kNN index with two known documents.
insertKnnData :: BH IO ()
insertKnnData = do
  _ <- tryEsError resetKnnIndex
  _ <- performBHRequest $ indexDocument knnTestIndex defaultIndexDocumentSettings exampleKnnDocA (DocId "a")
  _ <- performBHRequest $ indexDocument knnTestIndex defaultIndexDocumentSettings exampleKnnDocB (DocId "b")
  _ <- performBHRequest $ refreshIndex knnTestIndex
  return ()

-- ---------------------------------------------------------------------------
-- - OpenSearch k-NN test fixtures
-- ---------------------------------------------------------------------------

-- | Dedicated index for OpenSearch kNN tests. OS uses a different field type
-- (@knn_vector@ vs ES's @dense_vector@) and requires @index.knn: true@.
osKnnTestIndex :: IndexName
osKnnTestIndex = [qqIndexName|bloodhound-tests-os-knn-1|]

-- | Mapping declaring @vector@ as an OpenSearch @knn_vector@. The
-- @index.knn: true@ setting is applied via 'osKnnIndexSettings'.
osKnnDocMapping :: Value
osKnnDocMapping =
  object
    [ "properties"
        .= object
          [ "label" .= object ["type" .= ("keyword" :: Text)],
            "vector"
              .= object
                [ "type" .= ("knn_vector" :: Text),
                  "dimension" .= (5 :: Int)
                ]
          ]
    ]

-- | OpenSearch-specific index settings that enable the kNN plugin.
osKnnIndexSettings :: Value
osKnnIndexSettings =
  object
    [ "settings"
        .= object ["index.knn" .= True],
      "mappings" .= osKnnDocMapping
    ]

-- | Reset the OpenSearch kNN test index. Uses a raw JSON body because the
-- common 'IndexSettings' type does not model @index.knn@.
resetOsKnnIndex :: BH IO ()
resetOsKnnIndex = do
  _ <- tryEsError deleteOsKnnExampleIndex
  _ <- createOsKnnExampleIndex
  return ()

deleteOsKnnExampleIndex :: BH IO (BHResponse StatusDependant Acknowledged, Acknowledged)
deleteOsKnnExampleIndex =
  performBHRequest $ keepBHResponse $ deleteIndex osKnnTestIndex

createOsKnnExampleIndex :: BH IO (BHResponse StatusDependant Value, Value)
createOsKnnExampleIndex = do
  result <-
    tryPerformBHRequest $
      keepBHResponse $
        mkFullRequest @StatusDependant @Value NHTM.methodPut (mkEndpoint [unIndexName osKnnTestIndex]) (encode osKnnIndexSettings)
  case result of
    Left e
      | T.isSuffixOf "already exists" (errorMessage e) -> return (error "TODO rewrite this part too", object ["acknowledged" .= True])
      | otherwise -> throwEsError e
    Right ack -> return ack

-- | Seed the OS kNN index with two known documents.
insertOsKnnData :: BH IO ()
insertOsKnnData = do
  _ <- tryEsError resetOsKnnIndex
  _ <- performBHRequest $ indexDocument osKnnTestIndex defaultIndexDocumentSettings exampleKnnDocA (DocId "a")
  _ <- performBHRequest $ indexDocument osKnnTestIndex defaultIndexDocumentSettings exampleKnnDocB (DocId "b")
  _ <- performBHRequest $ refreshIndex osKnnTestIndex
  return ()

-- ---------------------------------------------------------------------------
-- - OpenSearch neural-sparse test fixtures
-- ---------------------------------------------------------------------------

-- | Dedicated index for OpenSearch neural-sparse warmup/cache tests. OS 3.x
-- requires @index.sparse: true@ for the neural warmup and clear-cache
-- endpoints to accept an index.
osNeuralSparseTestIndex :: IndexName
osNeuralSparseTestIndex = [qqIndexName|bloodhound-tests-os-neural-sparse-1|]

-- | Mapping declaring @sparse_embedding@ as an OpenSearch @sparse_vector@
-- field. OS 3.x requires a @method@ parameter on @sparse_vector@ fields.
osNeuralSparseDocMapping :: Value
osNeuralSparseDocMapping =
  object
    [ "properties"
        .= object
          [ "sparse_embedding"
              .= object
                [ "type" .= ("sparse_vector" :: Text),
                  "method" .= object ["name" .= ("seismic" :: Text)]
                ]
          ]
    ]

-- | OpenSearch-specific index settings that enable neural sparse search.
-- The @index.sparse@ setting is required for the warmup and clear-cache
-- endpoints (live-verified on OS 3.7.0, bead bloodhound-at5).
osNeuralSparseIndexSettings :: Value
osNeuralSparseIndexSettings =
  object
    [ "settings"
        .= object ["index" .= object ["sparse" .= True]],
      "mappings" .= osNeuralSparseDocMapping
    ]

-- | Reset the OpenSearch neural-sparse test index. Uses a raw JSON body
-- because the common 'IndexSettings' type does not model @index.sparse@.
resetOsNeuralSparseIndex :: BH IO ()
resetOsNeuralSparseIndex = do
  _ <- tryEsError deleteOsNeuralSparseExampleIndex
  _ <- createOsNeuralSparseExampleIndex
  return ()

deleteOsNeuralSparseExampleIndex :: BH IO (BHResponse StatusDependant Acknowledged, Acknowledged)
deleteOsNeuralSparseExampleIndex =
  performBHRequest $ keepBHResponse $ deleteIndex osNeuralSparseTestIndex

createOsNeuralSparseExampleIndex :: BH IO (BHResponse StatusDependant Value, Value)
createOsNeuralSparseExampleIndex = do
  result <-
    tryPerformBHRequest $
      keepBHResponse $
        mkFullRequest @StatusDependant @Value NHTM.methodPut (mkEndpoint [unIndexName osNeuralSparseTestIndex]) (encode osNeuralSparseIndexSettings)
  case result of
    Left e
      | T.isSuffixOf "already exists" (errorMessage e) -> return (error "TODO rewrite this part too", object ["acknowledged" .= True])
      | otherwise -> throwEsError e
    Right ack -> return ack

is :: Versions.Version -> IO Bool
is v = getServerVersion >>= \x -> return $ x == v

esOnlyIT :: (HasCallStack, Example a) => IO (String -> a -> SpecWith (Arg a))
esOnlyIT = withMajorVersionIT (>= 6)

es9OnlyIT :: (HasCallStack, Example a) => IO (String -> a -> SpecWith (Arg a))
es9OnlyIT = withMajorVersionIT (== 9)

os1OnlyIT :: (HasCallStack, Example a) => IO (String -> a -> SpecWith (Arg a))
os1OnlyIT = withMajorVersionIT (== 1)

os2OnlyIT :: (HasCallStack, Example a) => IO (String -> a -> SpecWith (Arg a))
os2OnlyIT = withMajorVersionIT (== 2)

os3OnlyIT :: (HasCallStack, Example a) => IO (String -> a -> SpecWith (Arg a))
os3OnlyIT = withMajorVersionIT (== 3)

withMajorVersionIT :: (HasCallStack, Example a) => (Word -> Bool) -> IO (String -> a -> SpecWith (Arg a))
withMajorVersionIT p = do
  majoreVersion <- fetchMajorVersion
  return $
    if p majoreVersion
      then it
      else xit

fetchMajorVersion :: IO Word
fetchMajorVersion =
  withTestEnv $ do
    x <- performBHRequest $ getNodesInfo LocalNode
    let majoreVersion = versionNumber $ nodeInfoESVersion $ head $ nodesInfo x
    return $ head $ toListOf Versions.major majoreVersion

-- | Detect backend type from the running server
-- | Returns Nothing if server is unreachable or unknown
detectBackendType :: IO (Maybe BackendType)
detectBackendType = do
  result <- Control.Exception.try (withTestEnv $ performBHRequest $ getNodesInfo LocalNode) :: IO (Either Control.Exception.SomeException NodesInfo)
  case result of
    Left _ -> pure Nothing
    Right nodesResponse -> do
      let versionNum = versionNumber $ nodeInfoESVersion $ head $ nodesInfo nodesResponse
      let major = head $ toListOf Versions.major versionNum
      return $ case major of
        7 -> Just ElasticSearch7
        8 -> Just ElasticSearch8
        9 -> Just ElasticSearch9
        1 -> Just OpenSearch1
        2 -> Just OpenSearch2
        3 -> Just OpenSearch3
        _ -> Nothing

-- | Run spec only if current backend matches one of the supported backends.
-- | If not, all tests are marked pending with an informative message.
backendSpecific :: (HasCallStack) => [BackendType] -> SpecWith a -> SpecWith a
backendSpecific supported spec = do
  detected <- runIO detectBackendType
  case detected of
    Just backend
      | backend `elem` supported ->
          spec
    Just other ->
      before_ (pendingWith $ "Requires " <> formatBackends supported <> " but running on " <> show other) spec
    Nothing ->
      before_ (pendingWith "Could not detect backend (server unreachable?)") spec
  where
    formatBackends = L.intercalate ", " . map show