{-# LANGUAGE CPP #-}
{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS_HADDOCK not-home #-}

-- | Internal to the arbiter packages. Not covered by the PVP, except for the names
-- that "Arbiter.Hasql.HasqlDb" re-exports.
--
-- Every hasql version difference that arbiter-hasql depends on.
module Arbiter.Hasql.Compat
  ( runSQL
  , connectionInTransaction
  , withHasqlListenConn
  , hasqlSettings
  , HasqlSettings
  , HasqlConnect
  , toHasqlConnect
  , acquireConnect
  , withDedicatedListenConn
  ) where

import Arbiter.Core.Enum (showText)
import Arbiter.Core.Exceptions (orThrowInternal, throwInternal)
import Arbiter.Core.Listen (ListenConn (..))
import Data.Bifunctor (first)
import Data.ByteString (ByteString)
import Data.Either (fromRight)
import Data.Text (Text)
import Data.Text.Encoding qualified as TE
import Data.Text.Encoding.Error qualified as TE
import Hasql.Connection qualified as Hasql
import Hasql.Connection.Settings qualified as Settings
import Hasql.Session qualified as Session

#if MIN_VERSION_hasql(2,0,0)
import Arbiter.Core.Listen (Notification (..))
import Arbiter.Core.Listen.Driver
  ( ConnStatus (..)
  , ConnectDriver (..)
  , Polling (..)
  , execOutcome
  , withDriverListenConn
  )
import Control.Monad ((>=>))
import Data.Either (isRight)
import Pqi qualified as PQ
#else
import Arbiter.LibPQ (libPQListenConn, withLibPQListenConn)
import Database.PostgreSQL.LibPQ qualified as PQ
#endif

-- | Run a bare SQL command, such as @BEGIN@ or @COMMIT@.
runSQL :: Hasql.Connection -> ByteString -> IO ()
runSQL :: Connection -> ByteString -> IO ()
runSQL Connection
conn ByteString
sql =
  Connection -> Session () -> IO (Either SessionError ())
forall a. Connection -> Session a -> IO (Either SessionError a)
Hasql.use Connection
conn (Text -> Session ()
Session.script (ByteString -> Text
lenientText ByteString
sql))
    IO (Either SessionError ())
-> (Either SessionError () -> IO ()) -> IO ()
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 () -> IO ()
forall (m :: * -> *) e a.
MonadIO m =>
(e -> Text) -> Either e a -> m a
orThrowInternal ((Text
"hasql runSQL error: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>) (Text -> Text) -> (SessionError -> Text) -> SessionError -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SessionError -> Text
forall a. Show a => a -> Text
showText)

-- | Decode UTF-8, replacing invalid bytes.
lenientText :: ByteString -> Text
lenientText :: ByteString -> Text
lenientText = OnDecodeError -> ByteString -> Text
TE.decodeUtf8With OnDecodeError
TE.lenientDecode

#if MIN_VERSION_hasql(2,0,0)
-- | Connection settings for hasql 2: transport adapter (for example @Pqi.Ffi.adapter@) and connection string.
data HasqlConnect = HasqlConnect PQ.Adapter ByteString

-- | Connection settings from a pqi adapter and a connection string. With the @hasql2@
-- flag off, the hasql 1.x shape takes only the connection string.
toHasqlConnect :: PQ.Adapter -> ByteString -> HasqlConnect
toHasqlConnect :: Adapter -> ByteString -> HasqlConnect
toHasqlConnect = Adapter -> ByteString -> HasqlConnect
HasqlConnect

-- | Open a connection. A failure comes back as its description.
acquireConnect :: HasqlConnect -> IO (Either String Hasql.Connection)
acquireConnect :: HasqlConnect -> IO (Either String Connection)
acquireConnect (HasqlConnect Adapter
adapter ByteString
connStr) = (ConnectionError -> String)
-> Either ConnectionError Connection -> Either String Connection
forall a b c. (a -> b) -> Either a c -> Either b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ConnectionError -> String
forall a. Show a => a -> String
show (Either ConnectionError Connection -> Either String Connection)
-> IO (Either ConnectionError Connection)
-> IO (Either String Connection)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Adapter -> Settings -> IO (Either ConnectionError Connection)
Hasql.acquire Adapter
adapter (ByteString -> Settings
hasqlSettings ByteString
connStr)

-- | Run the listener loop on a driver connection of its own.
withDedicatedListenConn :: HasqlConnect -> (ListenConn -> IO a) -> IO a
withDedicatedListenConn :: forall a. HasqlConnect -> (ListenConn -> IO a) -> IO a
withDedicatedListenConn (HasqlConnect Adapter
adapter ByteString
connStr) = ConnectDriver Connection
-> (Connection -> ListenConn)
-> ByteString
-> (ListenConn -> IO a)
-> IO a
forall conn a.
ConnectDriver conn
-> (conn -> ListenConn)
-> ByteString
-> (ListenConn -> IO a)
-> IO a
withDriverListenConn (Adapter -> ConnectDriver Connection
pqiConnectDriver Adapter
adapter) Connection -> ListenConn
toListenConn ByteString
connStr
#else
-- | Connection settings for hasql 1.x: connection string.
newtype HasqlConnect = HasqlConnect ByteString

-- | Connection settings from a connection string. With the @hasql2@ flag on, the
-- hasql 2 shape also takes a pqi adapter first.
toHasqlConnect :: ByteString -> HasqlConnect
toHasqlConnect = HasqlConnect

-- | Open a connection. A failure comes back as its description.
acquireConnect :: HasqlConnect -> IO (Either String Hasql.Connection)
acquireConnect (HasqlConnect connStr) = first show <$> Hasql.acquire (hasqlSettings connStr)

-- | Run the listener loop on a driver connection of its own.
withDedicatedListenConn :: HasqlConnect -> (ListenConn -> IO a) -> IO a
withDedicatedListenConn (HasqlConnect connStr) = withLibPQListenConn connStr
#endif

-- | Whether the connection is in a transaction block, valid or aborted.
connectionInTransaction :: Hasql.Connection -> IO Bool
connectionInTransaction :: Connection -> IO Bool
connectionInTransaction Connection
conn =
  Bool -> Either SessionError Bool -> Bool
forall b a. b -> Either a b -> b
fromRight Bool
False
    (Either SessionError Bool -> Bool)
-> IO (Either SessionError Bool) -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Connection -> Session Bool -> IO (Either SessionError Bool)
forall a. Connection -> Session a -> IO (Either SessionError a)
Hasql.use
      Connection
conn
      ( (Connection -> IO (Either SessionError Bool, Connection))
-> Session Bool
forall a.
(Connection -> IO (Either SessionError a, Connection)) -> Session a
Session.onLibpqConnection ((Connection -> IO (Either SessionError Bool, Connection))
 -> Session Bool)
-> (Connection -> IO (Either SessionError Bool, Connection))
-> Session Bool
forall a b. (a -> b) -> a -> b
$ \Connection
libpq -> do
          status <- Connection -> IO TransactionStatus
PQ.transactionStatus Connection
libpq
          pure (Right (txStatusNeedsRollback status), libpq)
      )

-- | Run the listener loop on the connection's driver handle. The loop runs outside the
-- session.
withHasqlListenConn :: Hasql.Connection -> (ListenConn -> IO a) -> IO a
withHasqlListenConn :: forall a. Connection -> (ListenConn -> IO a) -> IO a
withHasqlListenConn Connection
conn ListenConn -> IO a
action = do
  result <- Connection
-> Session Connection -> IO (Either SessionError Connection)
forall a. Connection -> Session a -> IO (Either SessionError a)
Hasql.use Connection
conn (Session Connection -> IO (Either SessionError Connection))
-> Session Connection -> IO (Either SessionError Connection)
forall a b. (a -> b) -> a -> b
$ (Connection -> IO (Either SessionError Connection, Connection))
-> Session Connection
forall a.
(Connection -> IO (Either SessionError a, Connection)) -> Session a
Session.onLibpqConnection ((Connection -> IO (Either SessionError Connection, Connection))
 -> Session Connection)
-> (Connection -> IO (Either SessionError Connection, Connection))
-> Session Connection
forall a b. (a -> b) -> a -> b
$ \Connection
libpq -> (Either SessionError Connection, Connection)
-> IO (Either SessionError Connection, Connection)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Connection -> Either SessionError Connection
forall a b. b -> Either a b
Right Connection
libpq, Connection
libpq)
  either (const (throwInternal "connection lost")) (action . toListenConn) result

#if MIN_VERSION_hasql(2,0,0)
-- | A 'ListenConn' over a pqi connection. The native transport reads the socket only
-- inside a query, so an empty query drains it.
toListenConn :: PQ.Connection -> ListenConn
toListenConn :: Connection -> ListenConn
toListenConn Connection
conn =
  ListenConn
    { listenNotifies :: IO (Maybe Notification)
listenNotifies = (Notify -> Notification) -> Maybe Notify -> Maybe Notification
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ByteString -> ByteString -> Notification
Notification (ByteString -> ByteString -> Notification)
-> (Notify -> ByteString) -> Notify -> ByteString -> Notification
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Notify -> ByteString
PQ.notifyRelname (Notify -> ByteString -> Notification)
-> (Notify -> ByteString) -> Notify -> Notification
forall a b. (Notify -> a -> b) -> (Notify -> a) -> Notify -> b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Notify -> ByteString
PQ.notifyExtra) (Maybe Notify -> Maybe Notification)
-> IO (Maybe Notify) -> IO (Maybe Notification)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Connection -> IO (Maybe Notify)
PQ.notifies Connection
conn
    , listenSocket :: IO (Maybe Fd)
listenSocket = Connection -> IO (Maybe Fd)
PQ.socket Connection
conn
    , listenConsumeInput :: IO Bool
listenConsumeInput = Connection -> ByteString -> IO (Maybe Result)
PQ.exec Connection
conn ByteString
"" IO (Maybe Result) -> (Maybe Result -> IO Bool) -> IO Bool
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Either Text () -> Bool) -> IO (Either Text ()) -> IO Bool
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Either Text () -> Bool
forall a b. Either a b -> Bool
isRight (IO (Either Text ()) -> IO Bool)
-> (Maybe Result -> IO (Either Text ())) -> Maybe Result -> IO Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExecStatus
-> (Result -> IO ExecStatus) -> Maybe Result -> IO (Either Text ())
forall status result.
(Eq status, Show status) =>
status
-> (result -> IO status) -> Maybe result -> IO (Either Text ())
execOutcome ExecStatus
PQ.EmptyQuery Result -> IO ExecStatus
PQ.resultStatus
    , listenExec :: ByteString -> IO (Either Text ())
listenExec = Connection -> ByteString -> IO (Maybe Result)
PQ.exec Connection
conn (ByteString -> IO (Maybe Result))
-> (Maybe Result -> IO (Either Text ()))
-> ByteString
-> IO (Either Text ())
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> ExecStatus
-> (Result -> IO ExecStatus) -> Maybe Result -> IO (Either Text ())
forall status result.
(Eq status, Show status) =>
status
-> (result -> IO status) -> Maybe result -> IO (Either Text ())
execOutcome ExecStatus
PQ.CommandOk Result -> IO ExecStatus
PQ.resultStatus
    , listenEscapeIdentifier :: ByteString -> IO (Maybe ByteString)
listenEscapeIdentifier = Connection -> ByteString -> IO (Maybe ByteString)
PQ.escapeIdentifier Connection
conn
    }

