{-# LANGUAGE OverloadedStrings #-}

-- | 'Arbiter.Core.MonadArbiter.MonadArbiter' primitives backed by postgresql-simple.
module Arbiter.Simple.MonadArbiter
  ( HasSimplePool (..)
  , SimpleConnectionPool (..)

    -- * MonadArbiter implementation
  , simpleExecuteQuery
  , simpleExecuteStatement
  , simpleWithDbTransaction
  , simpleWithConnection
  , simpleRunHandlerWithConnection
  ) where

import Arbiter.Core.Codec (Col (..), NullCol (..), runCodec)
import Arbiter.Core.Exceptions (throwInternal)
import Arbiter.Core.MonadArbiter hiding (Query (..))
import Arbiter.Core.MonadArbiter qualified as MA
import Control.Monad (void)
import Control.Monad.IO.Class (liftIO)
import Data.ByteString.Char8 qualified as BSC
import Data.Int (Int64)
import Data.Pool (Pool, withResource)
import Data.Text.Encoding qualified as T
import Database.PostgreSQL.Simple (Connection)
import Database.PostgreSQL.Simple qualified as PG
import Database.PostgreSQL.Simple.FromRow (RowParser, field)
import Database.PostgreSQL.Simple.ToField (Action, ToField (..), toField, toJSONField)
import Database.PostgreSQL.Simple.Types (PGArray (..), Query (..))
import UnliftIO (MonadUnliftIO, mask, onException, withRunInIO)

-- | Pool, pinned connection, and savepoint depth.
data SimpleConnectionPool = SimpleConnectionPool
  { SimpleConnectionPool -> Maybe (Pool Connection)
connectionPool :: Maybe (Pool Connection)
  , SimpleConnectionPool -> Maybe Connection
activeConn :: Maybe Connection
  , SimpleConnectionPool -> Int
transactionDepth :: Int
  }

-- | Ambient access to the connection pool.
class (Monad m) => HasSimplePool m where
  getSimplePool :: m SimpleConnectionPool
  localSimplePool :: (SimpleConnectionPool -> SimpleConnectionPool) -> m a -> m a

-- | Run a query, decoding rows.
simpleExecuteQuery
  :: (HasSimplePool m, MonadUnliftIO m)
  => MA.Query a
  -> m [a]
simpleExecuteQuery :: forall (m :: * -> *) a.
(HasSimplePool m, MonadUnliftIO m) =>
Query a -> m [a]
simpleExecuteQuery (MA.Query Text
sqlTemplate Params
params RowCodec a
codec) = do
  let sql :: Query
sql = ByteString -> Query
Query (ByteString -> Query) -> ByteString -> Query
forall a b. (a -> b) -> a -> b
$ Text -> ByteString
T.encodeUtf8 Text
sqlTemplate
      parser :: RowParser a
parser = (forall x. NullCol x -> RowParser x) -> RowCodec a -> RowParser a
forall (f :: * -> *) a.
Applicative f =>
(forall x. NullCol x -> f x) -> RowCodec a -> f a
runCodec NullCol x -> RowParser x
forall x. NullCol x -> RowParser x
interpretNullCol RowCodec a
codec
  (Connection -> m [a]) -> m [a]
forall (m :: * -> *) a.
(HasSimplePool m, MonadUnliftIO m) =>
(Connection -> m a) -> m a
withConn ((Connection -> m [a]) -> m [a]) -> (Connection -> m [a]) -> m [a]
forall a b. (a -> b) -> a -> b
$ \Connection
conn -> 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
$ case Params
params of
    [] -> RowParser a -> Connection -> Query -> IO [a]
forall r. RowParser r -> Connection -> Query -> IO [r]
PG.queryWith_ RowParser a
parser Connection
conn Query
sql
    Params
_ -> RowParser a -> Connection -> Query -> [Action] -> IO [a]
forall q r.
ToRow q =>
RowParser r -> Connection -> Query -> q -> IO [r]
PG.queryWith RowParser a
parser Connection
conn Query
sql ((SomeParam -> Action) -> Params -> [Action]
forall a b. (a -> b) -> [a] -> [b]
map SomeParam -> Action
someParamToAction Params
params)

-- | Run a statement, returning rows affected.
simpleExecuteStatement
  :: (HasSimplePool m, MonadUnliftIO m)
  => MA.Query a
  -> m Int64
simpleExecuteStatement :: forall (m :: * -> *) a.
(HasSimplePool m, MonadUnliftIO m) =>
Query a -> m Int64
simpleExecuteStatement (MA.Query Text
sqlTemplate Params
params RowCodec a
_) = do
  let sql :: Query
sql = ByteString -> Query
Query (ByteString -> Query) -> ByteString -> Query
forall a b. (a -> b) -> a -> b
$ Text -> ByteString
T.encodeUtf8 Text
sqlTemplate
  (Connection -> m Int64) -> m Int64
forall (m :: * -> *) a.
(HasSimplePool m, MonadUnliftIO m) =>
(Connection -> m a) -> m a
withConn ((Connection -> m Int64) -> m Int64)
-> (Connection -> m Int64) -> m Int64
forall a b. (a -> b) -> a -> b
$ \Connection
conn -> IO Int64 -> m Int64
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int64 -> m Int64) -> IO Int64 -> m Int64
forall a b. (a -> b) -> a -> b
$ case Params
params of
    [] -> Connection -> Query -> IO Int64
