{-# LANGUAGE OverloadedStrings #-}

-- | Orville primitives for writing a custom
-- 'Arbiter.Core.MonadArbiter.MonadArbiter' instance.
--
-- Handlers run in the application monad on Orville's own connection:
--
-- @
-- import Arbiter.Core qualified as Arb
-- import Arbiter.Orville qualified as ArbO
--
-- instance Arb.MonadArbiter MyApp where
--   type RegistryOf MyApp = MyRegistry
--   type Handler MyApp job result = job -> MyApp result
--   getSchema                = asks appSchema
--   executeQuery             = ArbO.orvilleExecuteQuery
--   executeStatement         = ArbO.orvilleExecuteStatement
--   withDbTransaction        = ArbO.orvilleWithDbTransaction
--   runHandlerWithConnection = ArbO.orvilleRunHandlerWithConnection
--   getListener              = asks appListener
-- @
--
-- The primitives also need a @MonadOrville MyApp@ instance.
--
-- 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.Orville.MonadArbiter
  ( -- * MonadArbiter implementation
    orvilleExecuteQuery
  , orvilleExecuteStatement
  , orvilleWithDbTransaction
  , orvilleRunHandlerWithConnection
  ) where

import Arbiter.Core.Enum (showText)
import Arbiter.Core.Exceptions (orThrowInternal)
import Arbiter.Core.MonadArbiter
  ( Col (..)
  , NullCol (..)
  , ParamType (..)
  , Query (..)
  , SomeParam (..)
  , Stored (..)
  , runCodec
  , storedBytes
  )
import Control.Monad ((<=<))
import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.Aeson (Value, eitherDecodeStrict', encode)
import Data.ByteString (ByteString)
import Data.Int (Int64)
import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.Encoding qualified as TE
import Data.Text.Lazy qualified as TL
import Data.Text.Lazy.Encoding qualified as TLE
import Database.PostgreSQL.LibPQ qualified as LibPQ
import Orville.PostgreSQL qualified as O
import Orville.PostgreSQL.Marshall.FieldDefinition qualified as FieldDef
import Orville.PostgreSQL.Marshall.SqlMarshaller qualified as O
import Orville.PostgreSQL.Raw.Connection qualified as Conn
import Orville.PostgreSQL.Raw.PgTextFormatValue (PgTextFormatValue)
import Orville.PostgreSQL.Raw.PgTextFormatValue qualified as PgText
import Orville.PostgreSQL.Raw.SqlValue (SqlValue)
import Orville.PostgreSQL.Raw.SqlValue qualified as SqlValue

import Arbiter.Orville.Array qualified as Array

-- | Run a query, decoding rows.
orvilleExecuteQuery
  :: (O.MonadOrville m)
  => Query a
  -> m [a]
orvilleExecuteQuery :: forall (m :: * -> *) a. MonadOrville m => Query a -> m [a]
orvilleExecuteQuery Query a
query = (Connection -> m [a]) -> m [a]
forall (m :: * -> *) a.
MonadOrville m =>
(Connection -> m a) -> m a
O.withConnection ((Connection -> m [a]) -> m [a]) -> (Connection -> m [a]) -> m [a]
forall a b. (a -> b) -> a -> b
$ \Connection
conn -> do
  result <- Connection -> Query a -> m Result
forall (m :: * -> *) a.
MonadIO m =>
Connection -> Query a -> m Result
runRaw Connection
conn Query a
query
  let marshaller = SqlMarshaller () a -> AnnotatedSqlMarshaller () a
forall writeEntity readEntity.
SqlMarshaller writeEntity readEntity
-> AnnotatedSqlMarshaller writeEntity readEntity
O.annotateSqlMarshallerEmptyAnnotation ((forall x. NullCol x -> SqlMarshaller () x)
-> RowCodec a -> SqlMarshaller () a
forall (f :: * -> *) a.
Applicative f =>
(forall x. NullCol x -> f x) -> RowCodec a -> f a
runCodec NullCol x -> SqlMarshaller () x
forall x. NullCol x -> SqlMarshaller () x
orvilleCol (Query a -> RowCodec a
forall a. Query a -> RowCodec a
qDecode Query a
query))
  decoded <- liftIO $ O.marshallResultFromSql O.defaultErrorDetailLevel marshaller result
  orThrowInternal (("orville decode error: " <>) . showText) decoded

-- | Run a statement, returning rows affected.
orvilleExecuteStatement
  :: (O.MonadOrville m)
  => Query a
  -> m Int64
orvilleExecuteStatement :: forall (m :: * -> *) a. MonadOrville m => Query a -> m Int64
orvilleExecuteStatement Query a
query = (Connection -> m Int64) -> m Int64
forall (m :: * -> *) a.
MonadOrville m =>
(Connection -> m a) -> m a
O.withConnection ((Connection -> m Int64) -> m Int64)
-> (Connection -> m Int64) -> m Int64
forall a b. (a -> b) -> a -> b
$ \Connection
conn -> Connection -> Query a -> m Result
forall (m :: * -> *) a.
MonadIO m =>
Connection -> Query a -> m Result
runRaw Connection
conn Query a
query m Result -> (Result -> m Int64) -> m Int64
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IO Int64 -> m Int64
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int64 -> m Int64) -> (Result -> IO Int64) -> Result -> m Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Result -> IO Int64
readRowCount

runRaw :: (MonadIO m) => Conn.Connection -> Query a -> m LibPQ.Result
runRaw :: forall (m :: * -> *) a.
MonadIO m =>
Connection -> Query a -> m Result
runRaw Connection
conn Query a
query = do
  pgParams <- [SomeParam] -> m [Maybe PgTextFormatValue]
forall (m :: * -> *).
MonadIO m =>
[SomeParam] -> m [Maybe PgTextFormatValue]
encodeParams (Query a -> [SomeParam]
forall a. Query a -> [SomeParam]
qParams Query a
query)
  liftIO $ Conn.executeRaw conn (TE.encodeUtf8 (qPositional query)) pgParams

-- | Transaction bracket. Nests via savepoints.
orvilleWithDbTransaction :: (O.MonadOrville m) => m a -> m a
orvilleWithDbTransaction :: forall (m :: * -> *) a. MonadOrville m => m a -> m a
orvilleWithDbTransaction = m a -> m a
forall (m :: * -> *) a. MonadOrville m => m a -> m a
O.withTransaction

-- | Run a handler. Orville manages its own connection.
orvilleRunHandlerWithConnection :: (job -> m result) -> job -> m result
orvilleRunHandlerWithConnection :: forall {k} job (m :: k -> *) (result :: k).
(job -> m result) -> job -> m result
orvilleRunHandlerWithConnection = (job -> m result) -> job -> m result
forall a. a -> a
id

someParamToSqlValue :: SomeParam -> Either Text SqlValue
someParamToSqlValue :: SomeParam -> Either Text SqlValue
someParamToSqlValue (SomeParam (PScalar Col a
col) a
value) =
  SqlValue -> Either Text SqlValue
forall a b. b -> Either a b
Right (SqlValue -> Either Text SqlValue)
-> SqlValue -> Either Text SqlValue
forall a b. (a -> b) -> a -> b
$ FieldDefinition NotNull a -> a -> SqlValue
forall nullability a.
FieldDefinition nullability a -> a -> SqlValue
FieldDef.fieldValueToSqlValue (Text -> Col a -> FieldDefinition NotNull a
forall a. Text -> Col a -> FieldDefinition NotNull a
colFieldDef Text
"" Col a
col) a
value
someParamToSqlValue (SomeParam (PNullable Col a1
col) a
value) =
  SqlValue -> Either Text SqlValue
forall a b. b -> Either a b
Right (SqlValue -> Either Text SqlValue)
-> SqlValue -> Either Text SqlValue
forall a b. (a -> b) -> a -> b
$ FieldDefinition Nullable (Maybe a1) -> Maybe a1 -> SqlValue
forall nullability a.
FieldDefinition nullability a -> a -> SqlValue
FieldDef.fieldValueToSqlValue (FieldDefinition NotNull a1 -> FieldDefinition Nullable (Maybe a1)
forall a.
FieldDefinition NotNull a -> FieldDefinition Nullable (Maybe a)
O.nullableField (Text -> Col a1 -> FieldDefinition NotNull a1
forall a. Text -> Col a -> FieldDefinition NotNull a
colFieldDef Text
"" Col a1
col)) a
Maybe a1
value
someParamToSqlValue (SomeParam (PArray Col a1
col) a
values) =
  ByteString -> SqlValue
SqlValue.fromRawBytes (ByteString -> SqlValue)
-> ([ByteString] -> ByteString) -> [ByteString] -> SqlValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ByteString] -> ByteString
Array.fmtArray ([ByteString] -> SqlValue)
-> Either Text [ByteString] -> Either Text SqlValue
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a1 -> Either Text ByteString) -> [a1] -> Either Text [ByteString]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (Col a1 -> a1 -> Either Text ByteString
forall a. Col a -> a -> Either Text ByteString
colToBytes Col a1
col) a
[a1]
values
someParamToSqlValue (SomeParam (PNullArray Col a1
col) a
values) =
  ByteString -> SqlValue