pqiConnectDriver :: PQ.Adapter -> ConnectDriver PQ.Connection
pqiConnectDriver :: Adapter -> ConnectDriver Connection
pqiConnectDriver Adapter
adapter =
  ConnectDriver
    { connectStart :: ByteString -> IO Connection
connectStart = Adapter -> ByteString -> IO Connection
PQ.connectStart Adapter
adapter
    , connectPoll :: Connection -> IO Polling
connectPoll = (PollingStatus -> Polling) -> IO PollingStatus -> IO Polling
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap PollingStatus -> Polling
polling (IO PollingStatus -> IO Polling)
-> (Connection -> IO PollingStatus) -> Connection -> IO Polling
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Connection -> IO PollingStatus
PQ.connectPoll
    , connectStatus :: Connection -> IO ConnStatus
connectStatus = (ConnStatus -> ConnStatus) -> IO ConnStatus -> IO ConnStatus
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ConnStatus -> ConnStatus
connStatus (IO ConnStatus -> IO ConnStatus)
-> (Connection -> IO ConnStatus) -> Connection -> IO ConnStatus
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Connection -> IO ConnStatus
PQ.status
    , connectFinish :: Connection -> IO ()
connectFinish = Connection -> IO ()
PQ.finish
    , connectErrorMessage :: Connection -> IO (Maybe ByteString)
connectErrorMessage = Connection -> IO (Maybe ByteString)
PQ.errorMessage
    }
  where
    polling :: PollingStatus -> Polling