PG.execute_ Connection
conn Query
sql
    Params
_ -> Connection -> Query -> [Action] -> IO Int64
forall q. ToRow q => Connection -> Query -> q -> IO Int64
PG.execute Connection
conn Query
sql ((SomeParam -> Action) -> Params -> [Action]
forall a b. (a -> b) -> [a] -> [b]
map SomeParam -> Action
someParamToAction Params
params)

interpretNullCol :: NullCol a -> RowParser a
interpretNullCol :: forall x. NullCol x -> RowParser x
interpretNullCol (NotNull Text
_ Col a
col) = Col a -> RowParser a
forall a. Col a -> RowParser a
colField Col a
col
interpretNullCol (Nullable Text
_ Col a1
col) = Col a1 -> RowParser (Maybe a1)
forall a. Col a -> RowParser (Maybe a)
colFieldNullable Col a1
col

colField :: Col a -> RowParser a
colField :: forall a. Col a -> RowParser a
colField Col a
CInt4 = RowParser a
forall a. FromField a => RowParser a
field
colField Col a
CInt8 = RowParser a
forall a. FromField a => RowParser a
field
colField Col a
CText = RowParser a
forall a. FromField a => RowParser a
field
colField Col a
CBool = RowParser a
forall a. FromField a => RowParser a
field
colField Col a
CTimestamptz = RowParser a
forall a. FromField a => RowParser a
field
colField Col a
CJsonb = RowParser a
forall a. FromField a => RowParser a
field
colField Col a
CFloat8 = RowParser a
forall a. FromField a => RowParser a
field
colField Col a
CUuid = RowParser a
forall a. FromField a => RowParser a
field

colFieldNullable :: Col a -> RowParser (Maybe a)
colFieldNullable :: forall a. Col a -> RowParser (Maybe a)
colFieldNullable Col a
CInt4 = RowParser (Maybe a)
forall a. FromField a => RowParser a
field
colFieldNullable Col a
CInt8 = RowParser (Maybe a)
forall a. FromField a => RowParser a
field
colFieldNullable Col a
CText = RowParser (Maybe a)
forall a. FromField a => RowParser a
field
colFieldNullable Col a
CBool = RowParser (Maybe a)
forall a. FromField a => RowParser a
field
colFieldNullable Col a
CTimestamptz = RowParser (Maybe a)
forall a. FromField a => RowParser a
field
colFieldNullable Col a
CJsonb = RowParser (Maybe a)
forall a. FromField a => RowParser a
field
colFieldNullable Col a
CFloat8 = RowParser (Maybe a)
forall a. FromField a => RowParser a
field
colFieldNullable Col a
CUuid = RowParser (Maybe a)
forall a. FromField a => RowParser a
field

-- | Transaction bracket. Nests via savepoints.
simpleWithDbTransaction
  :: (HasSimplePool m, MonadUnliftIO m)
  => m a
  -> m a
simpleWithDbTransaction :: forall (m :: * -> *) a.
(HasSimplePool m, MonadUnliftIO m) =>
m a -> m a
simpleWithDbTransaction m a
action = do
  pool <- m SimpleConnectionPool
