-- | postgresql-simple primitives for writing a custom
-- 'Arbiter.Core.MonadArbiter.MonadArbiter' instance.
--
-- Handlers receive a @Database.PostgreSQL.Simple.Connection@ for running queries
-- inside the worker transaction:
--
-- @
-- import Arbiter.Core qualified as Arb
-- import Arbiter.Simple qualified as ArbS
-- import Database.PostgreSQL.Simple (Connection)
--
-- instance Arb.MonadArbiter MyApp where
--   type RegistryOf MyApp = MyRegistry
--   type Handler MyApp job result = Connection -> job -> MyApp result
--   getSchema                = asks appSchema
--   executeQuery             = ArbS.simpleExecuteQuery
--   executeStatement         = ArbS.simpleExecuteStatement
--   withDbTransaction        = ArbS.simpleWithDbTransaction
--   runHandlerWithConnection = ArbS.simpleRunHandlerWithConnection
--   getListener              = asks appListener
-- @
--
-- The primitives also need a @HasPoolState Connection MyApp@ instance. See
-- 'Arbiter.Simple.SimpleDb.HasPoolState'.
--
-- Write a handler's own signature as @JobHandler MyApp MyPayload MyResult@, which is
-- 'Arbiter.Core.MonadArbiter.Handler' at that queue's job and declared result types.
module Arbiter.Simple.MonadArbiter
  ( -- * MonadArbiter implementation
    simpleExecuteQuery
  , simpleExecuteStatement
  , simpleWithDbTransaction
  , simpleRunHandlerWithConnection

    -- * Connection
  , simpleWithConnection
  ) where

import Arbiter.Core.Backend (HasPoolState, pinConnection, withConn, withSavepointTransaction)
import Arbiter.Core.MonadArbiter
  ( Col (..)
  , NullCol (..)
  , ParamType (..)
  , SomeParam (..)
  , Stored (..)
  , runCodec
  , storedBytes
  )
import Arbiter.Core.MonadArbiter qualified as MA
import Control.Monad (void)
import Control.Monad.IO.Class (liftIO)
import Data.ByteString (ByteString)
import Data.Int (Int64)
import Data.Text.Encoding qualified as TE
import Database.PostgreSQL.Simple (Connection)
import Database.PostgreSQL.Simple qualified as PG
import Database.PostgreSQL.Simple.FromField (FromField (..), ResultError (..), returnError, typeOid)
import Database.PostgreSQL.Simple.FromRow (RowParser, field)
import Database.PostgreSQL.Simple.ToField (Action (..), ToField (..), toField)
import Database.PostgreSQL.Simple.TypeInfo.Static qualified as TI
import Database.PostgreSQL.Simple.Types (PGArray (..), Query (..))
import UnliftIO (MonadUnliftIO)

-- | Pin one pooled connection for the action.
simpleWithConnection :: (HasPoolState Connection m, MonadUnliftIO m) => m a -> m a
simpleWithConnection :: forall (m :: * -> *) a.
(HasPoolState Connection m, MonadUnliftIO m) =>
m a -> m a
simpleWithConnection = m a -> m a
forall conn (m :: * -> *) a.
(HasPoolState conn m, MonadUnliftIO m) =>
m a -> m a
pinConnection

-- | Run a query, decoding rows.
simpleExecuteQuery
  :: (HasPoolState Connection m, MonadUnliftIO m)
  => MA.Query a
  -> m [a]
simpleExecuteQuery :: forall (m :: * -> *) a.
(HasPoolState Connection m, MonadUnliftIO m) =>
Query a -> m [a]
simpleExecuteQuery Query a
query =
  (Connection -> Query -> IO [a])
-> (Connection -> Query -> [Action] -> IO [a]) -> Query a -> m [a]
forall (m :: * -> *) r a.
(HasPoolState Connection m, MonadUnliftIO m) =>
(Connection -> Query -> IO r)
-> (Connection -> Query -> [Action] -> IO r) -> Query a -> m r
runSplit (RowParser a -> Connection -> Query -> IO [a]
forall r. RowParser r -> Connection -> Query -> IO [r]
PG.queryWith_ RowParser a
parser) (RowParser a -> Connection -> Query -> [Action] -> IO [a]
forall q r.
ToRow q =>
RowParser r -> Connection -> Query -> q -> IO [r]
PG.queryWith RowParser a
parser) Query a
query
  where
    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 (Query a -> RowCodec a
forall a. Query a -> RowCodec a
MA.qDecode Query a
query)

-- | Run a statement, returning rows affected.
simpleExecuteStatement
  :: (HasPoolState Connection m, MonadUnliftIO m)
  => MA.Query a
  -> m Int64