polling PollingStatus
PQ.PollingReading = Polling
PollReading
    polling PollingStatus
PQ.PollingWriting = Polling
PollWriting
    polling PollingStatus
_ = Polling
PollDone
    connStatus :: ConnStatus -> ConnStatus
connStatus ConnStatus
PQ.ConnectionOk = ConnStatus
ConnOk
    connStatus ConnStatus
PQ.ConnectionBad = ConnStatus
ConnBad
    connStatus ConnStatus
_ = ConnStatus
ConnPending
#else
toListenConn :: PQ.Connection -> ListenConn
toListenConn = libPQListenConn
#endif

-- | @TransInTrans@ and @TransInError@ accept a @ROLLBACK@ without warning.
txStatusNeedsRollback :: PQ.TransactionStatus -> Bool
txStatusNeedsRollback :: TransactionStatus -> Bool
txStatusNeedsRollback TransactionStatus
PQ.TransInTrans = Bool
True
txStatusNeedsRollback TransactionStatus
PQ.TransInError = Bool
True
txStatusNeedsRollback TransactionStatus
_ = Bool
False

-- | hasql connection settings.
type HasqlSettings = Settings.Settings

-- | Convert a connection string ByteString to hasql settings.
hasqlSettings :: ByteString -> HasqlSettings
hasqlSettings :: ByteString -> Settings
hasqlSettings = Text -> Settings
Settings.connectionString (Text -> Settings)
-> (ByteString -> Text) -> ByteString -> Settings
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Text
lenientText