{-# LANGUAGE CPP #-}

-- |
--  Module      : Database.PostgreSQL.Entity.DBT
--  Copyright   : © Clément Delafargue, 2018
--                  Théophile Choutri, 2021
--  License     : MIT
--  Maintainer  : theophile@choutri.eu
--  Stability   : stable
--
--  The 'Database.PostgreSQL.Transact.DBT' plumbing module to handle database queries and pools
module Database.PostgreSQL.Entity.DBT
  ( mkPool
  , withPool
  , execute
  , executeMany
  , query
  , query_
  , queryOne
  , queryOne_
  , QueryNature (..)
  )
where

#ifdef PROD
#else
import Colourista.IO (cyanMessage, redMessage, yellowMessage)
import Data.ByteString (ByteString)
import Data.Text.Encoding (decodeUtf8)
import qualified Database.PostgreSQL.Simple as Simple
#endif

import Control.Monad.IO.Class
import Data.Int
import Data.Maybe (listToMaybe)
import Data.Pool (Pool, createPool, withResource)
import Data.Time (NominalDiffTime)
import Data.Vector (Vector)
import qualified Data.Vector as V

import Database.PostgreSQL.Simple as PG (ConnectInfo, Connection, FromRow, Query, ToRow, close, connect)
import qualified Database.PostgreSQL.Transact as PGT

-- | Create a Pool Connection with the appropriate parameters
--
-- @since 0.0.1.0
mkPool ::
  ConnectInfo -> -- Database access information
  Int -> -- Number of sub-pools
  NominalDiffTime -> -- Allowed timeout
  Int -> -- Number of connections
  IO (Pool Connection)
mkPool :: ConnectInfo
-> Int -> NominalDiffTime -> Int -> IO (Pool Connection)
mkPool ConnectInfo
connectInfo Int
subPools NominalDiffTime
timeout Int
connections =
  IO Connection
-> (Connection -> IO ())
-> Int
-> NominalDiffTime
-> Int
-> IO (Pool Connection)
forall a.
IO a
-> (a -> IO ()) -> Int -> NominalDiffTime -> Int -> IO (Pool a)
createPool (ConnectInfo -> IO Connection
connect ConnectInfo
connectInfo) Connection -> IO ()
close Int
subPools NominalDiffTime
timeout Int
connections

-- | Run a DBT action with no explicit error handling.
--
-- This functions is suited for using 'MonadError' error handling.
--
-- === __Example__
--
-- > let e1 = E 1 True True
-- > result <- runExceptT @EntityError $ do
-- >   withPool pool $ insertEntity e1
-- >   withPool pool $ markForProcessing 1
-- > case result of
-- >   Left err -> print err
-- >   Right _  -> putStrLn "Everything went well"
--
-- See the code in the @example/@ directory on GitHub
--
-- @since 0.0.1.0
#if MIN_VERSION_resource_pool(0,3,0)
withPool :: (MonadIO m) => Pool Connection -> PGT.DBT IO a -> m a
withPool :: Pool Connection -> DBT IO a -> m a
withPool Pool Connection
pool DBT IO a
action = IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ Pool Connection -> (Connection -> IO a) -> IO a
forall a r. Pool a -> (a -> IO r) -> IO r
withResource Pool Connection
pool (\Connection
conn -> DBT IO a -> Connection -> IO a
forall (m :: * -> *) a.
MonadBaseControl IO m =>
DBT m a -> Connection -> m a
PGT.runDBTSerializable DBT IO a
action Connection
conn)
#else
withPool :: (MonadBaseControl IO m) => Pool Connection -> PGT.DBT m a -> m a
withPool pool action = withResource pool (\conn -> PGT.runDBTSerializable action conn)
#endif

-- | Query wrapper that returns a 'Vector' of results
--
-- @since 0.0.1.0
query ::
  (ToRow params, FromRow result, MonadIO m) =>
  QueryNature ->
  Query ->
  params ->
  PGT.DBT m (Vector result)
query :: QueryNature -> Query -> params -> DBT m (Vector result)
query QueryNature
queryNature Query
q params
params = do
  QueryNature -> Query -> params -> DBT m ()
forall (m :: * -> *) params.
Monad m =>
QueryNature -> Query -> params -> DBT m ()
logQueryFormat QueryNature
queryNature Query
q params
params
  [result] -> Vector result
forall a. [a] -> Vector a
V.fromList ([result] -> Vector result)
-> DBT m [result] -> DBT m (Vector result)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Query -> params -> DBT m [result]
forall a b (m :: * -> *).
(ToRow a, FromRow b, MonadIO m) =>
Query -> a -> DBT m [b]
PGT.query Query
q params
params

-- | Query wrapper that returns a 'Vector' of results and does not take an argument
--
-- @since 0.0.1.0
query_ ::
  (FromRow result, MonadIO m) =>
  QueryNature ->
  Query ->
  PGT.DBT m (Vector result)
query_ :: QueryNature -> Query -> DBT m (Vector result)
query_ QueryNature
queryNature Query
q = do
  QueryNature -> Query -> () -> DBT m ()
forall (m :: * -> *) params.
Monad m =>
QueryNature -> Query -> params -> DBT m ()
logQueryFormat QueryNature
queryNature Query
q ()
  [result] -> Vector result
forall a. [a] -> Vector a
V.fromList ([result] -> Vector result)
-> DBT m [result] -> DBT m (Vector result)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Query -> DBT m [result]
forall b (m :: * -> *).
(FromRow b, MonadIO m) =>
Query -> DBT m [b]
PGT.query_ Query
q

-- | Query wrapper that returns one result.
--
-- @since 0.0.1.0
queryOne ::
  (ToRow params, FromRow result, MonadIO m) =>
  QueryNature ->
  Query ->
  params ->
  PGT.DBT m (Maybe result)
queryOne :: QueryNature -> Query -> params -> DBT m (Maybe result)
queryOne QueryNature
queryNature Query
q params
params = do
  QueryNature -> Query -> params -> DBT m ()
forall (m :: * -> *) params.
Monad m =>
QueryNature -> Query -> params -> DBT m ()
logQueryFormat QueryNature
queryNature Query
q params
params
  [result] -> Maybe result
forall a. [a] -> Maybe a
listToMaybe ([result] -> Maybe result)
-> DBT m [result] -> DBT m (Maybe result)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Query -> params -> DBT m [result]
forall a b (m :: * -> *).
(ToRow a, FromRow b, MonadIO m) =>
Query -> a -> DBT m [b]
PGT.query Query
q params
params

--

-- | Query wrapper that returns one result and does not take an argument
--
-- @since 0.0.2.0
queryOne_ ::
  (FromRow result, MonadIO m) =>
  QueryNature ->
  Query ->
  PGT.DBT m (Maybe result)
queryOne_ :: QueryNature -> Query -> DBT m (Maybe result)
queryOne_ QueryNature
queryNature Query
q = do
  QueryNature -> Query -> () -> DBT m ()
forall (m :: * -> *) params.
Monad m =>
QueryNature -> Query -> params -> DBT m ()
logQueryFormat QueryNature
queryNature Query
q ()
  [result] -> Maybe result
forall a. [a] -> Maybe a
listToMaybe ([result] -> Maybe result)
-> DBT m [result] -> DBT m (Maybe result)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Query -> DBT m [result]
forall b (m :: * -> *).
(FromRow b, MonadIO m) =>
Query -> DBT m [b]
PGT.query_ Query
q

-- | Query wrapper for SQL statements which do not return.
--
-- @since 0.0.1.0
execute ::
  (ToRow params, MonadIO m) =>
  QueryNature ->
  Query ->
  params ->
  PGT.DBT m Int64
execute :: QueryNature -> Query -> params -> DBT m Int64
execute QueryNature
queryNature Query
q params
params = do
  QueryNature -> Query -> params -> DBT m ()
forall (m :: * -> *) params.
Monad m =>
QueryNature -> Query -> params -> DBT m ()
logQueryFormat QueryNature
queryNature Query
q params
params
  Query -> params -> DBT m Int64
forall q (m :: * -> *).
(ToRow q, MonadIO m) =>
Query -> q -> DBT m Int64
PGT.execute Query
q params
params

-- | Query wrapper for SQL statements that operate on multiple rows which do not return.
--
-- @since 0.0.2.0
executeMany ::
  (ToRow params, MonadIO m) =>
  QueryNature ->
  Query ->
  [params] ->
  PGT.DBT m Int64
executeMany :: QueryNature -> Query -> [params] -> DBT m Int64
executeMany QueryNature
queryNature Query
q [params]
params = do
  QueryNature -> Query -> [params] -> DBT m ()
forall (m :: * -> *) params.
Monad m =>
QueryNature -> Query -> [params] -> DBT m ()
logQueryFormatMany QueryNature
queryNature Query
q [params]
params
  Query -> [params] -> DBT m Int64
forall q (m :: * -> *).
(ToRow q, MonadIO m) =>
Query -> [q] -> DBT m Int64
PGT.executeMany Query
q [params]
params

#ifdef PROD
logQueryFormat :: (Monad m) => QueryNature -> Query -> params -> PGT.DBT m ()
logQueryFormat :: QueryNature -> Query -> params -> DBT m ()
logQueryFormat QueryNature
_ Query
_ params
_ = () -> DBT m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
#else
logQueryFormat :: (ToRow params, MonadIO m)
               => QueryNature -> Query -> params -> PGT.DBT m ()
logQueryFormat queryNature q params = do
  msg <- PGT.formatQuery q params
  case queryNature of
    Select -> liftIO $ cyanMessage   $ decodeUtf8 msg
    Update -> liftIO $ yellowMessage $ decodeUtf8 msg
    Insert -> liftIO $ yellowMessage $ decodeUtf8 msg
    Delete -> liftIO $ redMessage    $ decodeUtf8 msg
#endif

#ifdef PROD
logQueryFormatMany :: (Monad m) => QueryNature -> Query -> [params] -> PGT.DBT m ()
logQueryFormatMany :: QueryNature -> Query -> [params] -> DBT m ()
logQueryFormatMany QueryNature
_ Query
_ [params]
_ = () -> DBT m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
#else
logQueryFormatMany :: (ToRow params, MonadIO m) => QueryNature -> Query -> [params] -> PGT.DBT m ()
logQueryFormatMany queryNature q params = do
  msg <- formatMany q params
  case queryNature of
    Select -> liftIO $ cyanMessage   $ decodeUtf8 msg
    Update -> liftIO $ yellowMessage $ decodeUtf8 msg
    Insert -> liftIO $ yellowMessage $ decodeUtf8 msg
    Delete -> liftIO $ redMessage    $ decodeUtf8 msg

formatMany :: (ToRow q, MonadIO m) => Query -> [q] -> PGT.DBT m ByteString
formatMany q xs = PGT.getConnection >>= \conn -> liftIO $ Simple.formatMany conn q xs
#endif

-- | This sum type is given to the 'query', 'queryOne' and 'execute' functions to help
-- with logging.
--
-- @since 0.0.1.0
data QueryNature = Select | Insert | Update | Delete deriving (QueryNature -> QueryNature -> Bool
(QueryNature -> QueryNature -> Bool)
-> (QueryNature -> QueryNature -> Bool) -> Eq QueryNature
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: QueryNature -> QueryNature -> Bool
$c/= :: QueryNature -> QueryNature -> Bool
== :: QueryNature -> QueryNature -> Bool
$c== :: QueryNature -> QueryNature -> Bool
Eq, Int -> QueryNature -> ShowS
[QueryNature] -> ShowS
QueryNature -> String
(Int -> QueryNature -> ShowS)
-> (QueryNature -> String)
-> ([QueryNature] -> ShowS)
-> Show QueryNature
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [QueryNature] -> ShowS
$cshowList :: [QueryNature] -> ShowS
show :: QueryNature -> String
$cshow :: QueryNature -> String
showsPrec :: Int -> QueryNature -> ShowS
$cshowsPrec :: Int -> QueryNature -> ShowS
Show)