SqlValue.fromRawBytes (ByteString -> SqlValue)
-> ([Maybe ByteString] -> ByteString)
-> [Maybe ByteString]
-> SqlValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe ByteString] -> ByteString
Array.fmtNullableArray ([Maybe ByteString] -> SqlValue)
-> Either Text [Maybe ByteString] -> Either Text SqlValue
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Maybe a1 -> Either Text (Maybe ByteString))
-> [Maybe a1] -> Either Text [Maybe ByteString]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse ((a1 -> Either Text ByteString)
-> Maybe a1 -> Either Text (Maybe ByteString)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse (Col a1 -> a1 -> Either Text ByteString
forall a. Col a -> a -> Either Text ByteString
colToBytes Col a1
col)) a
[Maybe a1]
values

colToBytes :: Col a -> a -> Either Text ByteString
colToBytes :: forall a. Col a -> a -> Either Text ByteString
colToBytes Col a
col a
value = SqlValue -> Either Text ByteString
sqlValueToBytes (SqlValue -> Either Text ByteString)
-> SqlValue -> Either Text ByteString
forall a b. (a -> b) -> a -> b
$ FieldDefinition NotNull a -> a -> SqlValue
forall nullability a.
FieldDefinition nullability a -> a -> SqlValue
FieldDef.fieldValueToSqlValue (Text -> Col a -> FieldDefinition NotNull a
forall a. Text -> Col a -> FieldDefinition NotNull a
colFieldDef Text
"" Col a
col) a
value

