module Hasql.Pool
(   Pool
,   Settings(..)
,   UsageError(..)
,   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 connections to DB.
newtype Pool =
    Pool (ResourcePool.Pool (Either Hasql.Connection.ConnectionError Hasql.Connection.Connection))
    deriving (Int -> Pool -> ShowS
[Pool] -> ShowS
Pool -> String
(Int -> Pool -> ShowS)
-> (Pool -> String) -> ([Pool] -> ShowS) -> Show Pool
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Pool] -> ShowS
$cshowList :: [Pool] -> ShowS
show :: Pool -> String
$cshow :: Pool -> String
showsPrec :: Int -> Pool -> ShowS
$cshowsPrec :: Int -> Pool -> ShowS
Show)


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

-- |
-- Connection getter action that allows for obtaining Postgres connection settings via external rsources 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@(Int
_size, ResidenceTimeout
_timeout, Settings
connectionSettings) =
    Int -> ConnectionGetter -> Settings -> IO Pool
acquireWith Int
forall p. Num p => p
stripes (Settings -> ConnectionGetter
Hasql.Connection.acquire Settings
connectionSettings) Settings
settings
    where
        stripes :: p
stripes = p
1


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


-- |
-- Release the connection-pool.
release :: Pool -> IO ()
release :: Pool -> IO ()
release (Pool Pool (Either ConnectionError Connection)
pool) =
    Pool (Either ConnectionError Connection) -> IO ()
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 (Int -> UsageError -> ShowS
[UsageError] -> ShowS
UsageError -> String
(Int -> UsageError -> ShowS)
-> (UsageError -> String)
-> ([UsageError] -> ShowS)
-> Show UsageError
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [UsageError] -> ShowS
$cshowList :: [UsageError] -> ShowS
show :: UsageError -> String
$cshow :: UsageError -> String
showsPrec :: Int -> UsageError -> ShowS
$cshowsPrec :: Int -> UsageError -> ShowS
Show, UsageError -> UsageError -> Bool
(UsageError -> UsageError -> Bool)
-> (UsageError -> UsageError -> Bool) -> Eq UsageError
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 :: Pool -> Session a -> IO (Either UsageError a)
use = Maybe ObserverAction
-> Pool -> Session a -> IO (Either UsageError a)
forall a.
Maybe ObserverAction
-> Pool -> Session a -> IO (Either UsageError a)
useWithObserver Maybe ObserverAction
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 :: Maybe ObserverAction
-> Pool -> Session a -> IO (Either UsageError a)
useWithObserver Maybe ObserverAction
observer (Pool Pool (Either ConnectionError Connection)
pool) Session a
session =
    (Either ConnectionError (Either QueryError a)
 -> Either UsageError a)
-> IO (Either ConnectionError (Either QueryError a))
-> IO (Either UsageError a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((ConnectionError -> Either UsageError a)
-> (Either QueryError a -> Either UsageError a)
-> Either ConnectionError (Either QueryError a)
-> Either UsageError a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (UsageError -> Either UsageError a
forall a b. a -> Either a b
Left (UsageError -> Either UsageError a)
-> (ConnectionError -> UsageError)
-> ConnectionError
-> Either UsageError a
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) ((QueryError -> Either UsageError a)
-> (a -> Either UsageError a)
-> Either QueryError a
-> Either UsageError a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (UsageError -> Either UsageError a
forall a b. a -> Either a b
Left (UsageError -> Either UsageError a)
-> (QueryError -> UsageError) -> QueryError -> Either UsageError a
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) a -> Either UsageError a
forall a b. b -> Either a b
Right)) (IO (Either ConnectionError (Either QueryError a))
 -> IO (Either UsageError a))
-> IO (Either ConnectionError (Either QueryError a))
-> IO (Either UsageError a)
forall a b. (a -> b) -> a -> b
$
    Pool (Either ConnectionError Connection)
-> (Either ConnectionError Connection
    -> IO (Either ConnectionError (Either QueryError a)))
-> IO (Either ConnectionError (Either QueryError a))
forall resource failure success.
Pool resource
-> (resource -> IO (Either failure success))
-> IO (Either failure success)
ResourcePool.withResourceOnEither Pool (Either ConnectionError Connection)
pool ((Either ConnectionError Connection
  -> IO (Either ConnectionError (Either QueryError a)))
 -> IO (Either ConnectionError (Either QueryError a)))
-> (Either ConnectionError Connection
    -> IO (Either ConnectionError (Either QueryError a)))
-> IO (Either ConnectionError (Either QueryError a))
forall a b. (a -> b) -> a -> b
$
    (Connection -> IO (Either QueryError a))
-> Either ConnectionError Connection
-> IO (Either ConnectionError (Either QueryError a))
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 = IO (Either QueryError a)
-> (ObserverAction -> IO (Either QueryError a))
-> Maybe ObserverAction
-> IO (Either QueryError a)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe IO (Either QueryError a)
action (IO (Either QueryError a)
-> ObserverAction -> IO (Either QueryError a)
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 = Session a -> Connection -> IO (Either QueryError a)
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 :: p
nsRatio  = p
1000000000
                observed :: Observed
observed = Observed :: Rational -> Observed
Observed {   latency :: Rational
latency = Rational -> Rational
forall a. Real a => a -> Rational
toRational (TimeSpec -> Integer
toNanoSecs (TimeSpec
end TimeSpec -> TimeSpec -> TimeSpec
`diffTimeSpec` TimeSpec
start) Integer -> Integer -> Rational
forall a. Integral a => a -> a -> Ratio a
% Integer
forall p. Num p => p
nsRatio)
                                    }
            Observed -> IO a
doObserve Observed
observed IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> b -> IO 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) = Pool (Either ConnectionError Connection) -> Bool -> IO Stats
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 Int
getPoolUsageStat Pool
pool = Pool -> IO Stats
getPoolStats Pool
pool IO Stats -> (Stats -> IO Int) -> IO Int
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Stats -> IO Int
forall (f :: * -> *). Applicative f => Stats -> f Int
gather where
    gather :: Stats -> f Int
gather = Int -> f Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> f Int) -> (Stats -> Int) -> Stats -> f Int
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. PoolStats -> Int
ResourcePool.currentUsage (PoolStats -> Int) -> (Stats -> PoolStats) -> Stats -> Int
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