{-|
  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
  , withPool'
  , execute
  , query
  , query_
  , queryOne
  , QueryNature(..)
  ) where

import Colourista.IO (cyanMessage, redMessage, yellowMessage)
import Control.Monad.IO.Class
import Control.Monad.Trans.Control (MonadBaseControl)
import Data.Int
import Data.Maybe (listToMaybe)
import Data.Pool (Pool, createPool, withResource)
import Data.Text.Encoding (decodeUtf8)
import Data.Time (NominalDiffTime)
import Data.Vector (Vector)
import qualified Data.Vector as V

import Control.Monad.Catch (Exception, MonadCatch, try)
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
withPool :: (MonadBaseControl IO m)
         => Pool Connection -> PGT.DBT m a -> m a
withPool :: Pool Connection -> DBT m a -> m a
withPool Pool Connection
pool DBT m a
action = Pool Connection -> (Connection -> m a) -> m a
forall (m :: * -> *) a b.
MonadBaseControl IO m =>
Pool a -> (a -> m b) -> m b
withResource Pool Connection
pool ((Connection -> m a) -> m a) -> (Connection -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ DBT m a -> Connection -> m a
forall (m :: * -> *) a.
MonadBaseControl IO m =>
DBT m a -> Connection -> m a
PGT.runDBTSerializable DBT m a
action

-- | Run a DBT action while handling errors as Exceptions.
--
-- This function wraps the DBT actions in a 'try', so that exceptions
-- raised will be converted to the Left branch of the Either.
--
-- @since 0.0.1.0
withPool' :: forall errorType result m
          . (Exception errorType, MonadCatch m, MonadBaseControl IO m)
         => Pool Connection
         -> PGT.DBT m result
         -> m (Either errorType result)
withPool' :: Pool Connection -> DBT m result -> m (Either errorType result)
withPool' Pool Connection
pool DBT m result
action = m result -> m (Either errorType result)
forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> m (Either e a)
try (m result -> m (Either errorType result))
-> m result -> m (Either errorType result)
forall a b. (a -> b) -> a -> b
$ Pool Connection -> DBT m result -> m result
forall (m :: * -> *) a.
MonadBaseControl IO m =>
Pool Connection -> DBT m a -> m a
withPool Pool Connection
pool DBT m result
action

-- | 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 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 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 params (m :: * -> *).
(ToRow params, MonadIO m) =>
QueryNature -> Query -> params -> DBT m ()
logQueryFormat QueryNature
queryNature Query
q params
params
  [result]
result <- 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
  Maybe result -> DBT m (Maybe result)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe result -> DBT m (Maybe result))
-> Maybe result -> DBT m (Maybe result)
forall a b. (a -> b) -> a -> b
$ [result] -> Maybe result
forall a. [a] -> Maybe a
listToMaybe [result]
result

-- | 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 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 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 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

logQueryFormat :: (ToRow params, MonadIO m) => QueryNature -> Query -> params -> PGT.DBT m ()
logQueryFormat :: 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 -> IO () -> DBT m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> DBT m ()) -> IO () -> DBT m ()
forall a b. (a -> b) -> a -> b
$ Text -> IO ()
cyanMessage   (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$ Text
"[SELECT] " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> ByteString -> Text
decodeUtf8 ByteString
msg
    QueryNature
Update -> IO () -> DBT m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> DBT m ()) -> IO () -> DBT m ()
forall a b. (a -> b) -> a -> b
$ Text -> IO ()
yellowMessage (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$ Text
"[UPDATE] " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> ByteString -> Text
decodeUtf8 ByteString
msg
    QueryNature
Insert -> IO () -> DBT m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> DBT m ()) -> IO () -> DBT m ()
forall a b. (a -> b) -> a -> b
$ Text -> IO ()
yellowMessage (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$ Text
"[INSERT] " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> ByteString -> Text
decodeUtf8 ByteString
msg
    QueryNature
Delete -> IO () -> DBT m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> DBT m ()) -> IO () -> DBT m ()
forall a b. (a -> b) -> a -> b
$ Text -> IO ()
redMessage    (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$ Text
"[DELETE] " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> ByteString -> Text
decodeUtf8 ByteString
msg

-- | 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)