{-# LANGUAGE OverloadedStrings #-}

-- | hasql primitives for writing a custom 'Arbiter.Core.MonadArbiter.MonadArbiter' instance.
--
-- Handlers receive a @Hasql.Connection.Connection@ for running typed hasql
-- queries inside the worker transaction:
--
-- @
-- import Arbiter.Core qualified as Arb
-- import Arbiter.Hasql qualified as ArbH
-- import Hasql.Connection qualified as Hasql
--
-- instance Arb.MonadArbiter MyApp where
--   type RegistryOf MyApp = MyRegistry
--   type Handler MyApp job result = Hasql.Connection -> job -> MyApp result
--   getSchema                = asks appSchema
--   executeQuery             = ArbH.hasqlExecuteQuery
--   executeQueryPrepared     = ArbH.hasqlExecuteQueryPrepared True
--   executeStatement         = ArbH.hasqlExecuteStatement
--   withDbTransaction        = ArbH.hasqlWithDbTransaction
--   runHandlerWithConnection = ArbH.hasqlRunHandlerWithConnection
--   getListener              = asks appListener
-- @
--
-- The primitives also need a @HasPoolState Hasql.Connection MyApp@ instance. See
-- 'Arbiter.Hasql.HasqlDb.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.Hasql.MonadArbiter
  ( -- * MonadArbiter implementation
    hasqlExecuteQuery
  , hasqlExecuteQueryPrepared
  , hasqlExecuteStatement
  , hasqlWithDbTransaction
  , hasqlRunHandlerWithConnection

    -- * Connection
  , hasqlWithConnection
  ) where

import Arbiter.Core.Backend (HasPoolState, pinConnection, withConn, withSavepointTransaction)
import Arbiter.Core.Enum (showText)
import Arbiter.Core.Exceptions (orThrowInternal)
import Arbiter.Core.MonadArbiter (Query (..))
import Control.Monad (void, when)
import Control.Monad.IO.Class (liftIO)
import Data.Int (Int64)
import Data.Text (Text)
import Hasql.Connection qualified as Hasql
import Hasql.Decoders qualified as D
import Hasql.Session qualified as Session
import Hasql.Statement qualified as S
import UnliftIO (MonadUnliftIO, mask, onException, tryAny)

import Arbiter.Hasql.Compat (connectionInTransaction, runSQL)
import Arbiter.Hasql.Decode qualified as Decode
import Arbiter.Hasql.Encode qualified as Encode

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

-- | Run a query unprepared, decoding rows.
hasqlExecuteQuery
  :: (HasPoolState Hasql.Connection m, MonadUnliftIO m)
  => Query a
  -> m [a]
hasqlExecuteQuery :: forall (m :: * -> *) a.
(HasPoolState Connection m, MonadUnliftIO m) =>
Query a -> m [a]
hasqlExecuteQuery = Bool -> Query a -> m [a]
forall (m :: * -> *) a.
(HasPoolState Connection m, MonadUnliftIO m) =>
Bool -> Query a -> m [a]
hasqlExecuteQueryPrepared Bool
False

-- | 'hasqlExecuteQuery', prepared once per connection when the flag is on.
hasqlExecuteQueryPrepared
  :: (HasPoolState Hasql.Connection m, MonadUnliftIO m)
  => Bool
  -> Query a
  -> m [a]
hasqlExecuteQueryPrepared :: forall (m :: * -> *) a.
(HasPoolState Connection m, MonadUnliftIO m) =>
Bool -> Query a -> m [a]
hasqlExecuteQueryPrepared Bool
prepare Query a
query = (Connection -> m [a]) -> m [a]
forall conn (m :: * -> *) a.
(HasPoolState conn m, MonadUnliftIO m) =>
(conn -> 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])
-> (Statement () [a] -> IO [a]) -> Statement () [a] -> m [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Connection -> Statement () [a] -> IO [a]
forall r. Text -> Connection -> Statement () r -> IO r
runStatement Text
"query" Connection
conn (Statement () [a] -> m [a]) -> Statement () [a] -> m [a]
forall a b. (a -> b) -> a -> b
$
    Text -> Params () -> Result [a] -> Statement () [a]
mkStatement (Query a -> Text
forall a. Query a -> Text
qPositional Query a
query) (Params -> Params ()
Encode.buildEncoder (Query a -> Params
forall a. Query a -> Params
qParams Query a
query)) (RowCodec a -> Result [a]
forall a. RowCodec a -> Result [a]
Decode.hasqlRowDecoder (Query a -> RowCodec a
forall a. Query a -> RowCodec a
qDecode Query a
query))
  where
    mkStatement :: Text -> Params () -> Result [a] -> Statement () [a]
mkStatement = if Bool
prepare then Text -> Params () -> Result [a] -> Statement () [a]
forall params result.
Text -> Params params -> Result result -> Statement params result
S.preparable else Text -> Params () -> Result [a] -> Statement () [a]
forall params result.
Text -> Params params -> Result result -> Statement params result
S.unpreparable

-- | Run a statement unprepared, returning rows affected.
hasqlExecuteStatement
  :: (HasPoolState Hasql.Connection m, MonadUnliftIO m)
  => Query a
  -> m Int64
hasqlExecuteStatement :: forall (m :: * -> *) a.
(HasPoolState Connection m, MonadUnliftIO m) =>
Query a -> m Int64
hasqlExecuteStatement Query a
query = (Connection -> m Int64) -> m Int64
forall conn (m :: * -> *) a.
(HasPoolState conn m, MonadUnliftIO m) =>
(conn -> 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)
-> (Statement () Int64 -> IO Int64)
-> Statement () Int64
-> m Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Connection -> Statement () Int64 -> IO Int64
forall r. Text -> Connection -> Statement () r -> IO r
runStatement Text
"statement" Connection
conn (Statement () Int64 -> m Int64) -> Statement () Int64 -> m Int64
forall a b. (a -> b) -> a -> b
$
    Text -> Params () -> Result Int64 -> Statement () Int64
forall params result.
Text -> Params params -> Result result -> Statement params result
S.unpreparable (Query a -> Text
forall a. Query a -> Text
qPositional Query a
query) (Params -> Params ()
Encode.buildEncoder (Query a -> Params
forall a. Query a -> Params
qParams Query a
query)) Result Int64
D.rowsAffected

runStatement :: Text -> Hasql.Connection -> S.Statement () r -> IO r
runStatement :: forall r. Text -> Connection -> Statement () r -> IO r
runStatement Text
label Connection
conn Statement () r
stmt =
  Connection -> Session r -> IO (Either SessionError r)
forall a. Connection -> Session a -> IO (Either SessionError a)
Hasql.use Connection
conn (() -> Statement () r -> Session r
forall params result.
params -> Statement params result -> Session result
Session.statement () Statement () r
stmt)
    IO (Either SessionError r)
-> (Either SessionError r -> IO r) -> IO r
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (SessionError -> Text) -> Either SessionError r -> IO r
forall (m :: * -> *) e a.
MonadIO m =>
(e -> Text) -> Either e a -> m a
orThrowInternal (\SessionError
err -> Text
"hasql " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
label Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" error: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> SessionError -> Text
forall a. Show a => a -> Text
showText SessionError
err)

-- | Transaction bracket. Nests via savepoints.
hasqlWithDbTransaction :: (HasPoolState Hasql.Connection m, MonadUnliftIO m) => m a -> m a
hasqlWithDbTransaction :: forall (m :: * -> *) a.
(HasPoolState Connection m, MonadUnliftIO m) =>
m a -> m a
hasqlWithDbTransaction = (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 -> ByteString -> IO ()
runSQL Connection -> IO a -> IO a
forall a. Connection -> IO a -> IO a
beginCommitOrRollback

beginCommitOrRollback :: Hasql.Connection -> IO a -> IO a
beginCommitOrRollback :: forall a. Connection -> IO a -> IO a
beginCommitOrRollback Connection
conn IO a
action = ((forall a. IO a -> IO a) -> IO a) -> IO a
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> m a) -> m b) -> m b
mask (((forall a. IO a -> IO a) -> IO a) -> IO a)
-> ((forall a. IO a -> IO a) -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
restore -> do
  Connection -> ByteString -> IO ()
runSQL Connection
conn ByteString
"BEGIN"
  result <- IO a -> IO a
forall a. IO a -> IO a
restore IO a
action IO a -> IO () -> IO a
forall (m :: * -> *) a b. MonadUnliftIO m => m a -> m b -> m a
`onException` IO ()
rollbackSafely
  runSQL conn "COMMIT"
  pure result
  where
    rollbackSafely :: IO ()
    rollbackSafely :: IO ()
rollbackSafely = do
      inTx <- Connection -> IO Bool
connectionInTransaction Connection
conn
      when inTx $ void (tryAny (runSQL conn "ROLLBACK"))

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