Cluster.hs

{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

module Database.Bloodhound.Client.Cluster
  ( module ReexportCompat,
    BH (..),
    BHEnv (..),
    MonadBH (..),
    BackendType (..),
    WithBackend,
    WithBackendType,
    StaticBH (..),
    SBackendType (..),
    emptyBody,
    mkBHEnv,
    mkHttpRequest,
    performBHRequest,
    runBH,
    tryPerformBHRequest,
    unsafePerformBH,
    withDynamicBH,
    withStreamingResponse,
  )
where

import Control.Monad.Catch
import Control.Monad.Except
import Data.ByteString.Lazy.Char8 qualified as L
import Data.Kind
import Data.Text qualified as T
import Database.Bloodhound.Internal.Client.BHRequest
import Database.Bloodhound.Internal.Utils.Imports
import Database.Bloodhound.Internal.Versions.Common.Types.Bulk as ReexportCompat
import Database.Bloodhound.Internal.Versions.Common.Types.Count as ReexportCompat
import Database.Bloodhound.Internal.Versions.Common.Types.Indices as ReexportCompat
import Database.Bloodhound.Internal.Versions.Common.Types.Nodes as ReexportCompat
import Database.Bloodhound.Internal.Versions.Common.Types.Snapshots as ReexportCompat
import Database.Bloodhound.Internal.Versions.Common.Types.Units as ReexportCompat
import Network.HTTP.Client qualified as HTTP
import Network.URI qualified as URI

-- | Common environment for Elasticsearch calls. Connections will be
--   pipelined according to the provided HTTP connection manager.
data BHEnv = BHEnv
  { bhServer :: Server,
    bhManager :: HTTP.Manager,
    -- | Low-level hook that is run before every request is sent. Used to implement custom authentication strategies. Defaults to 'return' with 'mkBHEnv'.
    bhRequestHook :: HTTP.Request -> IO HTTP.Request,
    -- | Low-level hook that is run after every response is received. Used to debug. Defaults to 'return' with 'mkBHEnv'.
    bhResponseHook :: HTTP.Response L.ByteString -> IO (HTTP.Response L.ByteString)
  }

-- | All API calls to Elasticsearch operate within
--   MonadBH
--   . The idea is that it can be easily embedded in your
--   own monad transformer stack. A default instance for a ReaderT and
--   alias 'BH' is provided for the simple case.
class (Functor m, Applicative m, MonadIO m, MonadCatch m) => MonadBH m where
  type Backend m :: BackendType
  dispatch :: BHRequest contextualized body -> m (BHResponse contextualized body)
  tryEsError :: m a -> m (Either EsError a)
  throwEsError :: EsError -> m a

-- | Backend (i.e. implementation) the queries are ran against
data BackendType
  = ElasticSearch7
  | ElasticSearch8
  | ElasticSearch9
  | OpenSearch1
  | OpenSearch2
  | OpenSearch3
  | -- | unknown, can be anything
    Dynamic
  deriving stock (Eq, Show)

-- | Best-effort, by-passed for 'Dynamic', statically enforced implementation
type family WithBackendType (expected :: BackendType) (actual :: BackendType) :: Constraint where
  WithBackendType e 'Dynamic = ()
  WithBackendType e a = e ~ a

-- | Helper for 'WithBackendType'
type WithBackend (backend :: BackendType) (m :: Type -> Type) = WithBackendType backend (Backend m)

-- | Create a 'BHEnv' with all optional fields defaulted. HTTP hook
-- will be a noop. You can use the exported fields to customize
-- it further, e.g.:
--
-- >> (mkBHEnv myServer myManager) { bhRequestHook = customHook }
mkBHEnv :: Server -> HTTP.Manager -> BHEnv
mkBHEnv s m = BHEnv s m return return

-- | Basic BH implementation
newtype BH m a = BH
  { unBH :: ReaderT BHEnv (ExceptT EsError m) a
  }
  deriving newtype
    ( Functor,
      Applicative,
      Monad,
      MonadIO,
      MonadState s,
      MonadWriter w,
      Alternative,
      MonadPlus,
      MonadFix,
      MonadThrow,
      MonadCatch,
      MonadFail,
      MonadMask
    )

instance MonadTrans BH where
  lift = BH . lift . lift

instance (MonadReader r m) => MonadReader r (BH m) where
  ask = lift ask
  local f (BH (ReaderT m)) = BH $
    ReaderT $ \r ->
      local f (m r)

instance (Functor m, Applicative m, MonadIO m, MonadCatch m, MonadThrow m) => MonadBH (BH m) where
  type Backend (BH m) = 'Dynamic
  dispatch request = BH $ do
    env <- ask @BHEnv
    liftIO $ do
      req <- mkHttpRequest env request
      let mgr = bhManager env
      rawResp <- HTTP.httpLbs req mgr
      BHResponse <$> bhResponseHook env rawResp
  tryEsError = try
  throwEsError = throwM

tryPerformBHRequest ::
  (MonadBH m, MonadThrow m, ParseBHResponse contextualized) =>
  BHRequest contextualized a ->
  m (ParsedEsResponse a)
tryPerformBHRequest req = dispatch req >>= either throwM return . bhRequestParser req

performBHRequest ::
  (MonadBH m, MonadThrow m, ParseBHResponse contextualized) =>
  BHRequest contextualized a ->
  m a
performBHRequest req = tryPerformBHRequest req >>= either throwEsError return

emptyBody :: L.ByteString
emptyBody = L.pack ""

-- | Build an @http-client@ 'HTTP.Request' from a 'BHEnv' and a
-- 'BHRequest'. This is the single source of truth for request
-- construction: the buffering 'dispatch' path and the streaming
-- 'withStreamingResponse' path both route through here, so any change
-- to header \/ body \/ query wiring applies uniformly.
--
-- The function preserves both query mechanisms a 'BHRequest' carries:
-- 'bhRequestEndpoint'\'s own 'getRawEndpointQueries' (rendered into the
-- URL string by 'getEndpoint') and the separate 'bhRequestQueryStrings'
-- (applied via 'HTTP.setQueryString' after parsing). The body defaults
-- to 'emptyBody' when 'Nothing' (matching the historical wire shape).
-- Status-code exceptions are suppressed
-- ('HTTP.setRequestIgnoreStatus') because Bloodhound does its own
-- status-based parsing in 'parseEsResponse'.
mkHttpRequest :: BHEnv -> BHRequest contextualized body -> IO HTTP.Request
mkHttpRequest env request = do
  initReq <- parseUrl' endpointText
  let setQueryStrings =
        case bhRequestQueryStrings request of
          [] -> id
          xs -> HTTP.setQueryString xs
  bhRequestHook env $
    HTTP.setRequestIgnoreStatus $
      setQueryStrings $
        initReq
          { HTTP.method = bhRequestMethod request,
            HTTP.requestHeaders =
              ("Content-Type", "application/json") : HTTP.requestHeaders initReq,
            HTTP.requestBody = HTTP.RequestBodyLBS $ fromMaybe emptyBody $ bhRequestBody request
          }
  where
    endpointText = getEndpoint (bhServer env) (bhRequestEndpoint request)

