{-# LANGUAGE OverloadedStrings #-}

module Arbiter.Simple.MonadArbiter
  ( HasSimplePool (..)
  , SimpleConnectionPool (..)

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

import Arbiter.Core.Codec (Col (..), NullCol (..), RowCodec, runCodec)
import Arbiter.Core.Exceptions (throwInternal)
import Arbiter.Core.MonadArbiter
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 (Text)
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)

data SimpleConnectionPool = SimpleConnectionPool
  { SimpleConnectionPool -> Maybe (Pool Connection)
connectionPool :: Maybe (Pool Connection)
  , SimpleConnectionPool -> Maybe Connection
activeConn :: Maybe Connection
  , SimpleConnectionPool -> Int
transactionDepth :: Int
  }

class (Monad m) => HasSimplePool m where
  getSimplePool :: m SimpleConnectionPool
  localSimplePool :: (SimpleConnectionPool -> SimpleConnectionPool) -> m a -> m a

simpleExecuteQuery
  :: (HasSimplePool m, MonadUnliftIO m)
  => Text
  -> Params
  -> RowCodec a
  -> m [a]
simpleExecuteQuery :: forall (m :: * -> *) a.
(HasSimplePool m, MonadUnliftIO m) =>
Text -> Params -> RowCodec a -> m [a]
simpleExecuteQuery 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
      pgParams :: [Action]
pgParams = (SomeParam -> Action) -> Params -> [Action]
forall a b. (a -> b) -> [a] -> [b]
map SomeParam -> Action
someParamToAction Params
params
      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
$ 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 [Action]
pgParams

simpleExecuteStatement
  :: (HasSimplePool m, MonadUnliftIO m)
  => Text
  -> Params
  -> m Int64
simpleExecuteStatement :: forall (m :: * -> *).
(HasSimplePool m, MonadUnliftIO m) =>
Text -> Params -> m Int64
simpleExecuteStatement Text
sqlTemplate Params
params = 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
      pgParams :: [Action]
pgParams = (SomeParam -> Action) -> Params -> [Action]
forall a b. (a -> b) -> [a] -> [b]
map SomeParam -> Action
someParamToAction Params
params
  (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
$ Connection -> Query -> [Action] -> IO Int64
forall q. ToRow q => Connection -> Query -> q -> IO Int64
PG.execute Connection
conn Query
sql [Action]
pgParams

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

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

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

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
p -> ((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
p ((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
sp -> SimpleConnectionPool
sp {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
p -> SimpleConnectionPool
p {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
      a <-
        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
p -> SimpleConnectionPool
p {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 a

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
p) -> ((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
p ((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
sp -> SimpleConnectionPool
sp {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"

simpleRunHandlerWithConnection
  :: (HasSimplePool m, MonadUnliftIO m)
  => (Connection -> jobs -> m result)
  -> jobs
  -> m result
simpleRunHandlerWithConnection :: forall (m :: * -> *) jobs result.
(HasSimplePool m, MonadUnliftIO m) =>
(Connection -> jobs -> m result) -> jobs -> m result
simpleRunHandlerWithConnection Connection -> jobs -> m result
handler jobs
jobs =
  (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 -> jobs -> m result
handler Connection
conn jobs
jobs

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
f = 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
f Connection
conn
    (Maybe Connection
Nothing, Just Pool Connection
p) -> ((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
p ((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
f 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
v) = a -> Action
forall a. ToJSON a => a -> Action
toJSONField a
v
someParamToAction (SomeParam (PScalar Col a
c) a
v) = Col a -> (ToField a => Action) -> Action
forall a r. Col a -> (ToField a => r) -> r
withColToField Col a
c (a -> Action
forall a. ToField a => a -> Action
toField a
v)
someParamToAction (SomeParam (PNullable Col a1
CJsonb) a
v) = 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
v
someParamToAction (SomeParam (PNullable Col a1
c) a
v) = Col a1 -> (ToField a1 => Action) -> Action
forall a r. Col a -> (ToField a => r) -> r
withColToField Col a1
c (a -> Action
forall a. ToField a => a -> Action
toField a
v)
someParamToAction (SomeParam (PArray Col a1
c) a
v) = Col a1 -> (ToField a1 => Action) -> Action
forall a r. Col a -> (ToField a => r) -> r
withColToField Col a1
c (PGArray a1 -> Action
forall a. ToField a => a -> Action
toField ([a1] -> PGArray a1
forall a. [a] -> PGArray a
PGArray a
[a1]
v))
someParamToAction (SomeParam (PNullArray Col a1
c) a
v) = Col a1 -> (ToField a1 => Action) -> Action
forall a r. Col a -> (ToField a => r) -> r
withColToField Col a1
c (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]
v))

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