{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralisedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

-- |
-- Description : Instrument servant-client with eventuo11y
-- Copyright   : Copyright 2022 Shea Levy.
-- License     : Apache-2.0
-- Maintainer  : shea@shealevy.com
--
-- This module offers a variant of servant-client's 'S.ClientM' which instruments
-- all requests with 'Event's. It also has miscellaneous helpers for instrumenting
-- servant-client functionality in other ways.
module Observe.Event.Servant.Client
  ( -- * ClientM
    ClientM,
    ClientEnv (..),
    runClientM,

    -- ** Instrumentation
    RunRequest (..),
    runRequestJSON,
    RunRequestField (..),
    runRequestFieldJSON,

    -- * Miscellaneous instrumentation
    clientErrorJSON,
    responseJSON,
  )
where

import Control.Monad.Except
import Control.Monad.Reader
import Data.Aeson
import Data.Binary.Builder
import Data.ByteString.Lazy hiding (null)
import Data.ByteString.Lazy.Internal (ByteString (..))
import Data.CaseInsensitive
import Data.Coerce
import Data.Map.Strict (mapKeys)
import Data.Text.Encoding
import Network.HTTP.Media.MediaType
import Network.HTTP.Media.RenderHeader
import Network.HTTP.Types.Status
import Network.HTTP.Types.Version
import Observe.Event
import Observe.Event.Render.JSON
import Servant.Client hiding (ClientEnv, ClientM, runClientM)
import qualified Servant.Client as S
import Servant.Client.Core.Request
import Servant.Client.Core.RunClient hiding (RunRequest)

-- | A monad to use in place of 'S.ClientM' to get instrumentation on requests.
--
-- Note that 'ClientM' is only 'RunClient' when the selector of the ambient 'EventBackend' is
-- @s@, see the instance under 'ClientEnv'.
type ClientM em s = TransEventMonad (ReaderT (ClientEnv s)) (TransEventMonad (ExceptT ClientError) em)

-- | An instrumented 'S.ClientEnv'
data ClientEnv s = ClientEnv
  { forall (s :: * -> *). ClientEnv s -> ClientEnv
env :: !S.ClientEnv,
    forall (s :: * -> *). ClientEnv s -> InjectSelector RunRequest s
injectRunRequest :: !(InjectSelector RunRequest s)
  }

instance (MonadIO (em r s), MonadWithEvent em) => RunClient (ClientM em s r s) where
  runRequestAcceptStatus :: Maybe [Status] -> Request -> ClientM em s r s Response
runRequestAcceptStatus Maybe [Status]
stats Request
req = do
    ClientEnv s
e <- forall r (m :: * -> *). MonadReader r m => m r
ask
    forall (s :: * -> *). ClientEnv s -> InjectSelector RunRequest s
injectRunRequest ClientEnv s
e RunRequest RunRequestField
RunRequest \s g
runReq RunRequestField -> g
injField -> forall (em :: * -> (* -> *) -> * -> *) (s :: * -> *) r a f.
MonadWithEvent em =>
s f -> (EnvEvent em r s f -> em r s a) -> em r s a
withEvent s g
runReq \EnvEvent
  (TransEventMonad
     (ReaderT (ClientEnv s)) (TransEventMonad (ExceptT ClientError) em))
  r
  s
  g
ev -> do
      forall (m :: * -> *) r f. Event m r f -> f -> m ()
addField EnvEvent
  (TransEventMonad
     (ReaderT (ClientEnv s)) (TransEventMonad (ExceptT ClientError) em))
  r
  s
  g
ev forall b c a. (b -> c) -> (a -> b) -> a -> c
. RunRequestField -> g
injField forall a b. (a -> b) -> a -> b
$ Request -> RunRequestField
ReqField Request
req
      Response
res <- coerce :: forall a b. Coercible a b => a -> b
coerce forall a b. (a -> b) -> a -> b
$ forall a b. a -> b -> a
const @_ @(ClientEnv s) (forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO @(em r s) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. ClientM a -> ClientEnv -> IO (Either ClientError a)
S.runClientM (forall (s :: * -> *). ClientEnv s -> ClientEnv
env ClientEnv s
e) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
RunClient m =>
Maybe [Status] -> Request -> m Response
runRequestAcceptStatus Maybe [Status]
stats Request
req)
      forall (m :: * -> *) r f. Event m r f -> f -> m ()
addField EnvEvent
  (TransEventMonad
     (ReaderT (ClientEnv s)) (TransEventMonad (ExceptT ClientError) em))
  r
  s
  g
ev forall b c a. (b -> c) -> (a -> b) -> a -> c
. RunRequestField -> g
injField forall a b. (a -> b) -> a -> b
$ Response -> RunRequestField
ResField Response
res
      forall (f :: * -> *) a. Applicative f => a -> f a
pure Response
res
  throwClientError :: forall a. ClientError -> ClientM em s r s a
throwClientError = forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError

-- | Instrumented version of 'S.runClientM'
runClientM :: ClientM em s r s' a -> ClientEnv s -> em r s' (Either ClientError a)
runClientM :: forall (em :: * -> (* -> *) -> * -> *) (s :: * -> *) r
       (s' :: * -> *) a.
ClientM em s r s' a
-> ClientEnv s -> em r s' (Either ClientError a)
runClientM = coerce :: forall a b. Coercible a b => a -> b
coerce

-- | Selector for events in 'ClientM'
data RunRequest f where
  RunRequest :: RunRequest RunRequestField

-- | Render a 'RunRequest' and the fields of its selected events as JSON
runRequestJSON :: RenderSelectorJSON RunRequest
runRequestJSON :: RenderSelectorJSON RunRequest
runRequestJSON RunRequest f
RunRequest = (Key
"run-request", RenderFieldJSON RunRequestField
runRequestFieldJSON)

-- | A field for v'RunRequest' events.
data RunRequestField
  = ReqField Request
  | ResField Response

-- | Render a 'RunRequestField' as JSON.
runRequestFieldJSON :: RenderFieldJSON RunRequestField
runRequestFieldJSON :: RenderFieldJSON RunRequestField
runRequestFieldJSON (ReqField Request {Maybe (RequestBody, MediaType)
ByteString
Builder
Seq QueryItem
Seq Header
Seq MediaType
HttpVersion
requestPath :: forall body path. RequestF body path -> path
requestQueryString :: forall body path. RequestF body path -> Seq QueryItem
requestBody :: forall body path. RequestF body path -> Maybe (body, MediaType)
requestAccept :: forall body path. RequestF body path -> Seq MediaType
requestHeaders :: forall body path. RequestF body path -> Seq Header
requestHttpVersion :: forall body path. RequestF body path -> HttpVersion
requestMethod :: forall body path. RequestF body path -> ByteString
requestMethod :: ByteString
requestHttpVersion :: HttpVersion
requestHeaders :: Seq Header
requestAccept :: Seq MediaType
requestBody :: Maybe (RequestBody, MediaType)
requestQueryString :: Seq QueryItem
requestPath :: Builder
..}) =
  ( Key
"request",
    Object -> Value
Object
      ( Key
"path" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= ByteString -> Text
decodeUtf8 (ByteString -> ByteString
toStrict forall a b. (a -> b) -> a -> b
$ Builder -> ByteString
toLazyByteString Builder
requestPath)
          forall a. Semigroup a => a -> a -> a
<> ( if forall (t :: * -> *) a. Foldable t => t a -> Bool
null Seq QueryItem
requestQueryString
                 then forall a. Monoid a => a
mempty
                 else
                   Key
"query"
                     forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= ( ( \(ByteString
k, Maybe ByteString
mv) ->
                              Object -> Value
Object
                                ( Key
"key" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= ByteString -> Text
decodeUtf8 ByteString
k
                                    forall a. Semigroup a => a -> a -> a
<> forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Monoid a => a
mempty ((Key
"value" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.=) forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Text
decodeUtf8) Maybe ByteString
mv
                                )
                          )
                            forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Seq QueryItem
requestQueryString
                        )
             )
          forall a. Semigroup a => a -> a -> a
<> case Maybe (RequestBody, MediaType)
requestBody of
            Maybe (RequestBody, MediaType)
Nothing -> forall a. Monoid a => a
mempty
            Just (RequestBody
body, MediaType
ty) ->
              ( Key
"content-type" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= ByteString -> Text
decodeUtf8 (forall h. RenderHeader h => h -> ByteString
renderHeader MediaType
ty)
                  forall a. Semigroup a => a -> a -> a
<> case RequestBody
body of
                    RequestBodyBS ByteString
bs -> Key
"body" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= ByteString -> Text
decodeUtf8 ByteString
bs
                    RequestBodyLBS ByteString
Empty -> Key
"body" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= Bool
False
                    RequestBodyLBS (Chunk ByteString
bs ByteString
Empty) -> Key
"body" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= ByteString -> Text
decodeUtf8 ByteString
bs
                    RequestBody
_ -> forall a. Monoid a => a
mempty
              )
          forall a. Semigroup a => a -> a -> a
<> ( if forall (t :: * -> *) a. Foldable t => t a -> Bool
null Seq MediaType
requestAccept
                 then forall a. Monoid a => a
mempty
                 else Key
"accept" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ByteString -> Text
decodeUtf8 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall h. RenderHeader h => h -> ByteString
renderHeader) Seq MediaType
requestAccept
             )
          forall a. Semigroup a => a -> a -> a
<> ( if forall (t :: * -> *) a. Foldable t => t a -> Bool
null Seq Header
requestHeaders
                 then forall a. Monoid a => a
mempty
                 else Key
"headers" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(CI ByteString
nm, ByteString
val) -> Object -> Value
Object (Key
"name" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= ByteString -> Text
decodeUtf8 (forall s. CI s -> s
original CI ByteString
nm) forall a. Semigroup a => a -> a -> a
<> (if CI ByteString
nm forall a. Eq a => a -> a -> Bool
== CI ByteString
"Authorization" then forall a. Monoid a => a
mempty else Key
"val" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= ByteString -> Text
decodeUtf8 ByteString
val))) Seq Header
requestHeaders
             )
          forall a. Semigroup a => a -> a -> a
<> Key
"http-version" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= Object -> Value
Object (Key
"major" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= HttpVersion -> Int
httpMajor HttpVersion
requestHttpVersion forall a. Semigroup a => a -> a -> a
<> Key
"minor" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= HttpVersion -> Int
httpMinor HttpVersion
requestHttpVersion)
          forall a. Semigroup a => a -> a -> a
<> Key
"method" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= ByteString -> Text
decodeUtf8 ByteString
requestMethod
      )
  )
