Simple.hs

{-|
Module      : Database.PGMQ.Simple
Description : Basic functionality
Copyright   : (c) Gargantext, 2024-Present
License     : AGPL
Maintainer  : gargantext@iscpif.fr
Stability   : experimental
Portability : POSIX

https://tembo.io/pgmq/api/sql/functions/
-}

{-# LANGUAGE QuasiQuotes #-}

module Database.PGMQ.Simple
  ( initialize
  , archiveMessage
  , archiveMessages
  , createQueue
  , deleteMessage
  , deleteMessages
  , dropQueue
  , getMetrics
  , getMetricsAll
  , listQueues
  , popMessage
  , purgeQueue
  , readMessage
  , readMessageFromArchive
  , readMessageById
  , readMessages
  , readMessageWithPoll
  , readMessagesWithPoll
  , sendMessage
  , sendMessages
  , setMessageVt
  , queueAvailableLength
  , queueAvailableIds )
where


-- Please note that message format should be JSON-serializable (the
-- 'pgmq.send' functions accept only a ::jsonb type).


-- import Control.Exception (Exception, SomeException(..), catch, fromException, throwIO, toException)
import Control.Monad (void)
import Data.Functor ((<&>))
import Data.String (fromString)
import Data.Text qualified as T
import Database.PGMQ.Schema (pgmqSchema)
import Database.PostgreSQL.Simple qualified as PSQL
import Database.PostgreSQL.Simple.Newtypes qualified as PSQL (Aeson(..))
import Database.PostgreSQL.Simple.SqlQQ (sql)
import Database.PostgreSQL.Simple.Types qualified as PSQL (QualifiedIdentifier(..), PGArray(..))
import Database.PGMQ.Types (Delay, MaxPollSeconds, Message, MessageCount, MessageId, Metrics, SerializableMessage, PollIntervalMs, Queue, QueueInfo, VisibilityTimeout)
import Safe (headMay)


{-| Initialize PGMQ given a PostgreSQL connection. Mainly concerned
    with creating the 'pgmq' extension. -}
initialize :: PSQL.Connection -> IO ()
initialize conn = do
  -- OK so this is a bit tricky because of the usage of IF NOT EXISTS:
  -- https://stackoverflow.com/questions/29900845/create-schema-if-not-exists-raises-duplicate-key-error
  -- PostgreSQL will complain badly if 'initialize' is called
  -- from multiple threads at once.
  -- Hence, we use 'pg_advisory_xact_lock' to lock ourselves
  -- out of this situation.
  pgmqSchema' <- pgmqSchema
  PSQL.withTransaction conn $ do
    let magicLockId = 1122334455 :: Int
    _ <- PSQL.query conn [sql| SELECT pg_advisory_xact_lock(?) |] (PSQL.Only magicLockId) :: IO [PSQL.Only ()]
    hasPgmq <- PSQL.query_ conn [sql|SELECT 1 FROM information_schema.schemata WHERE schema_name = 'pgmq'|] :: IO [PSQL.Only Int]
    case hasPgmq of
      [PSQL.Only 1] -> pure ()
      _ -> do
        -- NOTE This isn't exactly right because we initialize the
        -- extension by only running the SQL, not using `CREATE
        -- EXTENSION`. But it works and is quite hassle-free
        void $ PSQL.execute_ conn [sql|CREATE SCHEMA pgmq |]
        void $ PSQL.execute_ conn (fromString pgmqSchema')
    -- void $ PSQL.execute_ conn [sql| CREATE EXTENSION IF NOT EXISTS pgmq |]


{-| Archives message in given queue for given id

   https://tembo.io/pgmq/api/sql/functions/#archive-single -}
archiveMessage :: PSQL.Connection -> Queue -> MessageId -> IO ()
archiveMessage conn queue msgId =
  void (PSQL.query conn [sql| SELECT pgmq.archive(?, ?) |] (queue, msgId) :: IO [PSQL.Only Bool])

{-| Archives messages in given queue for given ids

   https://tembo.io/pgmq/api/sql/functions/#archive-batch -}
archiveMessages :: PSQL.Connection -> Queue -> [MessageId] -> IO ()
archiveMessages conn queue msgIds =
  void (PSQL.query conn [sql| SELECT pgmq.archive(?, ?::bigint[]) |] (queue, PSQL.PGArray msgIds) :: IO [PSQL.Only Int])

{-| Creates a queue

  https://tembo.io/pgmq/api/sql/functions/#create -}
createQueue :: PSQL.Connection -> Queue -> IO ()
createQueue conn queue =
  void (PSQL.query conn [sql| SELECT pgmq.create(?) |] (PSQL.Only queue) :: IO [PSQL.Only ()])

{-| Deletes given message from given queue

  https://tembo.io/pgmq/api/sql/functions/#delete-single -}
deleteMessage :: PSQL.Connection -> Queue -> MessageId -> IO ()
deleteMessage conn queue msgId =
  void (PSQL.query conn [sql| SELECT pgmq.delete(?, ?) |] (queue, msgId) :: IO [PSQL.Only Bool])

{-| Deletes given messages from given queue

  https://tembo.io/pgmq/api/sql/functions/#delete-batch -}
deleteMessages :: PSQL.Connection -> Queue -> [MessageId] -> IO ()
deleteMessages conn queue msgIds =
  void (PSQL.query conn [sql| SELECT pgmq.delete(?, ?) |] (queue, PSQL.PGArray msgIds) :: IO [PSQL.Only Int])

{-| Deletes given queue

  https://tembo.io/pgmq/api/sql/functions/#drop_queue -}
dropQueue :: PSQL.Connection -> Queue -> IO ()
dropQueue conn queue =
  void (PSQL.query conn [sql| SELECT pgmq.drop_queue(?) |] (PSQL.Only queue) :: IO [PSQL.Only Bool])

{-| Read metrics for a given queue

  https://tembo.io/pgmq/api/sql/functions/#metrics -}
getMetrics :: PSQL.Connection -> Queue -> IO (Maybe Metrics)
getMetrics conn queue =
  PSQL.query conn [sql| SELECT * FROM pgmq.metrics(?) |] (PSQL.Only queue) <&> headMay
  -- catch
  --   (PSQL.query conn [sql| SELECT * FROM pgmq.metrics(?) |] (PSQL.Only queue) >>= return . headMay)
  --   handleError
  -- where
  --   -- support the case when a table does not exist
  --   -- handleError :: Exception e => e -> IO (Maybe Metrics)
  --   handleError err@(SomeException _) = do
  --     putStrLn $ "Error: " <> show err
  --     return Nothing
    -- handleError :: SomeException -> IO (Maybe Metrics)
    -- handleError err = do
    --   putStrLn $ "Error: " <> show err
    --   putStrLn "x"
    --   return Nothing
      -- case fromException err of
      --   -- Just (PSQL.SomePostgreSqlException e) ->
      --   --   case fromException (toException e) of
      --   --     Just (PSQL.SqlError { sqlState = "42P01" }) -> return Nothing
      --   --     -- re-raise other errors
      --   --     _ -> throwIO err
      --   Just (PSQL.SqlError { sqlState = "42P01" }) -> return Nothing
      --   -- re-raise other errors
      --   _ -> throwIO err


-- | Read metrics for all queues
--   https://tembo.io/pgmq/api/sql/functions/#metrics_all
getMetricsAll :: PSQL.Connection -> IO [Metrics]
getMetricsAll conn =
  PSQL.query_ conn [sql| SELECT * FROM pgmq.metrics_all() |]

-- | List all pgmq queues
--   https://tembo.io/pgmq/api/sql/functions/#list_queues
listQueues :: PSQL.Connection -> IO [QueueInfo]
listQueues conn =
  PSQL.query_ conn [sql| SELECT * FROM pgmq.list_queues() |]

-- | Read a message and immediately delete it from the queue. Returns `None` if the queue is empty.
--   https://tembo.io/pgmq/api/sql/functions/#pop
popMessage :: (SerializableMessage a)
           => PSQL.Connection -> Queue -> IO (Maybe (Message a))
popMessage conn queue = do
  ret <- PSQL.query conn
           [sql| SELECT msg_id, read_ct, enqueued_at, NULL, vt, message FROM pgmq.pop(?) |]
           (PSQL.Only queue)

  return $ headMay ret

-- | Deletes all messages from a queue
--   https://tembo.io/pgmq/api/sql/functions/#purge_queue
purgeQueue :: PSQL.Connection -> Queue -> IO ()
purgeQueue conn queue = do
  void (PSQL.query conn [sql| SELECT * FROM pgmq.purge_queue(?) |] (PSQL.Only queue) :: IO [PSQL.Only Int])

-- | Read a message from given queue, with given visibility timeout (in seconds)
--   https://tembo.io/pgmq/api/sql/functions/#read
readMessage :: (SerializableMessage a)
            => PSQL.Connection -> Queue -> VisibilityTimeout -> IO (Maybe (Message a))
readMessage conn queue vt =
  readMessages conn queue vt 1 <&> headMay

-- | Read a message from given archive
readMessageFromArchive :: (SerializableMessage a)
            => PSQL.Connection -> Queue -> MessageId -> IO (Maybe (Message a))
readMessageFromArchive conn queue msgId = do
  let archiveTable = "a_" <> queue
  PSQL.query conn
    [sql| SELECT msg_id, read_ct, enqueued_at, archived_at, vt, message FROM ? WHERE msg_id = ?|]
      (PSQL.QualifiedIdentifier (Just "pgmq") $ T.pack archiveTable, msgId) <&> headMay

-- | Read a message from queue with given ID (for querying purposes, doesn't pop message from queue)
readMessageById :: (SerializableMessage a)
                => PSQL.Connection -> Queue -> MessageId -> IO (Maybe (Message a))
readMessageById conn queue msgId = do
  let queueTable = "q_" <> queue
  PSQL.query conn
    [sql| SELECT msg_id, read_ct, enqueued_at, NULL, vt, message FROM ? WHERE msg_id = ?|]
      (PSQL.QualifiedIdentifier (Just "pgmq") $ T.pack queueTable, msgId) <&> headMay

{-| Reads given number of messages from given queue

  https://tembo.io/pgmq/api/sql/functions/#read -}
readMessages :: (SerializableMessage a)
             => PSQL.Connection -> Queue -> VisibilityTimeout -> MessageCount -> IO [Message a]
readMessages conn queue vt count =
  PSQL.query conn
    [sql| SELECT msg_id, read_ct, enqueued_at, NULL, vt, message FROM pgmq.read(?, ?, ?) |]
      (queue, vt, count)

{-| Reads a single message, polling for given duration if the queue
  is empty.
  
  NOTE This is a blocking operation.
  
  https://tembo.io/pgmq/api/sql/functions/#read_with_poll -}
readMessageWithPoll :: (SerializableMessage a)
                    => PSQL.Connection
                    -> Queue
                    -> VisibilityTimeout
                    -> MaxPollSeconds
                    -> PollIntervalMs
                    -> IO (Maybe (Message a))
readMessageWithPoll conn queue vt maxPollSeconds pollIntervalMs =
  readMessagesWithPoll conn queue vt 1 maxPollSeconds pollIntervalMs <&> headMay

-- | Reads given number of messages, polling for given duration if the
--   queue is empty.
--   NOTE This is a blocking operation.
--   https://tembo.io/pgmq/api/sql/functions/#read_with_poll
readMessagesWithPoll :: (SerializableMessage a)
                     => PSQL.Connection
                     -> Queue
                     -> VisibilityTimeout
                     -> MessageCount
                     -> MaxPollSeconds
                     -> PollIntervalMs
                     -> IO [Message a]
readMessagesWithPoll conn queue vt count maxPollSeconds pollIntervalMs =
  PSQL.query conn
    [sql| SELECT msg_id, read_ct, enqueued_at, NULL, vt, message FROM pgmq.read_with_poll(?, ?, ?, ?, ?) |]
      (queue, vt, count, maxPollSeconds, pollIntervalMs)

-- | Sends one message to a queue
--   https://tembo.io/pgmq/api/sql/functions/#send
sendMessage :: (SerializableMessage a)
            => PSQL.Connection -> Queue -> a -> Delay -> IO Int
sendMessage conn queue msg delay = do
  [PSQL.Only msgId] <- PSQL.query conn [sql| SELECT pgmq.send(?, ?::jsonb, ?) |] (queue, PSQL.Aeson msg, delay)

  return msgId


-- | Sends a batch of messages
--   https://tembo.io/pgmq/api/sql/functions/#send_batch
sendMessages :: (SerializableMessage a)
             => PSQL.Connection -> Queue -> [a] -> Delay -> IO [Int]
sendMessages conn queue msgs delay = do
  ret <- PSQL.query conn [sql| SELECT pgmq.send_batch(?, ?::jsonb[], ?) |] (queue, PSQL.PGArray (PSQL.Aeson <$> msgs), delay)

  return $ (\(PSQL.Only id') -> id') <$> ret

-- | Sets the visibility timeout of a message for X seconds from now
--   https://tembo.io/pgmq/api/sql/functions/#set_vt
setMessageVt :: PSQL.Connection -> Queue -> MessageId -> VisibilityTimeout -> IO ()
setMessageVt conn queue msgId vt =
  void (PSQL.query conn [sql| SELECT 1 FROM pgmq.set_vt(?, ?, ?) |] (queue, msgId, vt) :: IO [PSQL.Only Int])



-- | This function returns the number of readily available messages.
-- 'metrics' returns total queue size, but even then, when vt of
-- these messages is set in the future, the read operation might
-- return nothing (see queueAvailableNow).
queueAvailableLength :: PSQL.Connection -> Queue -> IO Int
queueAvailableLength conn queue = do
  let queueTable = "q_" <> queue
  let tableId = PSQL.QualifiedIdentifier (Just "pgmq") $ T.pack queueTable
  [PSQL.Only len] <- PSQL.query conn [sql| SELECT COUNT(*) FROM ? WHERE vt <= now() |] (PSQL.Only tableId) :: IO [PSQL.Only Int]
  return len

-- | Same as 'queryAvailableLength' but returns the ids of messages.
queueAvailableIds :: PSQL.Connection -> Queue -> IO [MessageId]
queueAvailableIds conn queue = do
  let queueTable = "q_" <> queue
  let tableId = PSQL.QualifiedIdentifier (Just "pgmq") $ T.pack queueTable
  ret <- PSQL.query conn [sql| SELECT msg_id FROM ? WHERE vt <= now() |] (PSQL.Only tableId) :: IO [PSQL.Only MessageId]
  return $ (\(PSQL.Only msgId) -> msgId) <$> ret


{-|
 A utility function: sometimes pgmq throws an error that a table (for
 queue) doesn't exist and we want to ignore it as it's not critical in
 the given function.

 '42P01' means 'undefined_table' in postgres:
 https://www.postgresql.org/docs/current/errcodes-appendix.html
-}
-- withNoDoesNotExistError :: IO a -> IO a