{-# 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 =
  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 = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a r. Pool a -> (a -> IO r) -> IO r
withResource Pool Connection
pool (\Connection
conn -> 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
  forall params (m :: * -> *).
(ToRow params, MonadIO m) =>
QueryNature -> Query -> params -> DBT m ()
logQueryFormat QueryNature
queryNature Query
q params
params
  forall a. [a] -> Vector a
V.fromList forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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
  forall params (m :: * -> *).
(ToRow params, MonadIO m) =>
QueryNature -> Query -> params -> DBT m ()
logQueryFormat QueryNature
queryNature Query
q ()
  forall a. [a] -> Vector a
V.fromList forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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
  forall params (m :: * -> *).
(ToRow params, MonadIO m) =>
QueryNature -> Query -> params -> DBT m ()
logQueryFormat QueryNature
queryNature Query
q params
params
  forall a. [a] -> Maybe a
listToMaybe forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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
  forall params (m :: * -> *).
(ToRow params, MonadIO m) =>
QueryNature -> Query -> params -> DBT m ()
logQueryFormat QueryNature
queryNature Query
q ()
  forall a. [a] -> Maybe a
listToMaybe forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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
  forall params (m :: * -> *).
(ToRow params, MonadIO m) =>
QueryNature -> Query -> params -> DBT m ()
logQueryFormat QueryNature
queryNature Query
q params
params
  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
  forall params (m :: * -> *).
(ToRow params, MonadIO m) =>
QueryNature -> Query -> [params] -> DBT m ()
logQueryFormatMany QueryNature
queryNature Query
q [params]
params
  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 = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ Handle -> ByteString -> IO ()
BS.hPutStrLn Handle
stdout (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 <- forall q (m :: * -> *).
(ToRow q, MonadIO m) =>
Query -> q -> DBT m ByteString
PGT.formatQuery Query
q params
params
  case QueryNature
queryNature of
    QueryNature
Select -> forall (m :: * -> *).
MonadIO m =>
ByteString -> ByteString -> DBT m ()
displayColoured forall str. IsString str => str
cyan ByteString
msg
    QueryNature
Update -> forall (m :: * -> *).
MonadIO m =>
ByteString -> ByteString -> DBT m ()
displayColoured forall str. IsString str => str
yellow ByteString
msg
    QueryNature
Insert -> forall (m :: * -> *).
MonadIO m =>
ByteString -> ByteString -> DBT m ()
displayColoured forall str. IsString str => str
yellow ByteString
msg
    QueryNature
Delete -> forall (m :: * -> *).
MonadIO m =>
ByteString -> ByteString -> DBT m ()
displayColoured 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 <- forall q (m :: * -> *).
(ToRow q, MonadIO m) =>
Query -> [q] -> DBT m ByteString
formatMany Query
q [params]
params
  case QueryNature
queryNature of
    QueryNature
Select -> forall (m :: * -> *).
MonadIO m =>
ByteString -> ByteString -> DBT m ()
displayColoured forall str. IsString str => str
cyan ByteString
msg
    QueryNature
Update -> forall (m :: * -> *).
MonadIO m =>
ByteString -> ByteString -> DBT m ()
displayColoured forall str. IsString str => str
yellow ByteString
msg
    QueryNature
Insert -> forall (m :: * -> *).
MonadIO m =>
ByteString -> ByteString -> DBT m ()
displayColoured forall str. IsString str => str
yellow ByteString
msg
    QueryNature
Delete -> forall (m :: * -> *).
MonadIO m =>
ByteString -> ByteString -> DBT m ()
displayColoured 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 = forall (m :: * -> *). Monad m => DBT m Connection
PGT.getConnection forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Connection
conn -> forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ 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
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
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)