{-# LANGUAGE OverloadedStrings #-}

-- | @hasql@ implementation helpers for 'Arbiter.Core.MonadArbiter.MonadArbiter'.
--
-- Handlers receive a @Hasql.Connection.Connection@ for running typed hasql
-- queries inside the worker transaction:
--
-- @
-- import Arbiter.Hasql.MonadArbiter
-- import Hasql.Connection qualified as Hasql
--
-- instance MonadArbiter MyApp where
--   type RegistryOf MyApp = MyRegistry
--   type Handler MyApp job result = Hasql.Connection -> job -> MyApp result
--   getSchema                = asks appSchema
--   executeQuery             = hasqlExecuteQuery
--   executeStatement         = hasqlExecuteStatement
--   withDbTransaction        = hasqlWithDbTransaction
--   runHandlerWithConnection = hasqlRunHandlerWithConnection
--   getListener              = asks appListener
-- @
--
-- 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 pool management
  , HasqlConnectionPool (..)
  , HasHasqlPool (..)
  , localHasqlConnection
  ) where

import Arbiter.Core.Codec (RowCodec)
import Arbiter.Core.Exceptions (throwInternal)
import Arbiter.Core.MonadArbiter (Params, Query (..))
import Arbiter.Core.Sql.Query (numberPlaceholders)
import Control.Monad (when)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.ByteString.Char8 qualified as BSC
import Data.Int (Int64)
import Data.Pool qualified as Pool
import Data.Text (Text)
import Data.Text qualified as T
import Hasql.Connection qualified as Hasql
import Hasql.Session qualified as Session
import Hasql.Statement qualified as S
import UnliftIO (MonadUnliftIO, mask, onException, withRunInIO)
import UnliftIO.Exception (SomeException, try)

import Arbiter.Hasql.Compat qualified as Compat
import Arbiter.Hasql.Decode qualified as Decode
import Arbiter.Hasql.Encode qualified as Encode

-- | Connection pool state for hasql connections.
data HasqlConnectionPool = HasqlConnectionPool
  { HasqlConnectionPool -> Maybe (Pool Connection)
connectionPool :: Maybe (Pool.Pool Hasql.Connection)
  -- ^ The underlying resource pool. 'Nothing' when using connection-only mode
  -- via 'Arbiter.Hasql.HasqlDb.inTransaction'.
  , HasqlConnectionPool -> Maybe Connection
activeConn :: Maybe Hasql.Connection
  -- ^ Pinned connection when inside a transaction
  , HasqlConnectionPool -> Int
transactionDepth :: Int
  -- ^ Current nesting depth (0 = no active transaction)
  , HasqlConnectionPool -> Bool
preparedStatements :: Bool
  -- ^ Whether 'hasqlExecuteQueryPrepared' prepares server-side. Off by default.
  }

-- | Typeclass for monads that carry a hasql connection pool.
class (Monad m) => HasHasqlPool m where
  getHasqlPool :: m HasqlConnectionPool
  localHasqlPool :: (HasqlConnectionPool -> HasqlConnectionPool) -> m a -> m a

-- | Pin a hasql connection for the callback. Every arbiter operation inside it runs on
-- that connection. The caller has already issued its @BEGIN@.
localHasqlConnection :: (HasHasqlPool m) => Hasql.Connection -> m a -> m a
localHasqlConnection :: forall (m :: * -> *) a. HasHasqlPool m => Connection -> m a -> m a
localHasqlConnection Connection
conn = (HasqlConnectionPool -> HasqlConnectionPool) -> m a -> m a
forall a.
(HasqlConnectionPool -> HasqlConnectionPool) -> m a -> m a
forall (m :: * -> *) a.
HasHasqlPool m =>
(HasqlConnectionPool -> HasqlConnectionPool) -> m a -> m a
localHasqlPool (\HasqlConnectionPool
pool -> HasqlConnectionPool
pool {activeConn = Just conn, transactionDepth = 1})

-- | Run a query unprepared, decoding rows.
hasqlExecuteQuery
  :: (HasHasqlPool m, MonadIO m)
  => Query a
  -> m [a]
hasqlExecuteQuery :: forall (m :: * -> *) a.
(HasHasqlPool m, MonadIO m) =>
Query a -> m [a]
hasqlExecuteQuery (Query Text
sql Params
params RowCodec a
codec) = (Connection -> IO [a]) -> m [a]
forall (m :: * -> *) a.
(HasHasqlPool m, MonadIO m) =>
(Connection -> IO a) -> m a
withConn ((Connection -> IO [a]) -> m [a])
-> (Connection -> IO [a]) -> m [a]
forall a b. (a -> b) -> a -> b
$ \Connection
conn ->
  Bool -> Connection -> Text -> Params -> RowCodec a -> IO [a]
forall a.
Bool -> Connection -> Text -> Params -> RowCodec a -> IO [a]
runQueryStatement Bool
False Connection
conn Text
sql Params
params RowCodec a
codec

