{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Database.Bloodhound.OpenSearch2.Client
  ( module Reexport,
    pitSearch,
    openPointInTime,
    closePointInTime,
  )
where

import Control.Monad
import Data.Aeson
import Data.Monoid
import Database.Bloodhound.Client.Cluster
import Database.Bloodhound.Common.Client as Reexport
import Database.Bloodhound.Internal.Client.BHRequest
import qualified Database.Bloodhound.OpenSearch2.Requests as Requests
import Database.Bloodhound.OpenSearch2.Types
import Prelude hiding (filter, head)

-- | 'pitSearch' uses the point in time (PIT) API of elastic, for a given
-- 'IndexName'. Note that this will consume the
-- entire search result set and will be doing O(n) list appends so this may
-- not be suitable for large result sets. In that case, the point in time API
-- should be used directly with `openPointInTime` and `closePointInTime`.
--
-- Note that 'pitSearch' utilizes the 'search_after' parameter under the hood,
-- which requires a non-empty 'sortBody' field in the provided 'Search' value.
-- Otherwise, 'pitSearch' will fail to return all matching documents.
--
-- For more information see
-- <https://opensearch.org/docs/latest/search-plugins/point-in-time/>.
pitSearch ::
  forall a m.
  (FromJSON a, MonadBH m, WithBackend OpenSearch2 m) =>
  IndexName ->
  Search ->
  m [Hit a]
pitSearch :: forall a (m :: * -> *).
(FromJSON a, MonadBH m, WithBackend 'OpenSearch2 m) =>
IndexName -> Search -> m [Hit a]
pitSearch IndexName
indexName Search
search = do
  ParsedEsResponse OpenPointInTimeResponse
openResp <- IndexName -> m (ParsedEsResponse OpenPointInTimeResponse)
forall (m :: * -> *).
(MonadBH m, WithBackend 'OpenSearch2 m) =>
IndexName -> m (ParsedEsResponse OpenPointInTimeResponse)
openPointInTime IndexName
indexName
  case ParsedEsResponse OpenPointInTimeResponse
openResp of
    Left EsError
e -> EsError -> m [Hit a]
forall a. EsError -> m a
forall (m :: * -> *) a. MonadBH m => EsError -> m a
throwEsError EsError
e
    Right OpenPointInTimeResponse {Text
POSIXTime
ShardResult
oos2PitId :: Text
oos2Shards :: ShardResult
oos2CreationTime :: POSIXTime
oos2CreationTime :: OpenPointInTimeResponse -> POSIXTime
oos2Shards :: OpenPointInTimeResponse -> ShardResult
oos2PitId :: OpenPointInTimeResponse -> Text
..} -> do
      let searchPIT :: Search
searchPIT = Search
search {pointInTime = Just (PointInTime oos2PitId "1m")}
      [Hit a]
hits <- Search -> [Hit a] -> m [Hit a]
pitAccumulator Search
searchPIT []
      ParsedEsResponse ClosePointInTimeResponse
closeResp <- ClosePointInTime -> m (ParsedEsResponse ClosePointInTimeResponse)
forall (m :: * -> *).
(MonadBH m, WithBackend 'OpenSearch2 m) =>
ClosePointInTime -> m (ParsedEsResponse ClosePointInTimeResponse)
closePointInTime (Text -> ClosePointInTime
ClosePointInTime Text
oos2PitId)
      case ParsedEsResponse ClosePointInTimeResponse
closeResp of
        Left EsError
_ -> [Hit a] -> m [Hit a]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return []
        Right (ClosePointInTimeResponse Bool
False Int
_) ->
          [Char] -> m [Hit a]
forall a. HasCallStack => [Char] -> a
error [Char]
"failed to close point in time (PIT)"
        Right (ClosePointInTimeResponse Bool
True Int
_) -> [Hit a] -> m [Hit a]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return [Hit a]
hits
  where
    pitAccumulator :: Search -> [Hit a] -> m [Hit a]
    pitAccumulator :: Search -> [Hit a] -> m [Hit a]
pitAccumulator Search
search' [Hit a]
oldHits = do
      ParsedEsResponse (SearchResult a)
resp <- BHRequest StatusDependant (SearchResult a)
-> m (ParsedEsResponse (SearchResult a))
forall (m :: * -> *) contextualized a.
(MonadBH m, MonadThrow m, ParseBHResponse contextualized) =>
BHRequest contextualized a -> m (ParsedEsResponse a)
tryPerformBHRequest (BHRequest StatusDependant (SearchResult a)
 -> m (ParsedEsResponse (SearchResult a)))
-> BHRequest StatusDependant (SearchResult a)
-> m (ParsedEsResponse (SearchResult a))
forall a b. (a -> b) -> a -> b
$ Search -> BHRequest StatusDependant (SearchResult a)
forall a.
FromJSON a =>
Search -> BHRequest StatusDependant (SearchResult a)
Requests.searchAll Search
search'
      case ParsedEsResponse (SearchResult a)
resp of
        Left EsError
_ -> [Hit a] -> m [Hit a]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return []
        Right SearchResult a
searchResult -> case SearchHits a -> [Hit a]
forall a. SearchHits a -> [Hit a]
hits (SearchResult a -> SearchHits a
forall a. SearchResult a -> SearchHits a
searchHits SearchResult a
searchResult) of
          [] -> [Hit a] -> m [Hit a]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return [Hit a]
oldHits
          [Hit a]
newHits -> case (Hit a -> Maybe SearchAfterKey
forall a. Hit a -> Maybe SearchAfterKey
hitSort (Hit a -> Maybe SearchAfterKey) -> Hit a -> Maybe SearchAfterKey
forall a b. (a -> b) -> a -> b
$ [Hit a] -> Hit a
forall a. HasCallStack => [a] -> a
last [Hit a]
newHits, SearchResult a -> Maybe Text
forall a. SearchResult a -> Maybe Text
pitId SearchResult a
searchResult) of
            (Maybe SearchAfterKey
Nothing, Maybe Text
Nothing) ->
              [Char] -> m [Hit a]
forall a. HasCallStack => [Char] -> a
error [Char]
"no point in time (PIT) ID or last sort value"
            (Just SearchAfterKey
_, Maybe Text
Nothing) -> [Char] -> m [Hit a]
forall a. HasCallStack => [Char] -> a
error [Char]
"no point in time (PIT) ID"
            (Maybe SearchAfterKey
Nothing, Maybe Text
_) -> [Hit a] -> m [Hit a]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Hit a]
oldHits [Hit a] -> [Hit a] -> [Hit a]
forall a. Semigroup a => a -> a -> a
<> [Hit a]
newHits)
            (Just SearchAfterKey
lastSort, Just Text
pitId') -> do
              let newSearch :: Search
newSearch =
                    Search
search'
                      { pointInTime = Just (PointInTime pitId' "1m"),
                        searchAfterKey = Just lastSort
                      }
              Search -> [Hit a] -> m [Hit a]
pitAccumulator Search
newSearch ([Hit a]
oldHits [Hit a] -> [Hit a] -> [Hit a]
forall a. Semigroup a => a -> a -> a
<> [Hit a]
newHits)

-- | 'openPointInTime' opens a point in time for an index given an 'IndexName'.
-- Note that the point in time should be closed with 'closePointInTime' as soon
-- as it is no longer needed.
--
-- For more information see
-- <https://opensearch.org/docs/latest/search-plugins/point-in-time/>.
openPointInTime ::
  (MonadBH m, WithBackend OpenSearch2 m) =>
  IndexName ->
  m (ParsedEsResponse OpenPointInTimeResponse)
openPointInTime :: forall (m :: * -> *).
(MonadBH m, WithBackend 'OpenSearch2 m) =>
IndexName -> m (ParsedEsResponse OpenPointInTimeResponse)
openPointInTime IndexName
indexName = BHRequest
  StatusDependant (ParsedEsResponse OpenPointInTimeResponse)
-> m (ParsedEsResponse OpenPointInTimeResponse)
forall (m :: * -> *) contextualized a.
(MonadBH m, MonadThrow m, ParseBHResponse contextualized) =>
BHRequest contextualized a -> m a
performBHRequest (BHRequest
   StatusDependant (ParsedEsResponse OpenPointInTimeResponse)
 -> m (ParsedEsResponse OpenPointInTimeResponse))
-> BHRequest
     StatusDependant (ParsedEsResponse OpenPointInTimeResponse)
-> m (ParsedEsResponse OpenPointInTimeResponse)
forall a b. (a -> b) -> a -> b
$ IndexName
-> BHRequest
     StatusDependant (ParsedEsResponse OpenPointInTimeResponse)
Requests.openPointInTime IndexName
indexName

-- | 'closePointInTime' closes a point in time given a 'ClosePointInTime'.
--
-- For more information see
-- <https://opensearch.org/docs/latest/search-plugins/point-in-time/>.
closePointInTime ::
  (MonadBH m, WithBackend OpenSearch2 m) =>
  ClosePointInTime ->
  m (ParsedEsResponse ClosePointInTimeResponse)
closePointInTime :: forall (m :: * -> *).
(MonadBH m, WithBackend 'OpenSearch2 m) =>
ClosePointInTime -> m (ParsedEsResponse ClosePointInTimeResponse)
closePointInTime ClosePointInTime
q = BHRequest
  StatusDependant (ParsedEsResponse ClosePointInTimeResponse)
-> m (ParsedEsResponse ClosePointInTimeResponse)
forall (m :: * -> *) contextualized a.
(MonadBH m, MonadThrow m, ParseBHResponse contextualized) =>
BHRequest contextualized a -> m a
performBHRequest (BHRequest
   StatusDependant (ParsedEsResponse ClosePointInTimeResponse)
 -> m (ParsedEsResponse ClosePointInTimeResponse))
-> BHRequest
     StatusDependant (ParsedEsResponse ClosePointInTimeResponse)
-> m (ParsedEsResponse ClosePointInTimeResponse)
forall a b. (a -> b) -> a -> b
$ ClosePointInTime
-> BHRequest
     StatusDependant (ParsedEsResponse ClosePointInTimeResponse)
Requests.closePointInTime ClosePointInTime
q