runRequestFieldJSON (ResField Response
res) =
  ( Key
"response",
    Response -> Bool -> Value
responseJSON Response
res Bool
False
  )

-- | Render a 'ClientError', considered as an 'Event' field, as JSON
clientErrorJSON :: RenderFieldJSON ClientError
clientErrorJSON :: RenderFieldJSON ClientError
clientErrorJSON (FailureResponse RequestF () (BaseUrl, ByteString)
_ Response
res) = (Key
"failure-response", Response -> Bool -> Value
responseJSON Response
res Bool
True)
clientErrorJSON (DecodeFailure Text
err Response
res) = (Key
"decode-failure", Object -> Value
Object (Key
"response" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= Response -> Bool -> Value
responseJSON Response
res Bool
True forall a. Semigroup a => a -> a -> a
<> Key
"err" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= Text -> Value
String Text
err))
clientErrorJSON (UnsupportedContentType MediaType
ty Response
res) =
  ( Key
"unsupported-content-type",
    Object -> Value
Object
      ( Key
"response" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= Response -> Bool -> Value
responseJSON Response
res Bool
True
          forall a. Semigroup a => a -> a -> a
<> Key
"main-type" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= ByteString -> Text
decodeUtf8 (forall s. CI s -> s
original forall a b. (a -> b) -> a -> b
$ MediaType -> CI ByteString
mainType MediaType
ty)
          forall a. Semigroup a => a -> a -> a
<> Key
"sub-type" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= ByteString -> Text
decodeUtf8 (forall s. CI s -> s
original forall a b. (a -> b) -> a -> b
$ MediaType -> CI ByteString
subType MediaType
ty)
          forall a. Semigroup a => a -> a -> a
<> Key
"parameters" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ByteString -> Text
decodeUtf8 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s. CI s -> s
original) (forall k2 k1 a. Ord k2 => (k1 -> k2) -> Map k1 a -> Map k2 a
mapKeys (ByteString -> Text
decodeUtf8 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s. CI s -> s
original) forall a b. (a -> b) -> a -> b
$ MediaType -> Parameters
parameters MediaType
ty)
      )
  )