sqlValueToBytes :: SqlValue -> Either Text ByteString
sqlValueToBytes :: SqlValue -> Either Text ByteString
sqlValueToBytes SqlValue
value =
  SqlValue -> Either Text (Maybe PgTextFormatValue)
sqlValueToParam SqlValue
value
    Either Text (Maybe PgTextFormatValue)
-> (Maybe PgTextFormatValue -> Either Text ByteString)
-> Either Text ByteString
forall a b. Either Text a -> (a -> Either Text b) -> Either Text b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Either Text ByteString
-> (PgTextFormatValue -> Either Text ByteString)
-> Maybe PgTextFormatValue
-> Either Text ByteString
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Text -> Either Text ByteString
forall a b. a -> Either a b
Left Text
"sqlValueToBytes: got NULL, expected non-null scalar") (ByteString -> Either Text ByteString
forall a b. b -> Either a b
Right (ByteString -> Either Text ByteString)
-> (PgTextFormatValue -> ByteString)
-> PgTextFormatValue
-> Either Text ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PgTextFormatValue -> ByteString
PgText.toByteString)

sqlValueToParam :: SqlValue -> Either Text (Maybe PgTextFormatValue)
sqlValueToParam :: SqlValue -> Either Text (Maybe PgTextFormatValue)
sqlValueToParam =
  (PgTextFormatValue -> Either Text (Maybe PgTextFormatValue))
-> (NonEmpty (Either Text (Maybe PgTextFormatValue))
    -> Either Text (Maybe PgTextFormatValue))
