HasqlBroadcast.hs

-- |
-- Module      : PostgresWebsockets.Broadcast
-- Description : Build a Hasql.Notifications based producer 'Multiplexer'.
--
-- Uses Broadcast module adding database as a source producer.
-- This module provides a function to produce a 'Multiplexer' from a Hasql 'Connection'.
-- The producer issues a LISTEN command upon Open commands and UNLISTEN upon Close.
module PostgresWebsockets.HasqlBroadcast
  ( newHasqlBroadcaster,
    newHasqlBroadcasterOrError,
    -- re-export
    acquire,
    relayMessages,
    relayMessagesForever,
  )
where

import APrelude
import Control.Retry (RetryStatus (..), capDelay, exponentialBackoff, retrying)
import Data.Aeson (Value (..), decode)
import qualified Data.Aeson.Key as Key
import qualified Data.Aeson.KeyMap as JSON
import Data.Either.Combinators (mapBoth)
import Hasql.Connection (Connection)
import Hasql.Errors (ConnectionError)
import qualified Hasql.Connection as HC
import qualified Hasql.Connection.Settings as HC
import qualified Hasql.Decoders as HD
import qualified Hasql.Encoders as HE
import Hasql.Notifications
import qualified Hasql.Session as H
import qualified Hasql.Statement as H
import PostgresWebsockets.Broadcast

acquire :: ByteString -> IO (Either ConnectionError Connection)
acquire settings = HC.acquire $ HC.connectionString $ decodeUtf8 settings
-- | Returns a multiplexer from a connection URI, keeps trying to connect in case there is any error.
--   This function also spawns a thread that keeps relaying the messages from the database to the multiplexer's listeners
newHasqlBroadcaster :: IO () -> Text -> Int -> Maybe Int -> ByteString -> IO Multiplexer
newHasqlBroadcaster onConnectionFailure ch maxRetries checkInterval = newHasqlBroadcasterForConnection . tryUntilConnected maxRetries
  where
    newHasqlBroadcasterForConnection = newHasqlBroadcasterForChannel onConnectionFailure ch checkInterval

-- | Returns a multiplexer from a connection URI or an error message on the left case
--   This function also spawns a thread that keeps relaying the messages from the database to the multiplexer's listeners
newHasqlBroadcasterOrError :: IO () -> Text -> ByteString -> IO (Either ByteString Multiplexer)
newHasqlBroadcasterOrError onConnectionFailure ch =
  acquire >=> (sequence . mapBoth showBS (newHasqlBroadcasterForConnection . return))
  where
    newHasqlBroadcasterForConnection = newHasqlBroadcasterForChannel onConnectionFailure ch Nothing

tryUntilConnected :: Int -> ByteString -> IO Connection
tryUntilConnected maxRetries =
  fmap (either (panic "Failure on connection retry") id) . retryConnection
  where
    retryConnection conStr = retrying retryPolicy shouldRetry (const $ acquire conStr)
    maxDelayInMicroseconds = 32000000
    firstDelayInMicroseconds = 1000000
    retryPolicy = capDelay maxDelayInMicroseconds $ exponentialBackoff firstDelayInMicroseconds
    shouldRetry :: RetryStatus -> Either ConnectionError Connection -> IO Bool
    shouldRetry RetryStatus {..} con =
      case con of
        Left err -> do
          putErrLn $ "Error connecting notification listener to database: " <> showText err
          pure $ rsIterNumber < maxRetries - 1
        _ -> return False

-- | Returns a multiplexer from a channel and an IO Connection, listen for different database notifications on the provided channel using the connection produced.
--
--   This function also spawns a thread that keeps relaying the messages from the database to the multiplexer's listeners
--
--   To listen on channels *chat*
--
--   @
--   import Protolude
--   import PostgresWebsockets.HasqlBroadcast
--   import PostgresWebsockets.Broadcast
--   import Hasql.Connection
--
--   main = do
--    conOrError <- H.acquire "postgres://localhost/test_database"
--    let con = either (panic . show) id conOrError :: Connection
--    multi <- newHasqlBroadcaster con
--
--    onMessage multi "chat" (\ch ->
--      forever $ fmap print (atomically $ readTChan ch)
--   @
newHasqlBroadcasterForChannel :: IO () -> Text -> Maybe Int -> IO Connection -> IO Multiplexer
newHasqlBroadcasterForChannel onConnectionFailure ch checkInterval getCon = do
  multi <- newMultiplexer openProducer $ const onConnectionFailure
  case checkInterval of
    Just i -> superviseMultiplexer multi i shouldRestart
    _ -> pure ()
  void $ relayMessagesForever multi
  return multi
  where
    toMsg :: Text -> Text -> Message
    toMsg c m = case decode (fromStrict $ encodeUtf8 m) of
      Just v -> Message (channelDef c v) m
      Nothing -> Message c m

    lookupStringDef :: Text -> Text -> Value -> Text
    lookupStringDef key d (Object obj) =
      case lookupDefault (String d) key obj of
        String s -> s
        _ -> d
    lookupStringDef _ d _ = d

    lookupDefault d key obj = fromMaybe d $ JSON.lookup (Key.fromText key) obj

    channelDef = lookupStringDef "channel"
    shouldRestart = do
      con <- getCon
      not <$> isListening con ch

    openProducer msgQ = do
      con <- getCon
      listen con $ toPgIdentifier ch
      waitForNotifications
        (\c m -> atomically $ writeTQueue msgQ $ toMsg (decodeUtf8 c) (decodeUtf8 m))
        con

isListening :: Connection -> Text -> IO Bool
isListening con ch = do
  resultOrError <- HC.use con session
  pure $ fromRight False resultOrError
  where
    session = H.statement chPattern isListeningStatement
    chPattern = "listen%" <> ch <> "%"

isListeningStatement :: H.Statement Text Bool
isListeningStatement =
  H.preparable sql encoder decoder
  where
    sql = "select exists (select * from pg_stat_activity where datname = current_database() and query ilike $1);"
    encoder = HE.param $ HE.nonNullable HE.text
    decoder = HD.singleRow (HD.column (HD.nonNullable HD.bool))