-------------------------------------------------------------------------------
-- network-wait
-- Copyright 2022 Michael B. Gale (github@michael-gale.co.uk)
-------------------------------------------------------------------------------

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

-- | This module exports variants of the functions from "Network.Wait"
-- specialised for PostgreSQL servers. In addition to checking whether a
-- connection can be established, the functions in this module also check
-- whether the PostgreSQL server is ready to accept commands.
module Network.Wait.PostgreSQL (
    PostgreSqlConnectInfo(..),
    waitPostgreSql,
    waitPostgreSqlVerbose,
    waitPostgreSqlVerboseFormat,
    waitPostgreSqlWith
) where

-------------------------------------------------------------------------------

import Data.ByteString ( ByteString )

import Control.Monad
import Control.Monad.Catch
import Control.Monad.IO.Class
import Control.Retry

import Database.PostgreSQL.Simple
import Database.PostgreSQL.Simple.Internal

import Network.Wait

-------------------------------------------------------------------------------

-- | Used to abstract over different ways to describe a database connection.
class PostgreSqlConnectInfo a where
    -- | `connectDb` @info@ attempts to establish a database connection using
    -- a configuration given by @info@.
    connectDb :: a -> IO Connection

instance PostgreSqlConnectInfo ConnectInfo where
    connectDb :: ConnectInfo -> IO Connection
connectDb = ConnectInfo -> IO Connection
connect

instance PostgreSqlConnectInfo ByteString where
    connectDb :: ByteString -> IO Connection
connectDb = ByteString -> IO Connection
connectPostgreSQL

-- | `waitPostgreSql` @retryPolicy connectInfo@ is a variant of
-- `waitPostgresWith` which does not install any additional handlers.
waitPostgreSql
    :: (MonadIO m, MonadMask m, PostgreSqlConnectInfo info)
    => RetryPolicyM m -> info -> m Connection
