{-| This module represents an extension to 'hasql-pool', it allows for a mix of dynamic credentials
    and static settings for connections in a pool.
    Due to the fact that the module tries to extend 'hasql-pool' rather than rewrite it, the exposed API relies on
    combining existing 'hasql' types with new types, whereas a rewrite would just extend the original types.
    It is done so to simplify maintenance of the extended functionality and make it more compatible with any
    future development of 'hasql-pool'.
-}
module Hasql.Pool
(   Pool
,   Settings(..)
,   UsageError(..)
,   ConnectionGetter
,   acquire
,   acquireWith
,   release
,   use
,   useWithObserver
,   getPoolUsageStat
)
where

import qualified Data.Pool as ResourcePool
import           System.Clock (Clock(Monotonic), diffTimeSpec, getTime, toNanoSecs)

import           Hasql.Pool.Prelude
import qualified Hasql.Connection
import qualified Hasql.Session
import qualified Hasql.Pool.ResourcePool as ResourcePool
import           Hasql.Pool.Observer (Observed(..), ObserverAction)


-- |
-- A pool of open DB connections.
newtype Pool =
    Pool (ResourcePool.Pool (Either Hasql.Connection.ConnectionError Hasql.Connection.Connection))
    deriving (PoolStripes -> Pool -> ShowS
[Pool] -> ShowS
Pool -> String
forall a.
(PoolStripes -> a -> ShowS)
-> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Pool] -> ShowS
$cshowList :: [Pool] -> ShowS
show :: Pool -> String
$cshow :: Pool -> String
showsPrec :: PoolStripes -> Pool -> ShowS
$cshowsPrec :: PoolStripes -> Pool -> ShowS
Show)


type PoolSize         = Int
type PoolStripes      = Int
type ResidenceTimeout = NominalDiffTime

-- |
-- Connection getter action that allows for obtaining Postgres connection settings
-- via external resources such as AWS tokens etc.
type ConnectionGetter = IO (Either Hasql.Connection.ConnectionError Hasql.Connection.Connection)

-- |
-- Settings of the connection pool. Consist of:
--
-- * Pool-size.
--
-- * Timeout.
-- An amount of time for which an unused resource is kept open.
-- The smallest acceptable value is 0.5 seconds.
--
-- * Connection settings.
--
type Settings =
  (PoolSize, ResidenceTimeout, Hasql.Connection.Settings)

-- |
-- Given the pool-size, timeout and connection settings
-- create a connection-pool.
acquire :: Settings -> IO Pool
acquire :: Settings -> IO Pool
acquire settings :: Settings
settings@(PoolStripes
_size, ResidenceTimeout
_timeout, Settings
connectionSettings) =
    PoolStripes -> ConnectionGetter -> Settings -> IO Pool
acquireWith forall {a}. Num a => a
stripes (Settings -> ConnectionGetter
Hasql.Connection.acquire Settings
connectionSettings) Settings
settings
    where
        stripes :: a
stripes = a
1


-- |
-- Similar to 'acquire', allows for finer configuration.
acquireWith :: PoolStripes
            -> ConnectionGetter
            -> Settings
            -> IO Pool
acquireWith :: PoolStripes -> ConnectionGetter -> Settings -> IO Pool
acquireWith PoolStripes
stripes ConnectionGetter
connGetter (PoolStripes
size, ResidenceTimeout
timeout, Settings
connectionSettings) =
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Pool (Either ConnectionError Connection) -> Pool
Pool forall a b. (a -> b) -> a -> b
$
        forall a.
IO a
-> (a -> IO ())
-> PoolStripes
-> ResidenceTimeout
-> PoolStripes
-> IO (Pool a)
ResourcePool.createPool ConnectionGetter
connGetter forall {a}. Either a Connection -> IO ()
release PoolStripes
stripes ResidenceTimeout
timeout PoolStripes
size
    where
        release :: Either a Connection -> IO ()
release = forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (forall a b. a -> b -> a
const (forall (f :: * -> *) a. Applicative f => a -> f a
pure ())) Connection -> IO ()
Hasql.Connection.release


-- |
-- Release the connection-pool by closing and removing all connections.
release :: Pool -> IO ()
release :: Pool -> IO ()
release (Pool Pool (Either ConnectionError Connection)
pool) =
    forall a. Pool a -> IO ()
ResourcePool.destroyAllResources Pool (Either ConnectionError Connection)
pool