parseUrl' :: (MonadThrow m) => Text -> m HTTP.Request
parseUrl' t = HTTP.parseRequest (URI.escapeURIString URI.isAllowedInURI (T.unpack t))

-- | Execute a 'BHRequest' against the cluster and hand the caller the
-- streaming response body via a bracketed continuation, instead of
-- buffering the entire body into memory as 'dispatch'\/'performBHRequest'
-- do. This is the low-level primitive on which streaming endpoints
-- (e.g. the ML Commons Execute Stream Agent API, which emits
-- Server-Sent Events) are layered.
--
-- The continuation receives the @http-client@ 'HTTP.Response' whose
-- 'HTTP.responseBody' is a 'HTTP.BodyReader' (an @IO 'BS.ByteString'@
-- that yields the next chunk on each call, blocking until data is
-- available, and returning an empty 'BS.ByteString' at end-of-stream).
-- The bracket ensures the connection is released when the continuation
-- exits, even on exception. Use 'HTTP.brRead' to pull chunks manually,
-- or feed the 'HTTP.BodyReader' into a streaming parser (e.g. the SSE
-- parser in "Database.Bloodhound.Internal.Utils.SSE").
--
-- The continuation runs in 'IO' (not the underlying monad @m@): the
-- caller is responsible for any monad-lifting it needs inside the
-- continuation. Higher-level streaming sources built on this primitive
-- (e.g. a @conduit@ 'ConduitT') handle that lifting themselves.
--
-- /Behavioural differences from 'dispatch':/
--
-- * The 'bhResponseHook' is __not__ invoked. The hook's type
--   ('bhResponseHook' :: 'HTTP.Response' 'L.ByteString' -> @…@)
--   requires the body to be fully materialized, which defeats the
--   purpose of streaming. Debug logging that relied on the response
--   hook will not see streaming responses.
-- * 'HTTP.setRequestIgnoreStatus' is still applied (so non-2xx is not
--   raised as an 'HttpException' by @http-client@); the caller is
--   responsible for inspecting 'HTTP.responseStatus' inside the
--   continuation if it cares about the status.
-- * The 'bhRequestHook' __is__ applied (it operates on the request,
--   not the response), so custom authentication strategies continue to
--   work for streaming requests.
withStreamingResponse ::
  (MonadIO m) =>
  BHRequest contextualized body ->
  (HTTP.Response HTTP.BodyReader -> IO a) ->
  BH m a