-> Either Text (Maybe PgTextFormatValue)
-> SqlValue
-> Either Text (Maybe PgTextFormatValue)
forall r.
(PgTextFormatValue -> r) -> (NonEmpty r -> r) -> r -> SqlValue -> r
SqlValue.foldSqlValue
    (Maybe PgTextFormatValue -> Either Text (Maybe PgTextFormatValue)
forall a b. b -> Either a b
Right (Maybe PgTextFormatValue -> Either Text (Maybe PgTextFormatValue))
-> (PgTextFormatValue -> Maybe PgTextFormatValue)
-> PgTextFormatValue
-> Either Text (Maybe PgTextFormatValue)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PgTextFormatValue -> Maybe PgTextFormatValue
forall a. a -> Maybe a
Just)
    (Either Text (Maybe PgTextFormatValue)
-> NonEmpty (Either Text (Maybe PgTextFormatValue))
-> Either Text (Maybe PgTextFormatValue)
forall a b. a -> b -> a
const (Either Text (Maybe PgTextFormatValue)
 -> NonEmpty (Either Text (Maybe PgTextFormatValue))
 -> Either Text (Maybe PgTextFormatValue))
-> Either Text (Maybe PgTextFormatValue)
-> NonEmpty (Either Text (Maybe PgTextFormatValue))
-> Either Text (Maybe PgTextFormatValue)
forall a b. (a -> b) -> a -> b
$ Text -> Either Text (Maybe PgTextFormatValue)
forall a b. a -> Either a b
Left Text
"sqlValueToParam: got composite row, expected scalar")
    (Maybe PgTextFormatValue -> Either Text (Maybe PgTextFormatValue)
forall a b. b -> Either a b
Right Maybe PgTextFormatValue
forall a. Maybe a
Nothing)

-- | Encode the ordered params into the libpq values 'Conn.executeRaw' consumes.
encodeParams :: (MonadIO m) => [SomeParam] -> m [Maybe PgTextFormatValue]
encodeParams :: forall (m :: * -> *).
MonadIO m =>
[SomeParam] -> m [Maybe PgTextFormatValue]
encodeParams = (SomeParam -> m (Maybe PgTextFormatValue))
-> [SomeParam] -> m [Maybe PgTextFormatValue]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse ((Text -> Text)
-> Either Text (Maybe PgTextFormatValue)
-> m (Maybe PgTextFormatValue)
forall (m :: * -> *) e a.
MonadIO m =>
(e -> Text) -> Either e a -> m a
orThrowInternal (Text
"param encoding error: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>) (Either Text (Maybe PgTextFormatValue)
 -> m (Maybe PgTextFormatValue))
-> (SomeParam -> Either Text (Maybe PgTextFormatValue))
-> SomeParam
-> m (Maybe PgTextFormatValue)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SqlValue -> Either Text (Maybe PgTextFormatValue)
sqlValueToParam (SqlValue -> Either Text (Maybe PgTextFormatValue))
-> (SomeParam -> Either Text SqlValue)
-> SomeParam
-> Either Text (Maybe PgTextFormatValue)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< SomeParam -> Either Text SqlValue
someParamToSqlValue))

orvilleCol :: NullCol a -> O.SqlMarshaller () a
orvilleCol :: forall x. NullCol x -> SqlMarshaller () x
orvilleCol (NotNull Text
name Col a
col) = FieldDefinition NotNull a -> SqlMarshaller () a
forall nullability fieldValue writeEntity.
FieldDefinition nullability fieldValue
-> SqlMarshaller writeEntity fieldValue
O.marshallReadOnlyField (Text -> Col a -> FieldDefinition NotNull a
forall a. Text -> Col a -> FieldDefinition NotNull a
colFieldDef Text
name Col a
col)
orvilleCol (Nullable Text
name Col a1
col) = FieldDefinition Nullable a -> SqlMarshaller () a
forall nullability fieldValue writeEntity.
FieldDefinition nullability fieldValue
-> SqlMarshaller writeEntity fieldValue
O.marshallReadOnlyField (FieldDefinition NotNull a1 -> FieldDefinition Nullable (Maybe a1)
forall a.
FieldDefinition NotNull a -> FieldDefinition Nullable (Maybe a)
O.nullableField (Text -> Col a1 -> FieldDefinition NotNull a1
forall a. Text -> Col a -> FieldDefinition NotNull a
colFieldDef Text
name Col a1
col))

colFieldDef :: Text -> Col a -> O.FieldDefinition O.NotNull a
colFieldDef :: forall a. Text -> Col a -> FieldDefinition NotNull a
colFieldDef Text
name Col a
col = Col a -> String -> FieldDefinition NotNull a
forall a. Col a -> String -> FieldDefinition NotNull a
colField Col a
col (Text -> String
T.unpack Text
name)

