{-# 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 (cyan, red, yellow, formatWith)
import Data.ByteString (ByteString)
import qualified Database.PostgreSQL.Simple as Simple
import System.IO (stdout)
import qualified Data.ByteString.Char8 as BS
#endif

import Control.Monad.IO.Class
#if MIN_VERSION_resource_pool(0,3,0)
#else
import Control.Monad.Trans.Control
#endif
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 :: forall (m :: * -> *) a.
MonadIO m =>
Pool Connection -> DBT IO a -> m a
withPool Pool Connection
pool DBT IO a
action = IO a -> m a
forall a. 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 :: forall params result (m :: * -> *).
(ToRow params, FromRow result, MonadIO m) =>
QueryNature -> Query -> params -> DBT m (Vector result)
query QueryNature
queryNature Query
q params
params = do
  QueryNature -> Query -> params -> DBT m ()
forall params (m :: * -> *).
(ToRow params, MonadIO 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_ :: forall result (m :: * -> *).
(FromRow result, MonadIO m) =>
QueryNature -> Query -> DBT m (Vector result)
query_ QueryNature
queryNature Query
q = do
  QueryNature -> Query -> () -> DBT m ()
forall params (m :: * -> *).
(ToRow params, MonadIO 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 :: forall params result (m :: * -> *).
(ToRow params, FromRow result, MonadIO m) =>
QueryNature -> Query -> params -> DBT m (Maybe result)
queryOne QueryNature
queryNature Query
q params
params = do
  QueryNature -> Query -> params -> DBT m ()
forall params (m :: * -> *).
(ToRow params, MonadIO 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_ :: forall result (m :: * -> *).
(FromRow result, MonadIO m) =>
QueryNature -> Query -> DBT m (Maybe result)
queryOne_ QueryNature
queryNature Query
q = do
  QueryNature -> Query -> () -> DBT m ()
forall params (m :: * -> *).
(ToRow params, MonadIO 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 :: forall params (m :: * -> *).
(ToRow params, MonadIO m) =>
QueryNature -> Query -> params -> DBT m Int64
execute QueryNature
queryNature Query
q params
params = do
  QueryNature -> Query -> params -> DBT m ()
forall params (m :: * -> *).
(ToRow params, MonadIO 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 :: forall params (m :: * -> *).
(ToRow params, MonadIO m) =>
QueryNature -> Query -> [params] -> DBT m Int64
executeMany QueryNature
queryNature Query
q [params]
params = do
  QueryNature -> Query -> [params] -> DBT m ()
forall params (m :: * -> *).
(ToRow params, MonadIO 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

#ifndef PROD
displayColoured :: (MonadIO m) => ByteString -> ByteString -> PGT.DBT m ()
displayColoured :: forall (m :: * -> *).
MonadIO m =>
ByteString -> ByteString -> DBT m ()
displayColoured ByteString
colour ByteString
text = IO () -> DBT m ()
forall a. IO a -> DBT m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> DBT m ()) -> IO () -> DBT m ()
forall a b. (a -> b) -> a -> b
$ Handle -> ByteString -> IO ()
BS.hPutStrLn Handle
stdout ([ByteString] -> ByteString -> ByteString
forall str. (IsString str, Semigroup str) => [str] -> str -> str
formatWith [ByteString
colour] ByteString
text)
#endif

#ifdef PROD
logQueryFormat :: (Monad m) => QueryNature -> Query -> params -> PGT.DBT m ()
logQueryFormat _ _ _ = pure ()
#else
logQueryFormat :: (ToRow params, MonadIO m)
               => QueryNature -> Query -> params -> PGT.DBT m ()
logQueryFormat :: forall params (m :: * -> *).
(ToRow params, MonadIO m) =>
QueryNature -> Query -> params -> DBT m ()
logQueryFormat QueryNature
queryNature Query
q params
params = do
  ByteString
msg <- Query -> params -> DBT m ByteString
forall q (m :: * -> *).
(ToRow q, MonadIO m) =>
Query -> q -> DBT m ByteString
PGT.formatQuery Query
q params
params
  case QueryNature
queryNature of
    QueryNature
Select -> ByteString -> ByteString -> DBT m ()
forall (m :: * -> *).
MonadIO m =>
ByteString -> ByteString -> DBT m ()
displayColoured ByteString
forall str. IsString str => str
cyan ByteString
msg
    QueryNature
Update -> ByteString -> ByteString -> DBT m ()
forall (m :: * -> *).
MonadIO m =>
ByteString -> ByteString -> DBT m ()
displayColoured ByteString
forall str. IsString str => str
yellow ByteString
msg
    QueryNature
Insert -> ByteString -> ByteString -> DBT m ()
forall (m :: * -> *).
MonadIO m =>
ByteString -> ByteString -> DBT m ()
displayColoured ByteString
forall str. IsString str => str
yellow ByteString
msg
    QueryNature
Delete -> ByteString -> ByteString -> DBT m ()
forall (m :: * -> *).
MonadIO m =>
ByteString -> ByteString -> DBT m ()
displayColoured ByteString
forall str. IsString str => str
red ByteString
msg
#endif

#ifdef PROD
logQueryFormatMany :: (Monad m) => QueryNature -> Query -> [params] -> PGT.DBT m ()
logQueryFormatMany _ _ _ = pure ()
#else
logQueryFormatMany :: (ToRow params, MonadIO m) => QueryNature -> Query -> [params] -> PGT.DBT m ()
logQueryFormatMany :: forall params (m :: * -> *).
(ToRow params, MonadIO m) =>
QueryNature -> Query -> [params] -> DBT m ()
logQueryFormatMany QueryNature
queryNature Query
q [params]
params = do
  ByteString
msg <- Query -> [params] -> DBT m ByteString
forall q (m :: * -> *).
(ToRow q, MonadIO m) =>
Query -> [q] -> DBT m ByteString
formatMany Query
q [params]
params
  case QueryNature
queryNature of
    QueryNature
Select -> ByteString -> ByteString -> DBT m ()
forall (m :: * -> *).
MonadIO m =>
ByteString -> ByteString -> DBT m ()
displayColoured ByteString
forall str. IsString str => str
cyan ByteString
msg
    QueryNature
Update -> ByteString -> ByteString -> DBT m ()
forall (m :: * -> *).
MonadIO m =>
ByteString -> ByteString -> DBT m ()
displayColoured ByteString
forall str. IsString str => str
yellow ByteString
msg
    QueryNature
Insert -> ByteString -> ByteString -> DBT m ()
forall (m :: * -> *).
MonadIO m =>
ByteString -> ByteString -> DBT m ()
displayColoured ByteString
forall str. IsString str => str
yellow ByteString
msg
    QueryNature
Delete -> ByteString -> ByteString -> DBT m ()
forall (m :: * -> *).
MonadIO m =>
ByteString -> ByteString -> DBT m ()
displayColoured ByteString
forall str. IsString str => str
red ByteString
msg

formatMany :: (ToRow q, MonadIO m) => Query -> [q] -> PGT.DBT m ByteString
formatMany :: forall q (m :: * -> *).
(ToRow q, MonadIO m) =>
Query -> [q] -> DBT m ByteString
formatMany Query
q [q]
xs = DBT m Connection
forall (m :: * -> *). Monad m => DBT m Connection
PGT.getConnection DBT m Connection
-> (Connection -> DBT m ByteString) -> DBT m ByteString
forall a b. DBT m a -> (a -> DBT m b) -> DBT m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Connection
conn -> IO ByteString -> DBT m ByteString
forall a. IO a -> DBT m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ByteString -> DBT m ByteString)
-> IO ByteString -> DBT m ByteString
forall a b. (a -> b) -> a -> b
$ Connection -> Query -> [q] -> IO ByteString
forall q. ToRow q => Connection -> Query -> [q] -> IO ByteString
Simple.formatMany Connection
conn Query
q [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
$c== :: QueryNature -> QueryNature -> Bool
== :: QueryNature -> QueryNature -> Bool
$c/= :: QueryNature -> QueryNature -> Bool
/= :: 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
$cshowsPrec :: Int -> QueryNature -> ShowS
showsPrec :: Int -> QueryNature -> ShowS
$cshow :: QueryNature -> String
show :: QueryNature -> String
$cshowList :: [QueryNature] -> ShowS
showList :: [QueryNature] -> ShowS
Show)