{-|
Module: Squeal.PostgreSQL.Pool
Description: connection pools
Copyright: (c) Eitan Chatav, 2019
Maintainer: [email protected]
Stability: experimental

Connection pools.

Typical use case would be to create your pool using `createConnectionPool`
and run anything that requires the pool connection with `usingConnectionPool`.

Here's a simplified example:

>>> import Squeal.PostgreSQL

>>> :{
do
  let
    qry :: Query_ (Public '[]) () (Only Char)
    qry = values_ (inline 'a' `as` #fromOnly)
  pool <- createConnectionPool "host=localhost port=5432 dbname=exampledb user=postgres password=postgres" 1 0.5 10
  chr <- usingConnectionPool pool $ do
    result <- runQuery qry
    Just (Only a) <- firstRow result
    return a
  destroyConnectionPool pool
  putChar chr
:}
a
-}

{-# LANGUAGE
    CPP
  , DeriveFunctor
  , FlexibleContexts
  , FlexibleInstances
  , InstanceSigs
  , MultiParamTypeClasses
  , PolyKinds
  , RankNTypes
  , ScopedTypeVariables
  , TypeFamilies
  , DataKinds
  , PolyKinds
  , UndecidableInstances
#-}

module Squeal.PostgreSQL.Session.Pool
  ( -- * Pool
    Pool
  , createConnectionPool
  , usingConnectionPool
  , destroyConnectionPool
  ) where

import Control.Monad.Catch
import Control.Monad.IO.Class
import Data.ByteString
import Data.Time
import Data.Pool

import Squeal.PostgreSQL.Type.Schema
import Squeal.PostgreSQL.Session (PQ (..))
import Squeal.PostgreSQL.Session.Connection

-- | Create a striped pool of connections.
-- Although the garbage collector will destroy all idle connections when the pool is garbage collected it's recommended to manually `destroyConnectionPool` when you're done with the pool so that the connections are freed up as soon as possible.
createConnectionPool
  :: forall (db :: SchemasType) io. MonadIO io
  => ByteString
  -- ^ The passed string can be empty to use all default parameters, or it can
  -- contain one or more parameter settings separated by whitespace.
  -- Each parameter setting is in the form keyword = value. Spaces around the equal
  -- sign are optional. To write an empty value or a value containing spaces,
  -- surround it with single quotes, e.g., keyword = 'a value'. Single quotes and
  -- backslashes within the value must be escaped with a backslash, i.e., ' and \.
  -> Int
  -- ^ The number of stripes (distinct sub-pools) to maintain. The smallest acceptable value is 1.
  -> NominalDiffTime
  -- ^ Amount of time for which an unused connection is kept open. The smallest acceptable value is 0.5 seconds.
  -- The elapsed time before destroying a connection may be a little longer than requested, as the reaper thread wakes at 1-second intervals.
  -> Int
  -- ^ Maximum number of connections to keep open per stripe. The smallest acceptable value is 1.
  -- Requests for connections will block if this limit is reached on a single stripe, even if other stripes have idle connections available.
  -> io (Pool (K Connection db))
createConnectionPool :: forall (db :: SchemasType) (io :: * -> *).
MonadIO io =>
ByteString
-> Int -> NominalDiffTime -> Int -> io (Pool (K Connection db))
createConnectionPool ByteString
conninfo Int
stripes NominalDiffTime
idle Int
maxResrc =
#if MIN_VERSION_resource_pool(0,4,0)
  IO (Pool (K Connection db)) -> io (Pool (K Connection db))
forall a. IO a -> io a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Pool (K Connection db)) -> io (Pool (K Connection db)))
-> (PoolConfig (K Connection db) -> IO (Pool (K Connection db)))
-> PoolConfig (K Connection db)
-> io (Pool (K Connection db))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PoolConfig (K Connection db) -> IO (Pool (K Connection db))
forall a. PoolConfig a -> IO (Pool a)
newPool (PoolConfig (K Connection db) -> io (Pool (K Connection db)))
-> PoolConfig (K Connection db) -> io (Pool (K Connection db))
forall a b. (a -> b) -> a -> b
$ Maybe Int
-> PoolConfig (K Connection db) -> PoolConfig (K Connection db)
forall a. Maybe Int -> PoolConfig a -> PoolConfig a
setNumStripes
    (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
stripes)
    (IO (K Connection db)
-> (K Connection db -> IO ())
-> Double
-> Int
-> PoolConfig (K Connection db)
forall a. IO a -> (a -> IO ()) -> Double -> Int -> PoolConfig a
defaultPoolConfig (ByteString -> IO (K Connection db)
forall (db :: SchemasType) (io :: * -> *).
MonadIO io =>
ByteString -> io (K Connection db)
connectdb ByteString
conninfo) K Connection db -> IO ()
forall {k} (io :: * -> *) (db :: k).
MonadIO io =>
K Connection db -> io ()
finish (NominalDiffTime -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac NominalDiffTime
idle) Int
maxResrc)
#else
  liftIO $ createPool (connectdb conninfo) finish stripes idle maxResrc
#endif

{-|
Temporarily take a connection from a `Pool`, perform an action with it,
and return it to the pool afterwards.

If the pool has an idle connection available, it is used immediately.
Otherwise, if the maximum number of connections has not yet been reached,
a new connection is created and used.
If the maximum number of connections has been reached, this function blocks
until a connection becomes available.
-}
usingConnectionPool
  :: (MonadIO io, MonadMask io)
  => Pool (K Connection db) -- ^ pool
  -> PQ db db io x -- ^ session
  -> io x
usingConnectionPool :: forall (io :: * -> *) (db :: SchemasType) x.
(MonadIO io, MonadMask io) =>
Pool (K Connection db) -> PQ db db io x -> io x
usingConnectionPool Pool (K Connection db)
pool (PQ K Connection db -> io (K x db)
session) = ((forall a. io a -> io a) -> io x) -> io x
forall b.
HasCallStack =>
((forall a. io a -> io a) -> io b) -> io b
forall (m :: * -> *) b.
(MonadMask m, HasCallStack) =>
((forall a. m a -> m a) -> m b) -> m b
mask (((forall a. io a -> io a) -> io x) -> io x)
-> ((forall a. io a -> io a) -> io x) -> io x
forall a b. (a -> b) -> a -> b
$ \forall a. io a -> io a
restore -> do
  (K Connection db
conn, LocalPool (K Connection db)
local) <- IO (K Connection db, LocalPool (K Connection db))
-> io (K Connection db, LocalPool (K Connection db))
forall a. IO a -> io a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (K Connection db, LocalPool (K Connection db))
 -> io (K Connection db, LocalPool (K Connection db)))