simpleExecuteStatement :: forall (m :: * -> *) a.
(HasPoolState Connection m, MonadUnliftIO m) =>
Query a -> m Int64
simpleExecuteStatement = (Connection -> Query -> IO Int64)
-> (Connection -> Query -> [Action] -> IO Int64)
-> Query a
-> m Int64
forall (m :: * -> *) r a.
(HasPoolState Connection m, MonadUnliftIO m) =>
(Connection -> Query -> IO r)
-> (Connection -> Query -> [Action] -> IO r) -> Query a -> m r
runSplit Connection -> Query -> IO Int64
PG.execute_ Connection -> Query -> [Action] -> IO Int64
forall q. ToRow q => Connection -> Query -> q -> IO Int64
PG.execute

-- | Run a query on a pooled connection. The first runner takes a query with no
-- parameters. The second takes the encoded parameters.
runSplit
  :: (HasPoolState Connection m, MonadUnliftIO m)
  => (Connection -> Query -> IO r)
  -> (Connection -> Query -> [Action] -> IO r)
  -> MA.Query a
  -> m r
runSplit :: forall (m :: * -> *) r a.
(HasPoolState Connection m, MonadUnliftIO m) =>
(Connection -> Query -> IO r)
-> (Connection -> Query -> [Action] -> IO r) -> Query a -> m r
runSplit Connection -> Query -> IO r
bare Connection -> Query -> [Action] -> IO r
withParams Query a
query =
  (Connection -> m r) -> m r
forall conn (m :: * -> *) a.
(HasPoolState conn m, MonadUnliftIO m) =>
(conn -> m a) -> m a
withConn ((Connection -> m r) -> m r) -> (Connection -> m r) -> m r
forall a b. (a -> b) -> a -> b
$ \Connection
conn -> IO r -> m r
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO r -> m r) -> IO r -> m r
forall a b. (a -> b) -> a -> b
$ case Query a -> Params
forall a. Query a -> Params
MA.qParams Query a
query of
    [] -> Connection -> Query -> IO r
bare Connection
conn Query
sql
    Params
params -> Connection -> Query -> [Action] -> IO r
withParams Connection
conn Query
sql ((SomeParam -> Action) -> Params -> [Action]
forall a b. (a -> b) -> [a] -> [b]
map SomeParam -> Action
someParamToAction Params
params)
  where
    sql :: Query
sql = ByteString -> Query
Query (Text -> ByteString
TE.encodeUtf8 (Query a -> Text
forall a. Query a -> Text
MA.qSql Query a
query))

interpretNullCol :: NullCol a -> RowParser a
interpretNullCol :: forall x. NullCol x -> RowParser x
interpretNullCol (NotNull Text
_ Col a
col) = Col a
-> (forall b.
    (FromField b, ToField b) =>
    (b -> a) -> (a -> b) -> RowParser a)
-> RowParser a
forall a r.
Col a
-> (forall b.
    (FromField b, ToField b) =>
    (b -> a) -> (a -> b) -> r)
-> r
withColField Col a
col (\b -> a
from a -> b
_ -> b -> a
from (b -> a) -> RowParser b -> RowParser a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> RowParser b
forall a. FromField a => RowParser a
field)
interpretNullCol (Nullable Text
_ Col a1
col) = Col a1
-> (forall b.
    (FromField b, ToField b) =>
    (b -> a1) -> (a1 -> b) -> RowParser a)
-> RowParser a
forall a r.
Col a
-> (forall b.
    (FromField b, ToField b) =>
    (b -> a) -> (a -> b) -> r)
-> r
withColField Col a1
col (\b -> a1
from a1 -> b
_ -> (b -> a1) -> Maybe b -> Maybe a1
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> a1
from (Maybe b -> a) -> RowParser (Maybe b) -> RowParser a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> RowParser (Maybe b)
forall a. FromField a => RowParser a
field)

-- | Transaction bracket. Nests via savepoints.
simpleWithDbTransaction
  :: (HasPoolState Connection m, MonadUnliftIO m)
  => m a
  -> m a
simpleWithDbTransaction :: forall (m :: * -> *) a.
(HasPoolState Connection m, MonadUnliftIO m) =>
m a -> m a
simpleWithDbTransaction =
  (Connection -> ByteString -> IO ())
-> (Connection -> IO a -> IO a) -> m a -> m a
forall conn (m :: * -> *) a.
(HasPoolState conn m, MonadUnliftIO m) =>
(conn -> ByteString -> IO ())
-> (conn -> IO a -> IO a) -> m a -> m a
withSavepointTransaction (\Connection
conn ByteString
sql -> IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (Connection -> Query -> IO Int64
PG.execute_ Connection
conn (ByteString -> Query
Query ByteString
sql))) Connection -> IO a -> IO a
forall a. Connection -> IO a -> IO a
PG.withTransaction