clientErrorJSON (InvalidContentTypeHeader Response
res) = (Key
"invalid-content-type-header", Response -> Bool -> Value
responseJSON Response
res Bool
True)
clientErrorJSON (ConnectionError SomeException
e) = (Key
"connection-error", forall a. ToJSON a => a -> Value
toJSON forall a b. (a -> b) -> a -> b
$ forall a. Show a => a -> String
show SomeException
e)

-- | Render a 'Servant.Client.Core.Response' as JSON, optionally forcing rendering the body even if it's large.
responseJSON :: Response -> Bool -> Value
responseJSON :: Response -> Bool -> Value
responseJSON Response {ByteString
Seq Header
HttpVersion
Status
responseHttpVersion :: forall a. ResponseF a -> HttpVersion
responseStatusCode :: forall a. ResponseF a -> Status
responseBody :: forall a. ResponseF a -> a
responseHeaders :: forall a. ResponseF a -> Seq Header
responseBody :: ByteString
responseHttpVersion :: HttpVersion
responseHeaders :: Seq Header
responseStatusCode :: Status
..} Bool
forceBody =
  Object -> Value
Object
    ( Key
"status" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= Status -> Int
statusCode Status
responseStatusCode
        forall a. Semigroup a => a -> a -> a
<> ( if forall (t :: * -> *) a. Foldable t => t a -> Bool
null Seq Header
responseHeaders
               then forall a. Monoid a => a
mempty
               else Key
"headers" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(CI ByteString
nm, ByteString
val) -> Object -> Value
Object (Key
"name" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= ByteString -> Text
decodeUtf8 (forall s. CI s -> s
original CI ByteString
nm) forall a. Semigroup a => a -> a -> a
<> (if CI ByteString
nm forall a. Eq a => a -> a -> Bool
== CI ByteString
"Cookie" then forall a. Monoid a => a
mempty else Key
"val" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= ByteString -> Text
decodeUtf8 ByteString
val))) Seq Header
responseHeaders
           )
        forall a. Semigroup a => a -> a -> a
<> Key
"http-version" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= Object -> Value
Object (Key
"major" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= HttpVersion -> Int
httpMajor HttpVersion
responseHttpVersion forall a. Semigroup a => a -> a -> a
<> Key
"minor" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= HttpVersion -> Int
httpMinor HttpVersion
responseHttpVersion)
        forall a. Semigroup a => a -> a -> a
<> ( if Bool
forceBody
               then Key
"body" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= (ByteString -> Text
decodeUtf8 forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
toStrict ByteString
responseBody)
               else case ByteString
responseBody of
                 ByteString
Empty -> Key
"body" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= Bool
False
                 Chunk ByteString
bs ByteString
Empty -> Key
"body" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= ByteString -> Text
decodeUtf8 ByteString
bs
                 ByteString
_ -> forall a. Monoid a => a
mempty
           )
    )