-> IO (K Connection db, LocalPool (K Connection db))
-> io (K Connection db, LocalPool (K Connection db))
forall a b. (a -> b) -> a -> b
$ Pool (K Connection db)
-> IO (K Connection db, LocalPool (K Connection db))
forall a. Pool a -> IO (a, LocalPool a)
takeResource Pool (K Connection db)
pool
  K x db
ret <- io (K x db) -> io (K x db)
forall a. io a -> io a
restore (K Connection db -> io (K x db)
session K Connection db
conn) io (K x db) -> io () -> io (K x db)
forall (m :: * -> *) a b.
(HasCallStack, MonadCatch m) =>
m a -> m b -> m a
`onException`
            IO () -> io ()
forall a. IO a -> io a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Pool (K Connection db)
-> LocalPool (K Connection db) -> K Connection db -> IO ()
forall a. Pool a -> LocalPool a -> a -> IO ()
destroyResource Pool (K Connection db)
pool LocalPool (K Connection db)
local K Connection db
conn)
  IO () -> io ()
forall a. IO a -> io a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> io ()) -> IO () -> io ()
forall a b. (a -> b) -> a -> b
$ LocalPool (K Connection db) -> K Connection db -> IO ()
forall a. LocalPool a -> a -> IO ()
putResource LocalPool (K Connection db)
local K Connection db
conn
  x -> io x
forall a. a -> io a
forall (m :: * -> *) a. Monad m => a -> m a
return (x -> io x) -> x -> io x
forall a b. (a -> b) -> a -> b
$ K x db -> x
forall {k} a (b :: k). K a b -> a
unK K x db
ret

{- |
Destroy all connections in all stripes in the pool.
Note that this will ignore any exceptions in the destroy function.

This function is useful when you detect that all connections
in the pool are broken. For example after a database has been
restarted all connections opened before the restart will be broken.
In that case it's better to close those connections so that
`usingConnectionPool` won't take a broken connection from the pool
but will open a new connection instead.

Another use-case for this function is that when you know you are done
with the pool you can destroy all idle connections immediately
instead of waiting on the garbage collector to destroy them,
thus freeing up those connections sooner.
-}
destroyConnectionPool
  :: MonadIO io
  => Pool (K Connection db) -- ^ pool
  -> io ()
destroyConnectionPool :: forall {k} (io :: * -> *) (db :: k).
MonadIO io =>
Pool (K Connection db) -> io ()
destroyConnectionPool = IO () -> io ()
forall a. IO a -> io a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> io ())
-> (Pool (K Connection db) -> IO ())
-> Pool (K Connection db)
-> io ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pool (K Connection db) -> IO ()
forall a. Pool a -> IO ()
destroyAllResources