forall (m :: * -> *). HasSimplePool m => m SimpleConnectionPool
getSimplePool
  let depth = SimpleConnectionPool -> Int
transactionDepth SimpleConnectionPool
pool
  case (activeConn pool, depth) of
    (Maybe Connection
Nothing, Int
_) -> case SimpleConnectionPool -> Maybe (Pool Connection)
connectionPool SimpleConnectionPool
pool of
      Maybe (Pool Connection)
Nothing -> Text -> m a
forall (m :: * -> *) a. MonadIO m => Text -> m a
throwInternal Text
"No active connection and no connection pool available"
      Just Pool Connection
connPool -> ((forall a. m a -> IO a) -> IO a) -> m a
forall b. ((forall a. m a -> IO a) -> IO b) -> m b
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO (((forall a. m a -> IO a) -> IO a) -> m a)
-> ((forall a. m a -> IO a) -> IO a) -> m a
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run ->
        Pool Connection -> (Connection -> IO a) -> IO a
forall a r. Pool a -> (a -> IO r) -> IO r
withResource Pool Connection
connPool ((Connection -> IO a) -> IO a) -> (Connection -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Connection
conn ->
          Connection -> IO a -> IO a
forall a. Connection -> IO a -> IO a
PG.withTransaction Connection
conn
            (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$ m a -> IO a
forall a. m a -> IO a
run
            (m a -> IO a) -> m a -> IO a
forall a b. (a -> b) -> a -> b
$ (SimpleConnectionPool -> SimpleConnectionPool) -> m a -> m a
forall a.
(SimpleConnectionPool -> SimpleConnectionPool) -> m a -> m a
forall (m :: * -> *) a.
HasSimplePool m =>
(SimpleConnectionPool -> SimpleConnectionPool) -> m a -> m a
localSimplePool (\SimpleConnectionPool
spool -> SimpleConnectionPool
spool {activeConn = Just conn, transactionDepth = 1}) m a
action
    (Just Connection
conn, Int
0) -> ((forall a. m a -> IO a) -> IO a) -> m a
forall b. ((forall a. m a -> IO a) -> IO b) -> m b
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO (((forall a. m a -> IO a) -> IO a) -> m a)
-> ((forall a. m a -> IO a) -> IO a) -> m a
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run ->
      Connection -> IO a -> IO a
forall a. Connection -> IO a -> IO a
PG.withTransaction Connection
conn
        (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$ m a -> IO a
forall a. m a -> IO a
run
        (m a -> IO a) -> m a -> IO a
forall a b. (a -> b) -> a -> b
$ (SimpleConnectionPool -> SimpleConnectionPool) -> m a -> m a
forall a.
(SimpleConnectionPool -> SimpleConnectionPool) -> m a -> m a
forall (m :: * -> *) a.
HasSimplePool m =>
(SimpleConnectionPool -> SimpleConnectionPool) -> m a -> m a
localSimplePool (\SimpleConnectionPool
spool -> SimpleConnectionPool
spool {transactionDepth = 1}) m a
action
    (Just Connection
conn, Int
_) -> ((forall a. m a -> m a) -> m a) -> m a
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> m a) -> m b) -> m b
mask (((forall a. m a -> m a) -> m a) -> m a)
-> ((forall a. m a -> m a) -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> m a
restore -> do
      let spName :: Query
spName = ByteString -> Query
Query (ByteString -> Query) -> ByteString -> Query
forall a b. (a -> b) -> a -> b
$ ByteString
"arbiter_sp_" ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> String -> ByteString
BSC.pack (Int -> String
forall a. Show a => a -> String
show Int
depth)
      m Int64 -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m Int64 -> m ()) -> (IO Int64 -> m Int64) -> IO Int64 -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO Int64 -> m Int64
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int64 -> m ()) -> IO Int64 -> m ()
forall a b. (a -> b) -> a -> b
$ Connection -> Query -> IO Int64
PG.execute_ Connection
conn (Query -> IO Int64) -> Query -> IO Int64
forall a b. (a -> b) -> a -> b
$ Query
"SAVEPOINT " Query -> Query -> Query
forall a. Semigroup a => a -> a -> a
<> Query
spName
      result <-
        m a -> m a