withStreamingResponse request inner =
  BH $ do
    env <- ask @BHEnv
    liftIO $ do
      req <- mkHttpRequest env request
      HTTP.withResponse req (bhManager env) inner

runBH :: BHEnv -> BH m a -> m (Either EsError a)
runBH e f = runExceptT $ runReaderT (unBH f) e

-- | Statically-type backend.
--
-- It's also an useful wrapper for 'DerivingVia'
newtype StaticBH (backend :: BackendType) m a = StaticBH
  {runStaticBH :: m a}
  deriving newtype
    ( Functor,
      Applicative,
      Monad,
      MonadIO,
      MonadReader r,
      MonadState s,
      MonadWriter w,
      Alternative,
      MonadPlus,
      MonadFix,
      MonadThrow,
      MonadCatch,
      MonadFail,
      MonadMask
    )

instance MonadTrans (StaticBH backend) where
  lift = StaticBH

instance (MonadBH m) => MonadBH (StaticBH backend m) where
  type Backend (StaticBH backend m) = backend
  dispatch = StaticBH . dispatch
  tryEsError = StaticBH . tryEsError . runStaticBH
  throwEsError = StaticBH . throwEsError

-- | Run a piece of code as-if we are on a given backend
unsafePerformBH :: StaticBH backend m a -> m a
unsafePerformBH = runStaticBH

-- | Dependently-typed version of 'BackendType'
data SBackendType :: BackendType -> Type where
  SElasticSearch7 :: SBackendType 'ElasticSearch7
  SElasticSearch8 :: SBackendType 'ElasticSearch8
  SElasticSearch9 :: SBackendType 'ElasticSearch9
  SOpenSearch1 :: SBackendType 'OpenSearch1
  SOpenSearch2 :: SBackendType 'OpenSearch2
  SOpenSearch3 :: SBackendType 'OpenSearch3

-- | Run an action given an actual backend
withDynamicBH ::
  (MonadBH m) =>
  BackendType ->
  (forall backend. SBackendType backend -> StaticBH backend m a) ->
  m a
withDynamicBH backend f =
  case backend of
    ElasticSearch7 -> unsafePerformBH $ f SElasticSearch7
    ElasticSearch8 -> unsafePerformBH $ f SElasticSearch8
    ElasticSearch9 -> unsafePerformBH $ f SElasticSearch9
    OpenSearch1 -> unsafePerformBH $ f SOpenSearch1
    OpenSearch2 -> unsafePerformBH $ f SOpenSearch2
    OpenSearch3 -> unsafePerformBH $ f SOpenSearch3
    Dynamic -> throwEsError $ EsError Nothing "Cannot perform on a 'Dynamic' backend"