-- | Run a handler on the pinned connection, or on a borrowed pool connection when none is pinned.
simpleRunHandlerWithConnection
  :: (HasPoolState Connection m, MonadUnliftIO m)
  => (Connection -> job -> m result)
  -> job
  -> m result
simpleRunHandlerWithConnection :: forall (m :: * -> *) job result.
(HasPoolState Connection m, MonadUnliftIO m) =>
(Connection -> job -> m result) -> job -> m result
simpleRunHandlerWithConnection Connection -> job -> m result
handler job
job =
  (Connection -> m result) -> m result
forall conn (m :: * -> *) a.
(HasPoolState conn m, MonadUnliftIO m) =>
(conn -> 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

someParamToAction :: SomeParam -> Action
someParamToAction :: SomeParam -> Action
someParamToAction (SomeParam (PScalar Col a
col) a
value) = Col a
-> (forall b.
    (FromField b, ToField b) =>
    (b -> a) -> (a -> b) -> Action)
-> Action
forall a r.
Col a
-> (forall b.
    (FromField b, ToField b) =>
    (b -> a) -> (a -> b) -> r)
-> r
withColField Col a
col (\b -> a
_ a -> b
to -> b -> Action
forall a. ToField a => a -> Action
toField (a -> b
to a
value))
someParamToAction (SomeParam (PNullable Col a1
col) a
value) = Col a1
-> (forall b.
    (FromField b, ToField b) =>
    (b -> a1) -> (a1 -> b) -> Action)
-> Action
forall a r.
Col a
-> (forall b.
    (FromField b, ToField b) =>
    (b -> a) -> (a -> b) -> r)
-> r
withColField Col a1
col (\b -> a1
_ a1 -> b
to -> Maybe b -> Action
forall a. ToField a => a -> Action
toField (a1 -> b
to (a1 -> b) -> Maybe a1 -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a
Maybe a1
value))
someParamToAction (SomeParam (PArray Col a1
col) a
value) = Col a1
-> (forall b.
    (FromField b, ToField b) =>
    (b -> a1) -> (a1 -> b) -> Action)
-> Action
forall a r.
Col a
-> (forall b.
    (FromField b, ToField b) =>
    (b -> a) -> (a -> b) -> r)
-> r
withColField Col a1
col (\b -> a1
_ a1 -> b
to -> PGArray b -> Action
forall a. ToField a => a -> Action
toField ([b] -> PGArray b
forall a. [a] -> PGArray a
PGArray ((a1 -> b) -> [a1] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map a1 -> b
to a
[a1]
value)))
someParamToAction (SomeParam (PNullArray Col a1
col) a
value) = Col a1
-> (forall b.
    (FromField b, ToField b) =>
    (b -> a1) -> (a1 -> b) -> Action)
-> Action
forall a r.
Col a
-> (forall b.
    (FromField b, ToField b) =>
    (b -> a) -> (a -> b) -> r)
-> r
withColField Col a1
col (\b -> a1
_ a1 -> b
to -> PGArray (Maybe b) -> Action
forall a. ToField a => a -> Action
toField ([Maybe b] -> PGArray (Maybe b)
forall a. [a] -> PGArray a
PGArray ((Maybe a1 -> Maybe b) -> [Maybe a1] -> [Maybe b]
forall a b. (a -> b) -> [a] -> [b]
map ((a1 -> b) -> Maybe a1 -> Maybe b
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a1 -> b
to) a
[Maybe a1]
value)))

-- | The field instances for a column, with raw JSON carried as 'RawJson'.
withColField :: Col a -> (forall b. (FromField b, ToField b) => (b -> a) -> (a -> b) -> r) -> r
withColField :: forall a r.
Col a
-> (forall b.
    (FromField b, ToField b) =>
    (b -> a) -> (a -> b) -> r)
-> r
withColField Col a
CInt4 forall b. (FromField b, ToField b) => (b -> a) -> (a -> b) -> r
k = (a -> a) -> (a -> a) -> r
forall b. (FromField b, ToField b) => (b -> a) -> (a -> b) -> r
k a -> a
forall a. a -> a
id a -> a
forall a. a -> a
id
withColField Col a
CInt8 forall b. (FromField b, ToField b) => (b -> a) -> (a -> b) -> r
k = (a -> a) -> (a -> a) -> r
forall b. (FromField b, ToField b) => (b -> a) -> (a -> b) -> r
k a -> a
forall a. a -> a
id a -> a
forall a. a -> a
id
withColField Col a
CText forall b. (FromField b, ToField b) => (b -> a) -> (a -> b) -> r
k = (a -> a) -> (a -> a) -> r
forall b. (FromField b, ToField b) => (b -> a) -> (a -> b) -> r
k a -> a
forall a. a -> a
id a -> a
forall a. a -> a
id
withColField Col a
CBool forall b. (FromField b, ToField b) => (b -> a) -> (a -> b) -> r
k = (a -> a) -> (a -> a) -> r
forall b. (FromField b, ToField b) => (b -> a) -> (a -> b) -> r
k a -> a
forall a. a -> a
id a -> a
forall a. a -> a
id
withColField Col a
CTimestamptz forall b. (FromField b, ToField b) => (b -> a) -> (a -> b) -> r
k = (a -> a) -> (a -> a) -> r
forall b. (FromField b, ToField b) => (b -> a) -> (a -> b) -> r
k a -> a
forall a. a -> a
id a -> a
forall a. a -> a
id
withColField Col a
CJsonb forall b. (FromField b, ToField b) => (b -> a) -> (a -> b) -> r
k = (a -> a) -> (a -> a) -> r
forall b. (FromField b, ToField b) => (b -> a) -> (a -> b) -> r
k a -> a
forall a. a -> a
id a -> a
forall a. a -> a
id
withColField Col a
CStored forall b. (FromField b, ToField b) => (b -> a) -> (a -> b) -> r
k = (RawJson -> a) -> (a -> RawJson) -> r
forall b. (FromField b, ToField b) => (b -> a) -> (a -> b) -> r
k (\(RawJson ByteString
bytes) -> ByteString -> Stored payload
forall payload. ByteString -> Stored payload
Stored ByteString
bytes) (ByteString -> RawJson
RawJson (ByteString -> RawJson) -> (a -> ByteString) -> a -> RawJson
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> ByteString
Stored payload -> ByteString
forall payload. Stored payload -> ByteString
storedBytes)
withColField Col a
CFloat8 forall b. (FromField b, ToField b) => (b -> a) -> (a -> b) -> r
k = (a -> a) -> (a -> a) -> r
forall b. (FromField b, ToField b) => (b -> a) -> (a -> b) -> r
k a -> a
forall a. a -> a
id a -> a
forall a. a -> a
id
withColField Col a
CUuid forall b. (FromField b, ToField b) => (b -> a) -> (a -> b) -> r
k = (a -> a) -> (a -> a) -> r
forall b. (FromField b, ToField b) => (b -> a) -> (a -> b) -> r
k a -> a
forall a. a -> a
id a -> a
forall a. a -> a
id

-- | JSON bytes passed through a @json@ or @jsonb@ column without parsing.
newtype RawJson = RawJson ByteString

instance FromField RawJson where
  fromField :: FieldParser RawJson
fromField Field
f Maybe ByteString
mdata
    | Field -> Oid
typeOid Field
f Oid -> Oid -> Bool
forall a. Eq a => a -> a -> Bool
/= Oid
TI.jsonbOid Bool -> Bool -> Bool
&& Field -> Oid
typeOid Field
f Oid -> Oid -> Bool
forall a. Eq a => a -> a -> Bool
/= Oid
TI.jsonOid = (String -> Maybe Oid -> String -> String -> String -> ResultError)
-> Field -> String -> Conversion RawJson
forall a err.
(Typeable a, Exception err) =>
(String -> Maybe Oid -> String -> String -> String -> err)
-> Field -> String -> Conversion a
returnError String -> Maybe Oid -> String -> String -> String -> ResultError
Incompatible Field
f String
""
    | Bool
otherwise = Conversion RawJson
-> (ByteString -> Conversion RawJson)
-> Maybe ByteString
-> Conversion RawJson
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ((String -> Maybe Oid -> String -> String -> String -> ResultError)
-> Field -> String -> Conversion RawJson
forall a err.
(Typeable a, Exception err) =>
(String -> Maybe Oid -> String -> String -> String -> err)
-> Field -> String -> Conversion a
returnError String -> Maybe Oid -> String -> String -> String -> ResultError
UnexpectedNull Field
f String
"") (RawJson -> Conversion RawJson
forall a. a -> Conversion a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RawJson -> Conversion RawJson)
-> (ByteString -> RawJson) -> ByteString -> Conversion RawJson
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> RawJson
RawJson) Maybe ByteString
mdata

instance ToField RawJson where
  toField :: RawJson -> Action
toField (RawJson ByteString
bytes) = ByteString -> Action
Escape ByteString
bytes