forall a. m a -> m a
restore ((SimpleConnectionPool -> SimpleConnectionPool) -> m a -> m a
forall a.
(SimpleConnectionPool -> SimpleConnectionPool) -> m a -> m a
forall (m :: * -> *) a.
HasSimplePool m =>
(SimpleConnectionPool -> SimpleConnectionPool) -> m a -> m a
localSimplePool (\SimpleConnectionPool
spool -> SimpleConnectionPool
spool {transactionDepth = depth + 1}) m a
action)
          m a -> m Int64 -> m a
forall (m :: * -> *) a b. MonadUnliftIO m => m a -> m b -> m a
`onException` IO Int64 -> m Int64
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Connection -> Query -> IO Int64
PG.execute_ Connection
conn (Query -> IO Int64) -> Query -> IO Int64
forall a b. (a -> b) -> a -> b
$ Query
"ROLLBACK TO SAVEPOINT " Query -> Query -> Query
forall a. Semigroup a => a -> a -> a
<> Query
spName)
      void . liftIO $ PG.execute_ conn $ "RELEASE SAVEPOINT " <> spName
      pure result

-- | Pin one pooled connection for the action.
simpleWithConnection
  :: (HasSimplePool m, MonadUnliftIO m)
  => m a
  -> m a
simpleWithConnection :: forall (m :: * -> *) a.
(HasSimplePool m, MonadUnliftIO m) =>
m a -> m a
simpleWithConnection m a
action = do
  pool <- m SimpleConnectionPool
forall (m :: * -> *). HasSimplePool m => m SimpleConnectionPool
getSimplePool
  case (activeConn pool, connectionPool pool) of
    (Just Connection
_, Maybe (Pool Connection)
_) -> m a
action
    (Maybe Connection
Nothing, Just Pool Connection
connPool) -> ((forall a. m a -> IO a) -> IO a) -> m a
forall b. ((forall a. m a -> IO a) -> IO b) -> m b
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO (((forall a. m a -> IO a) -> IO a) -> m a)
-> ((forall a. m a -> IO a) -> IO a) -> m a
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run ->
      Pool Connection -> (Connection -> IO a) -> IO a
