-- | LISTEN\/NOTIFY connections and listeners over libpq.
module Arbiter.LibPQ
  ( libPQListenConn
  , withLibPQListenConn
  , newLibPQListener
  ) where

import Arbiter.Core.Listen (ListenConn (..), Listener, Notification (..), newListener)
import Arbiter.Core.Listen.Driver
  ( ConnStatus (..)
  , ConnectDriver (..)
  , Polling (..)
  , execOutcome
  , withDriverListenConn
  )
import Control.Monad ((>=>))
import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.ByteString (ByteString)
import Database.PostgreSQL.LibPQ qualified as PQ

connectDriver :: ConnectDriver PQ.Connection
connectDriver :: ConnectDriver Connection
connectDriver =
  ConnectDriver
    { connectStart :: ByteString -> IO Connection
connectStart = ByteString -> IO Connection
PQ.connectStart
    , 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

-- | A 'ListenConn' over a libpq connection.
libPQListenConn :: PQ.Connection -> ListenConn
libPQListenConn :: Connection -> ListenConn
libPQListenConn 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 -> IO Bool
PQ.consumeInput Connection
conn
    , 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
    }

-- | Run an action on a libpq connection of its own, opened from a connection string.
withLibPQListenConn :: ByteString -> (ListenConn -> IO a) -> IO a
withLibPQListenConn :: forall a. ByteString -> (ListenConn -> IO a) -> IO a
withLibPQListenConn = ConnectDriver Connection
-> (Connection -> ListenConn)
-> ByteString
-> (ListenConn -> IO a)
-> IO a
forall conn a.
ConnectDriver conn
-> (conn -> ListenConn)
-> ByteString
-> (ListenConn -> IO a)
-> IO a
withDriverListenConn ConnectDriver Connection
connectDriver Connection -> ListenConn
libPQListenConn

-- | A 'Listener' over its own libpq connection, opened from a connection string.
newLibPQListener :: (MonadIO m) => ByteString -> m Listener
newLibPQListener :: forall (m :: * -> *). MonadIO m => ByteString -> m Listener
newLibPQListener ByteString
connStr = IO Listener -> m Listener
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (((ListenConn -> IO ()) -> IO ()) -> IO Listener
newListener (ByteString -> (ListenConn -> IO ()) -> IO ()
forall a. ByteString -> (ListenConn -> IO a) -> IO a
withLibPQListenConn ByteString
connStr))