-- | 'hasqlExecuteQuery' that prepares the statement once per connection and reuses
-- the plan, when the pool enables prepared statements.
hasqlExecuteQueryPrepared
  :: (HasHasqlPool m, MonadIO m)
  => Query a
  -> m [a]
hasqlExecuteQueryPrepared :: forall (m :: * -> *) a.
(HasHasqlPool m, MonadIO m) =>
Query a -> m [a]
hasqlExecuteQueryPrepared (Query Text
sql Params
params RowCodec a
codec) = do
  pool <- m HasqlConnectionPool
forall (m :: * -> *). HasHasqlPool m => m HasqlConnectionPool
getHasqlPool
  withConn $ \Connection
conn -> Bool -> Connection -> Text -> Params -> RowCodec a -> IO [a]
forall a.
Bool -> Connection -> Text -> Params -> RowCodec a -> IO [a]
runQueryStatement (HasqlConnectionPool -> Bool
preparedStatements HasqlConnectionPool
pool) Connection
conn Text
sql Params
params RowCodec a
codec

runQueryStatement :: Bool -> Hasql.Connection -> Text -> Params -> RowCodec a -> IO [a]
runQueryStatement :: forall a.
Bool -> Connection -> Text -> Params -> RowCodec a -> IO [a]
runQueryStatement Bool
prepare Connection
conn Text
sql Params
params RowCodec a
codec = do
  let 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
      stmt :: Statement () [a]
stmt = Text -> Params () -> Result [a] -> Statement () [a]
mkStatement (Text -> Text
numberPlaceholders Text
sql) (Params -> Params ()
Encode.buildEncoder Params
params) (RowCodec a -> Result [a]
forall a. RowCodec a -> Result [a]
Decode.hasqlRowDecoder RowCodec a
codec)
  result <- Connection -> Session [a] -> IO (Either SessionError [a])
forall a. Connection -> Session a -> IO (Either SessionError a)
Hasql.use Connection
conn (() -> Statement () [a] -> Session [a]
forall params result.
params -> Statement params result -> Session result
Session.statement () Statement () [a]
stmt)
  case result of
    Right [a]
rows -> [a] -> IO [a]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [a]
rows
    Left SessionError
err -> Text -> IO [a]
forall (m :: * -> *) a. MonadIO m => Text -> m a
throwInternal (Text -> IO [a]) -> Text -> IO [a]
forall a b. (a -> b) -> a -> b
$ Text
"hasql query error: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (SessionError -> String
forall a. Show a => a -> String
show SessionError
err)

-- | Run a statement unprepared, returning rows affected.
hasqlExecuteStatement
  :: (HasHasqlPool m, MonadIO m)
  => Query a
  -> m Int64
hasqlExecuteStatement :: forall (m :: * -> *) a.
(HasHasqlPool m, MonadIO m) =>
Query a -> m Int64
hasqlExecuteStatement (Query Text
sql Params
params RowCodec a
_) = (Connection -> IO Int64) -> m Int64
forall (m :: * -> *) a.
(HasHasqlPool m, MonadIO m) =>
(Connection -> IO a) -> m a
withConn ((Connection -> IO Int64) -> m Int64)
-> (Connection -> IO Int64) -> m Int64
forall a b. (a -> b) -> a -> b
$ \Connection
conn -> IO Int64 -> IO Int64
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int64 -> IO Int64) -> IO Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ do
  let stmt :: Statement () Int64
stmt = Text -> Params -> Statement () Int64
Encode.buildStatementRowCount Text
sql Params
params
  result <- Connection -> Session Int64 -> IO (Either SessionError Int64)
forall a. Connection -> Session a -> IO (Either SessionError a)
Hasql.use Connection
conn (() -> Statement () Int64 -> Session Int64
forall params result.
params -> Statement params result -> Session result
Session.statement () Statement () Int64
stmt)
  case result of
    Right Int64
rowCount -> Int64 -> IO Int64
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int64
rowCount
    Left SessionError
err -> Text -> IO Int64
forall (m :: * -> *) a. MonadIO m => Text -> m a
throwInternal (Text -> IO Int64) -> Text -> IO Int64
forall a b. (a -> b) -> a -> b
$ Text
"hasql statement error: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (SessionError -> String
forall a. Show a => a -> String
show SessionError
err)

-- | Transaction bracket. Nests via savepoints.
hasqlWithDbTransaction :: (HasHasqlPool m, MonadUnliftIO m) => m a -> m a
hasqlWithDbTransaction :: forall (m :: * -> *) a.
(HasHasqlPool m, MonadUnliftIO m) =>
m a -> m a
hasqlWithDbTransaction m a
action = do
  pool <- m HasqlConnectionPool
forall (m :: * -> *). HasHasqlPool m => m HasqlConnectionPool
getHasqlPool
  let depth = HasqlConnectionPool -> Int