forall a r. Pool a -> (a -> IO r) -> IO r
withResource Pool Connection
connPool ((Connection -> IO a) -> IO a) -> (Connection -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Connection
conn ->
        m a -> IO a
forall a. m a -> IO a
run (m a -> IO a) -> m a -> IO a
forall a b. (a -> b) -> a -> b
$ (SimpleConnectionPool -> SimpleConnectionPool) -> m a -> m a
forall a.
(SimpleConnectionPool -> SimpleConnectionPool) -> m a -> m a
forall (m :: * -> *) a.
HasSimplePool m =>
(SimpleConnectionPool -> SimpleConnectionPool) -> m a -> m a
localSimplePool (\SimpleConnectionPool
spool -> SimpleConnectionPool
spool {activeConn = Just conn}) m a
action
    (Maybe Connection
Nothing, Maybe (Pool Connection)
Nothing) -> Text -> m a
forall (m :: * -> *) a. MonadIO m => Text -> m a
throwInternal Text
"No active connection and no connection pool available"

-- | Run a handler on the pinned connection.
simpleRunHandlerWithConnection
  :: (HasSimplePool m, MonadUnliftIO m)
  => (Connection -> job -> m result)
  -> job
  -> m result
simpleRunHandlerWithConnection :: forall (m :: * -> *) job result.
(HasSimplePool m, MonadUnliftIO m) =>
(Connection -> job -> m result) -> job -> m result
simpleRunHandlerWithConnection Connection -> job -> m result
handler job
job =
  (Connection -> m result) -> m result
forall (m :: * -> *) a.
(HasSimplePool m, MonadUnliftIO m) =>
(Connection -> m a) -> m a
withConn ((Connection -> m result) -> m result)
-> (Connection -> m result) -> m result
forall a b. (a -> b) -> a -> b
$ \Connection
conn -> Connection -> job -> m result
handler Connection
conn job
job

withConn
  :: (HasSimplePool m, MonadUnliftIO m)
  => (Connection -> m a)
  -> m a
withConn :: forall (m :: * -> *) a.
(HasSimplePool m, MonadUnliftIO m) =>
(Connection -> m a) -> m a
withConn Connection -> m a
action = do
  pool <- m SimpleConnectionPool
forall (m :: * -> *). HasSimplePool m => m SimpleConnectionPool
getSimplePool
  case (activeConn pool, connectionPool pool) of
    (Just Connection
conn, Maybe (Pool Connection)
_) -> Connection -> m a
action Connection
conn
    (Maybe Connection
Nothing, Just Pool Connection
connPool) -> ((forall a. m a -> IO a) -> IO a) -> m a
forall b. ((forall a. m a -> IO a) -> IO b) -> m b
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO (((forall a. m a -> IO a) -> IO a) -> m a)
-> ((forall a. m a -> IO a) -> IO a) -> m a
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run ->
      Pool Connection -> (Connection -> IO a) -> IO a
forall a r. Pool a -> (a -> IO r) -> IO r
withResource Pool Connection
connPool ((Connection -> IO a) -> IO a) -> (Connection -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Connection
conn -> m a -> IO a
forall a. m a -> IO a
run (m a -> IO a) -> m a -> IO a
forall a b. (a -> b) -> a -> b
$ Connection -> m a
action Connection
conn
    (Maybe Connection
Nothing, Maybe (Pool Connection)
Nothing) -> Text -> m a
forall (m :: * -> *) a. MonadIO m => Text -> m a
throwInternal Text
"No active connection and no connection pool available"

someParamToAction :: SomeParam -> Action
someParamToAction :: SomeParam -> Action
someParamToAction (SomeParam (PScalar Col a
CJsonb) a
value) = a -> Action
forall a. ToJSON a => a -> Action
toJSONField a
value
someParamToAction (SomeParam (PScalar Col a
col) a
value) = Col a -> (ToField a => Action) -> Action
forall a r. Col a -> (ToField a => r) -> r
withColToField Col a
col (a -> Action
forall a. ToField a => a -> Action
toField a
value)
someParamToAction (SomeParam (PNullable Col a1
CJsonb) a
value) = Action -> (Value -> Action) -> Maybe Value -> Action
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Maybe Int -> Action
forall a. ToField a => a -> Action
toField (Maybe Int
forall a. Maybe a
Nothing :: Maybe Int)) Value -> Action
forall a. ToJSON a => a -> Action
toJSONField a
Maybe Value
value
someParamToAction (SomeParam (PNullable Col a1
col) a
value) = Col a1 -> (ToField a1 => Action) -> Action
forall a r. Col a -> (ToField a => r) -> r
withColToField Col a1
col (a -> Action
forall a. ToField a => a -> Action
toField a
value)
someParamToAction (SomeParam (PArray Col a1
col) a
value) = Col a1 -> (ToField a1 => Action) -> Action
forall a r. Col a -> (ToField a => r) -> r
withColToField Col a1
col (PGArray a1 -> Action
forall a. ToField a => a -> Action
toField ([a1] -> PGArray a1
forall a. [a] -> PGArray a
PGArray a
[a1]
value))
someParamToAction (SomeParam (PNullArray Col a1
col) a
value) = Col a1 -> (ToField a1 => Action) -> Action
forall a r. Col a -> (ToField a => r) -> r
withColToField Col a1
col (PGArray (Maybe a1) -> Action
forall a. ToField a => a -> Action
toField ([Maybe a1] -> PGArray (Maybe a1)
forall a. [a] -> PGArray a
PGArray a
[Maybe a1]
value))

withColToField :: Col a -> ((ToField a) => r) -> r
withColToField :: forall a r. Col a -> (ToField a => r) -> r
withColToField Col a
CInt4 ToField a => r
continuation = r
ToField a => r
continuation
withColToField Col a
CInt8 ToField a => r
continuation = r
ToField a => r
continuation
withColToField Col a
CText ToField a => r
continuation = r
ToField a => r
continuation
withColToField Col a
CBool ToField a => r
continuation = r
ToField a => r
continuation
withColToField Col a
CTimestamptz ToField a => r
continuation = r
ToField a => r
continuation
withColToField Col a
CJsonb ToField a => r
continuation = r
ToField a => r
continuation
withColToField Col a
CFloat8 ToField a => r
continuation = r
ToField a => r
continuation
withColToField Col a
CUuid ToField a => r
continuation = r
ToField a => r
continuation