colField :: Col a -> String -> O.FieldDefinition O.NotNull a
colField :: forall a. Col a -> String -> FieldDefinition NotNull a
colField Col a
CInt4 = String -> FieldDefinition NotNull a
String -> FieldDefinition NotNull Int32
O.integerField
colField Col a
CInt8 = String -> FieldDefinition NotNull a
String -> FieldDefinition NotNull Int64
O.bigIntegerField
colField Col a
CText = String -> FieldDefinition NotNull a
String -> FieldDefinition NotNull Text
O.unboundedTextField
colField Col a
CBool = String -> FieldDefinition NotNull a
String -> FieldDefinition NotNull Bool
O.booleanField
colField Col a
CTimestamptz = String -> FieldDefinition NotNull a
String -> FieldDefinition NotNull UTCTime
O.utcTimestampField
colField Col a
CJsonb = SqlType a -> String -> FieldDefinition NotNull a
forall a. SqlType a -> String -> FieldDefinition NotNull a
O.fieldOfType SqlType a
SqlType Value
jsonbValue
colField Col a
CStored = SqlType a -> String -> FieldDefinition NotNull a
forall a. SqlType a -> String -> FieldDefinition NotNull a
O.fieldOfType SqlType a
SqlType (Stored payload)
forall payload. SqlType (Stored payload)
storedJsonb
colField Col a
CFloat8 = String -> FieldDefinition NotNull a
String -> FieldDefinition NotNull Double
O.doubleField
colField Col a
CUuid = String -> FieldDefinition NotNull a
String -> FieldDefinition NotNull UUID
O.uuidField

jsonbValue :: O.SqlType Value
jsonbValue :: SqlType Value
jsonbValue =
  (Value -> Text)
-> (Text -> Either String Value) -> SqlType Text -> SqlType Value
forall b a.
(b -> a) -> (a -> Either String b) -> SqlType a -> SqlType b
O.tryConvertSqlType
    (LazyText -> Text
TL.toStrict (LazyText -> Text) -> (Value -> LazyText) -> Value -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> LazyText
TLE.decodeUtf8 (ByteString -> LazyText)
-> (Value -> ByteString) -> Value -> LazyText
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> ByteString
forall a. ToJSON a => a -> ByteString
encode)
    (ByteString -> Either String Value
forall a. FromJSON a => ByteString -> Either String a
eitherDecodeStrict' (ByteString -> Either String Value)
-> (Text -> ByteString) -> Text -> Either String Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ByteString
TE.encodeUtf8)
    SqlType Text
O.jsonb

-- | The JSON bytes as stored, through orville's text-typed @jsonb@.
storedJsonb :: O.SqlType (Stored payload)
storedJsonb :: forall payload. SqlType (Stored payload)
storedJsonb = (Stored payload -> Text)
-> (Text -> Stored payload)
-> SqlType Text
-> SqlType (Stored payload)
forall b a. (b -> a) -> (a -> b) -> SqlType a -> SqlType b
O.convertSqlType (ByteString -> Text
TE.decodeUtf8 (ByteString -> Text)
-> (Stored payload -> ByteString) -> Stored payload -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stored payload -> ByteString
forall payload. Stored payload -> ByteString
storedBytes) (ByteString -> Stored payload
forall payload. ByteString -> Stored payload
Stored (ByteString -> Stored payload)
-> (Text -> ByteString) -> Text -> Stored payload
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ByteString
TE.encodeUtf8) SqlType Text
O.jsonb

readRowCount :: LibPQ.Result -> IO Int64
readRowCount :: Result -> IO Int64
readRowCount Result
res =
  Int64 -> (ByteString -> Int64) -> Maybe ByteString -> Int64
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int64
0 ((String -> Int64) -> (Int -> Int64) -> Either String Int -> Int64
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Int64 -> String -> Int64
forall a b. a -> b -> a
const Int64
0) Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Either String Int -> Int64)
-> (ByteString -> Either String Int) -> ByteString -> Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SqlValue -> Either String Int
SqlValue.toInt (SqlValue -> Either String Int)
-> (ByteString -> SqlValue) -> ByteString -> Either String Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> SqlValue
SqlValue.fromRawBytes) (Maybe ByteString -> Int64) -> IO (Maybe ByteString) -> IO Int64
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Result -> IO (Maybe ByteString)
LibPQ.cmdTuples Result
res