-- |
-- A union over the connection establishment error and the session error.
data UsageError
    =   ConnectionError Hasql.Connection.ConnectionError
    |   SessionError    Hasql.Session.QueryError
    deriving (PoolStripes -> UsageError -> ShowS
[UsageError] -> ShowS
UsageError -> String
forall a.
(PoolStripes -> a -> ShowS)
-> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [UsageError] -> ShowS
$cshowList :: [UsageError] -> ShowS
show :: UsageError -> String
$cshow :: UsageError -> String
showsPrec :: PoolStripes -> UsageError -> ShowS
$cshowsPrec :: PoolStripes -> UsageError -> ShowS
Show, UsageError -> UsageError -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: UsageError -> UsageError -> Bool
$c/= :: UsageError -> UsageError -> Bool
== :: UsageError -> UsageError -> Bool
$c== :: UsageError -> UsageError -> Bool
Eq)

-- |
-- Use a connection from the pool to run a session and
-- return the connection to the pool, when finished.
use :: Pool -> Hasql.Session.Session a -> IO (Either UsageError a)
use :: forall a. Pool -> Session a -> IO (Either UsageError a)
use = forall a.
Maybe ObserverAction
-> Pool -> Session a -> IO (Either UsageError a)
useWithObserver forall a. Maybe a
Nothing

-- |
-- Same as 'use' but allows for a custom observer action. You can use it for gathering latency metrics.
useWithObserver :: Maybe ObserverAction
                -> Pool
                -> Hasql.Session.Session a
                -> IO (Either UsageError a)
useWithObserver :: forall a.
Maybe ObserverAction
-> Pool -> Session a -> IO (Either UsageError a)
useWithObserver Maybe ObserverAction
observer (Pool Pool (Either ConnectionError Connection)
pool) Session a
session =
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (forall a b. a -> Either a b
Left forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ConnectionError -> UsageError
ConnectionError) (forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (forall a b. a -> Either a b
Left forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. QueryError -> UsageError
SessionError) forall a b. b -> Either a b
Right)) forall a b. (a -> b) -> a -> b
$
    forall resource failure success.
Pool resource
-> (resource -> IO (Either failure success))
-> IO (Either failure success)
ResourcePool.withResourceOnEither Pool (Either ConnectionError Connection)
pool forall a b. (a -> b) -> a -> b
$
    forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Connection -> IO (Either QueryError a)
runQuery
    where
        runQuery :: Connection -> IO (Either QueryError a)
runQuery Connection
dbConn = forall b a. b -> (a -> b) -> Maybe a -> b
maybe IO (Either QueryError a)
action (forall {b} {a}. IO b -> (Observed -> IO a) -> IO b
runWithObserver IO (Either QueryError a)
action) Maybe ObserverAction
observer
            where
                action :: IO (Either QueryError a)
action = forall a. Session a -> Connection -> IO (Either QueryError a)
Hasql.Session.run Session a
session Connection
dbConn

        runWithObserver :: IO b -> (Observed -> IO a) -> IO b
runWithObserver IO b
action Observed -> IO a
doObserve = do
            let measure :: IO TimeSpec
measure = Clock -> IO TimeSpec
getTime Clock
Monotonic
            TimeSpec
start  <- IO TimeSpec
measure
            b
result <- IO b
action
            TimeSpec
end    <- IO TimeSpec
measure
            let nsRatio :: a
nsRatio  = a
1000000000
                observed :: Observed
observed = Observed {   latency :: Ratio Integer
latency = forall a. Real a => a -> Ratio Integer
toRational (TimeSpec -> Integer
toNanoSecs (TimeSpec
end TimeSpec -> TimeSpec -> TimeSpec
`diffTimeSpec` TimeSpec
start) forall a. Integral a => a -> a -> Ratio a
% forall {a}. Num a => a
nsRatio)
                                    }
            Observed -> IO a
doObserve Observed
observed forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (f :: * -> *) a. Applicative f => a -> f a
pure b
result


getPoolStats :: Pool -> IO ResourcePool.Stats
getPoolStats :: Pool -> IO Stats
getPoolStats (Pool Pool (Either ConnectionError Connection)
p) = forall a. Pool a -> Bool -> IO Stats
ResourcePool.stats Pool (Either ConnectionError Connection)
p Bool
performStatsReset
    where
        performStatsReset :: Bool
performStatsReset = Bool
False


getPoolUsageStat :: Pool -> IO PoolSize
getPoolUsageStat :: Pool -> IO PoolStripes
getPoolUsageStat Pool
pool = Pool -> IO Stats
getPoolStats Pool
pool forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall {f :: * -> *}. Applicative f => Stats -> f PoolStripes
gather where
    gather :: Stats -> f PoolStripes
gather = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. PoolStats -> PoolStripes
ResourcePool.currentUsage forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Stats -> PoolStats
ResourcePool.poolStats