waitPostgreSql :: forall (m :: * -> *) info.
(MonadIO m, MonadMask m, PostgreSqlConnectInfo info) =>
RetryPolicyM m -> info -> m Connection
waitPostgreSql = [RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> info -> m Connection
forall (m :: * -> *) info.
(MonadIO m, MonadMask m, PostgreSqlConnectInfo info) =>
[RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> info -> m Connection
waitPostgreSqlWith []

-- | `waitPostgreSqlVerbose` @outputHandler retryPolicy connectInfo@ is a variant
-- of `waitPostgreSqlVerboseFormat` which catches all exceptions derived from
-- `SomeException` and formats retry attempt information using `defaultLogMsg`
-- before passing the resulting `String` to @out@.
waitPostgreSqlVerbose
    :: (MonadIO m, MonadMask m, PostgreSqlConnectInfo info)
    => (String -> m ()) -> RetryPolicyM m -> info -> m Connection
waitPostgreSqlVerbose :: forall (m :: * -> *) info.
(MonadIO m, MonadMask m, PostgreSqlConnectInfo info) =>
(String -> m ()) -> RetryPolicyM m -> info -> m Connection
waitPostgreSqlVerbose String -> m ()
out =
    forall e (m :: * -> *) info.
(MonadIO m, MonadMask m, PostgreSqlConnectInfo info,
 Exception e) =>
(Bool -> e -> RetryStatus -> m ())
-> RetryPolicyM m -> info -> m Connection
waitPostgreSqlVerboseFormat @SomeException ((Bool -> SomeException -> RetryStatus -> m ())
 -> RetryPolicyM m -> info -> m Connection)
-> (Bool -> SomeException -> RetryStatus -> m ())
-> RetryPolicyM m
-> info
-> m Connection
forall a b. (a -> b) -> a -> b
$
    \Bool
b SomeException
ex RetryStatus
st -> String -> m ()
out (String -> m ()) -> String -> m ()
forall a b. (a -> b) -> a -> b
$ Bool -> SomeException -> RetryStatus -> String
forall e. Exception e => Bool -> e -> RetryStatus -> String
defaultLogMsg Bool
b SomeException
ex RetryStatus
st

-- | `waitPostgreSqlVerboseFormat` @outputHandler retryPolicy connectInfo@ is a
-- variant of `waitPostgreSqlWith` which installs an extra handler based on
-- `logRetries` which passes status information for each retry attempt
-- to @outputHandler@.
waitPostgreSqlVerboseFormat
    :: forall e m info. (MonadIO m, MonadMask m, PostgreSqlConnectInfo info, Exception e)
    => (Bool -> e -> RetryStatus -> m ())
    -> RetryPolicyM m
    -> info
    -> m Connection
waitPostgreSqlVerboseFormat :: forall e (m :: * -> *) info.
(MonadIO m, MonadMask m, PostgreSqlConnectInfo info,
 Exception e) =>
(Bool -> e -> RetryStatus -> m ())
-> RetryPolicyM m -> info -> m Connection
waitPostgreSqlVerboseFormat Bool -> e -> RetryStatus -> m ()
out = [RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> info -> m Connection
forall (m :: * -> *) info.
(MonadIO m, MonadMask m, PostgreSqlConnectInfo info) =>
[RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> info -> m Connection
waitPostgreSqlWith [RetryStatus -> Handler m Bool
h]
    where h :: RetryStatus -> Handler m Bool
h = (e -> m Bool)
-> (Bool -> e -> RetryStatus -> m ())
-> RetryStatus
-> Handler m Bool
forall (m :: * -> *) e.
(Monad m, Exception e) =>
(e -> m Bool)
-> (Bool -> e -> RetryStatus -> m ())
-> RetryStatus
-> Handler m Bool
logRetries (m Bool -> e -> m Bool
forall a b. a -> b -> a
const (m Bool -> e -> m Bool) -> m Bool -> e -> m Bool
forall a b. (a -> b) -> a -> b
$ Bool -> m Bool
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True) Bool -> e -> RetryStatus -> m ()
out

-- | `waitPostgreSqlWith` @extraHandlers retryPolicy connectInfo@ will attempt
-- to connect to the PostgreSQL server using @connectInfo@ and check that the
-- server is ready to accept commands. If this check fails, @retryPolicy@ is
-- used to determine whether (and how often) this function should attempt to
-- retry establishing the connection. By default, this function will retry
-- after all exceptions (except for those given by `skipAsyncExceptions`).
-- This behaviour may be customised with @extraHandlers@ which are installed
-- after `skipAsyncExceptions`, but before the default exception handler. The
--  @extraHandlers@ may also be used to report retry attempts to e.g. the
-- standard output or a logger.
waitPostgreSqlWith
    :: (MonadIO m, MonadMask m, PostgreSqlConnectInfo info)
    => [RetryStatus -> Handler m Bool] -> RetryPolicyM m -> info
    -> m Connection
waitPostgreSqlWith :: forall (m :: * -> *) info.
(MonadIO m, MonadMask m, PostgreSqlConnectInfo info) =>
[RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> info -> m Connection
waitPostgreSqlWith [RetryStatus -> Handler m Bool]
hs RetryPolicyM m
policy info
info =
    [RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> m Connection -> m Connection
forall (m :: * -> *) a.
(MonadIO m, MonadMask m) =>
[RetryStatus -> Handler m Bool] -> RetryPolicyM m -> m a -> m a
recoveringWith [RetryStatus -> Handler m Bool]
hs RetryPolicyM m
policy (m Connection -> m Connection) -> m Connection -> m Connection
forall a b. (a -> b) -> a -> b
$
    IO Connection -> m Connection
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Connection -> m Connection) -> IO Connection -> m Connection
forall a b. (a -> b) -> a -> b
$
    IO Connection
-> (Connection -> IO ())
-> (Connection -> IO Connection)
-> IO Connection
forall (m :: * -> *) a c b.
(HasCallStack, MonadMask m) =>
m a -> (a -> m c) -> (a -> m b) -> m b
bracket (info -> IO Connection
forall a. PostgreSqlConnectInfo a => a -> IO Connection
connectDb info
info) Connection -> IO ()
close ((Connection -> IO Connection) -> IO Connection)
-> (Connection -> IO Connection) -> IO Connection
forall a b. (a -> b) -> a -> b
$ \Connection
con -> do
        rs <- forall r. FromRow r => Connection -> Query -> IO [r]
query_ @[Int] Connection
con Query
"SELECT 1;"
        unless (rs == [[1]]) $ throwM $
            fatalError "Unexpected result for SELECT 1."
        pure con

-------------------------------------------------------------------------------