-- | Integration of "Freckle.App" tooling with "Yesod"
module Freckle.App.Yesod
  ( respondQueryCanceled
  , respondQueryCanceledHeaders
  , logExceptionsMiddleware
  ) where

import Freckle.App.Prelude

import Blammo.Logging
import Database.PostgreSQL.Simple (SqlError (..))
import Freckle.App.Exception
  ( AnnotatedException (..)
  , annotatedExceptionMessageFrom
  , fromException
  , withException
  )
import Freckle.App.Stats (HasStatsClient)
import qualified Freckle.App.Stats as Stats
import Network.HTTP.Types (ResponseHeaders, status503)
import qualified Network.Wai as W
import Yesod.Core.Handler (HandlerFor, sendWaiResponse)
import Yesod.Core.Types (HandlerContents)

-- | Catch 'SqlError' when queries are canceled due to timeout and respond 503
--
-- Also logs and increments a metric.
respondQueryCanceled
  :: HasStatsClient site => HandlerFor site res -> HandlerFor site res
respondQueryCanceled :: forall site res.
HasStatsClient site =>
HandlerFor site res -> HandlerFor site res
respondQueryCanceled = ResponseHeaders -> HandlerFor site res -> HandlerFor site res
forall site res.
HasStatsClient site =>
ResponseHeaders -> HandlerFor site res -> HandlerFor site res
respondQueryCanceledHeaders []

-- | 'respondQueryCanceledHeaders' but adding headers to the 503 response
respondQueryCanceledHeaders
  :: HasStatsClient site
  => ResponseHeaders
  -> HandlerFor site res
  -> HandlerFor site res
respondQueryCanceledHeaders :: forall site res.
HasStatsClient site =>
ResponseHeaders -> HandlerFor site res -> HandlerFor site res
respondQueryCanceledHeaders ResponseHeaders
headers HandlerFor site res
handler =
  (AnnotatedException SqlError
 -> Maybe (AnnotatedException SqlError))
-> HandlerFor site res
-> (AnnotatedException SqlError -> HandlerFor site res)
-> HandlerFor site res
forall e b (m :: * -> *) a.
(Exception e, MonadUnliftIO m, HasCallStack) =>
(e -> Maybe b) -> m a -> (b -> m a) -> m a
catchJust AnnotatedException SqlError -> Maybe (AnnotatedException SqlError)
queryCanceled HandlerFor site res
handler ((AnnotatedException SqlError -> HandlerFor site res)
 -> HandlerFor site res)
-> (AnnotatedException SqlError -> HandlerFor site res)
-> HandlerFor site res
forall a b. (a -> b) -> a -> b
$ \AnnotatedException SqlError
ex -> do
    LogSource -> Message -> HandlerFor site ()
forall (m :: * -> *).
(HasCallStack, MonadLogger m) =>
LogSource -> Message -> m ()
logErrorNS LogSource
"yesod" (Message -> HandlerFor site ()) -> Message -> HandlerFor site ()
forall a b. (a -> b) -> a -> b
$ (SqlError -> Message) -> AnnotatedException SqlError -> Message
forall ex.
Exception ex =>
(ex -> Message) -> AnnotatedException ex -> Message
annotatedExceptionMessageFrom (Message -> SqlError -> Message
forall a b. a -> b -> a
const Message
"Query canceled") AnnotatedException SqlError
ex
    LogSource -> HandlerFor site ()
forall (m :: * -> *) env.
(MonadUnliftIO m, MonadReader env m, HasStatsClient env) =>
LogSource -> m ()
Stats.increment LogSource
"query_canceled"
    Response -> HandlerFor site res
forall (m :: * -> *) b. MonadHandler m => Response -> m b
sendWaiResponse (Response -> HandlerFor site res)
-> Response -> HandlerFor site res
forall a b. (a -> b) -> a -> b
$ Status -> ResponseHeaders -> ByteString -> Response
W.responseLBS Status
status503 ResponseHeaders
headers ByteString
"Query canceled"

logExceptionsMiddleware :: (MonadUnliftIO m, MonadLogger m) => m a -> m a
logExceptionsMiddleware :: forall (m :: * -> *) a.
(MonadUnliftIO m, MonadLogger m) =>
m a -> m a
logExceptionsMiddleware m a
f =
  m a
f m a -> (AnnotatedException SomeException -> m ()) -> m a
forall e a (m :: * -> *) b.
(Exception e, MonadUnliftIO m, HasCallStack) =>
m a -> (e -> m b) -> m a
`withException` \AnnotatedException SomeException
ex ->
    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (AnnotatedException SomeException -> Bool
isHandlerContents AnnotatedException SomeException
ex) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
      LogSource -> Message -> m ()
forall (m :: * -> *).
(HasCallStack, MonadLogger m) =>
LogSource -> Message -> m ()
logErrorNS LogSource
"yesod" (Message -> m ()) -> Message -> m ()
forall a b. (a -> b) -> a -> b
$
        (SomeException -> Message)
-> AnnotatedException SomeException -> Message
forall ex.
Exception ex =>
(ex -> Message) -> AnnotatedException ex -> Message
annotatedExceptionMessageFrom (Message -> SomeException -> Message
forall a b. a -> b -> a
const Message
"Handler exception") AnnotatedException SomeException
ex

isHandlerContents :: AnnotatedException SomeException -> Bool
isHandlerContents :: AnnotatedException SomeException -> Bool
isHandlerContents = Maybe HandlerContents -> Bool
forall a. Maybe a -> Bool
isJust (Maybe HandlerContents -> Bool)
-> (AnnotatedException SomeException -> Maybe HandlerContents)
-> AnnotatedException SomeException
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall e. Exception e => SomeException -> Maybe e
fromException @HandlerContents (SomeException -> Maybe HandlerContents)
-> (AnnotatedException SomeException -> SomeException)
-> AnnotatedException SomeException
-> Maybe HandlerContents
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AnnotatedException SomeException -> SomeException
forall exception. AnnotatedException exception -> exception
exception

queryCanceled
  :: AnnotatedException SqlError -> Maybe (AnnotatedException SqlError)
queryCanceled :: AnnotatedException SqlError -> Maybe (AnnotatedException SqlError)
queryCanceled AnnotatedException SqlError
ex = AnnotatedException SqlError
ex AnnotatedException SqlError
-> Maybe () -> Maybe (AnnotatedException SqlError)
forall a b. a -> Maybe b -> Maybe a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (SqlError -> ByteString
sqlState (AnnotatedException SqlError -> SqlError
forall exception. AnnotatedException exception -> exception
exception AnnotatedException SqlError
ex) ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
== ByteString
"57014")