transactionDepth HasqlConnectionPool
pool
  case (activeConn pool, depth) of
    (Maybe Connection
Nothing, Int
_) -> case HasqlConnectionPool -> Maybe (Pool Connection)
connectionPool HasqlConnectionPool
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
Pool.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
beginCommitOrRollback 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 ((HasqlConnectionPool -> HasqlConnectionPool) -> m a -> m a
forall a.
(HasqlConnectionPool -> HasqlConnectionPool) -> m a -> m a
forall (m :: * -> *) a.
HasHasqlPool m =>
(HasqlConnectionPool -> HasqlConnectionPool) -> m a -> m a
localHasqlPool (\HasqlConnectionPool
hpool -> HasqlConnectionPool
hpool {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
beginCommitOrRollback 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 ((HasqlConnectionPool -> HasqlConnectionPool) -> m a -> m a
forall a.
(HasqlConnectionPool -> HasqlConnectionPool) -> m a -> m a
forall (m :: * -> *) a.
HasHasqlPool m =>
(HasqlConnectionPool -> HasqlConnectionPool) -> m a -> m a
localHasqlPool (\HasqlConnectionPool
hpool -> HasqlConnectionPool
hpool {transactionDepth = 1}) m a
action)
    (Just Connection
conn, Int
nestedDepth) -> ((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 :: ByteString
spName = 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
nestedDepth)
      IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Connection -> ByteString -> IO ()
forall (m :: * -> *).
MonadUnliftIO m =>
Connection -> ByteString -> m ()
Compat.runSQL Connection
conn (ByteString
"SAVEPOINT " ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
spName)
      result <-
        m a -> m a
forall a. m a -> m a
restore ((HasqlConnectionPool -> HasqlConnectionPool) -> m a -> m a
forall a.
(HasqlConnectionPool -> HasqlConnectionPool) -> m a -> m a
forall (m :: * -> *) a.
HasHasqlPool m =>
(HasqlConnectionPool -> HasqlConnectionPool) -> m a -> m a
localHasqlPool (\HasqlConnectionPool
hpool -> HasqlConnectionPool
hpool {transactionDepth = nestedDepth + 1}) m a
action)
          m a -> m () -> m a
forall (m :: * -> *) a b. MonadUnliftIO m => m a -> m b -> m a
`onException` IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Connection -> ByteString -> IO ()
forall (m :: * -> *).
MonadUnliftIO m =>
Connection -> ByteString -> m ()
Compat.runSQL Connection
conn (ByteString
"ROLLBACK TO SAVEPOINT " ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
spName))
      liftIO $ Compat.runSQL conn ("RELEASE SAVEPOINT " <> spName)
      pure result

beginCommitOrRollback :: forall a. 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 ()
forall (m :: * -> *).
MonadUnliftIO m =>
Connection -> ByteString -> m ()
Compat.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
  Compat.runSQL conn "COMMIT"
  pure result
  where
    rollbackSafely :: IO ()
    rollbackSafely :: IO ()
rollbackSafely = do
      inTx <- Connection -> IO Bool
Compat.connectionInTransaction Connection
conn
      when inTx $ do
        _ <- try (Compat.runSQL conn "ROLLBACK") :: IO (Either SomeException ())
        pure ()

-- | Run a handler on the active connection. The handler can issue typed hasql queries
-- inside the worker transaction.
hasqlRunHandlerWithConnection
  :: (HasHasqlPool m, MonadIO m)
  => (Hasql.Connection -> job -> m result)
  -> job
  -> m result
hasqlRunHandlerWithConnection :: forall (m :: * -> *) job result.
(HasHasqlPool m, MonadIO m) =>
(Connection -> job -> m result) -> job -> m result
hasqlRunHandlerWithConnection Connection -> job -> m result
handler job
job = do
  pool <- m HasqlConnectionPool
forall (m :: * -> *). HasHasqlPool m => m HasqlConnectionPool
getHasqlPool
  case activeConn pool of
    Just Connection
conn -> Connection -> job -> m result
handler Connection
conn job
job
    Maybe Connection
Nothing -> Text -> m result
forall (m :: * -> *) a. MonadIO m => Text -> m a
throwInternal Text
"hasqlRunHandlerWithConnection: no active connection"

-- ---------------------------------------------------------------------------
-- Internal
-- ---------------------------------------------------------------------------

-- | The pinned connection, or one checked out of the pool.
withConn :: (HasHasqlPool m, MonadIO m) => (Hasql.Connection -> IO a) -> m a
withConn :: forall (m :: * -> *) a.
(HasHasqlPool m, MonadIO m) =>
(Connection -> IO a) -> m a
withConn Connection -> IO a
action = do
  pool <- m HasqlConnectionPool
forall (m :: * -> *). HasHasqlPool m => m HasqlConnectionPool
getHasqlPool
  case (activeConn pool, connectionPool pool) of
    (Just Connection
conn, Maybe (Pool Connection)
_) -> 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
$ Connection -> IO a
action Connection
conn
    (Maybe Connection
Nothing, Just Pool Connection
connPool) -> 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
$ Pool Connection -> (Connection -> IO a) -> IO a
forall a r. Pool a -> (a -> IO r) -> IO r
Pool.withResource Pool Connection
connPool Connection -> IO 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"