diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,5 @@ +# Revision history for stakhanov + +## 0.0.1.0 -- 2026-01-12 + +* First version. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,28 @@ +BSD 3-Clause License + +Copyright (c) 2026, stakhanov, Michel Boucey + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/ReadMe.md b/ReadMe.md new file mode 100644 --- /dev/null +++ b/ReadMe.md @@ -0,0 +1,4 @@ +# Stakhanov ![CI](https://github.com/MichelBoucey/stakhanov/actions/workflows/haskell-ci.yml/badge.svg) [![Hackage](https://img.shields.io/hackage/v/stakhanov.svg)](https://hackage.haskell.org/package/stakhanov) + +The Haskell library `stakhanov`, built upon [Hasql](https://hackage.haskell.org/package/hasql)'s ecosystem and [Vector](https://hackage.haskell.org/package/vector), implements most of the functions of the [API](https://pgmq.github.io/pgmq/api/sql/functions/) of [PGMQ](https://pgmq.github.io/pgmq/), "a lightweight message queue, like AWS SQS and RSMQ but on Postgres". + diff --git a/src/Database/PostgreSQL/Stakhanov.hs b/src/Database/PostgreSQL/Stakhanov.hs new file mode 100644 --- /dev/null +++ b/src/Database/PostgreSQL/Stakhanov.hs @@ -0,0 +1,299 @@ +-- | This Haskell library, based upon [Hasql](https://hackage.haskell.org/package/hasql)'s ecosystem and [Vector](https://hackage.haskell.org/package/vector), implements the most of [PGMQ](https://github.com/pgmq/pgmq) [API functions](https://pgmq.github.io/pgmq/api/sql/functions/) and should be used qualified. + +module Database.PostgreSQL.Stakhanov + ( + + -- * Queue management + create + , createUnlogged + , declare + , purge + , drop + + -- * Sending Messages + , send + , send' + , batchSend + , batchSend' + + -- * Reading Messages + , read + , readWithPoll + , pop + + -- * Deleting/Archiving Messages + , archive + , delete + , batchArchive + , batchDelete + + -- * Utilities + , batchSetVT + , listQueues + , listQueues' + , details + + -- * Queue details getters + , getQName + , getCreatedAt + , getIsPartitioned + , getIsUnlogged + + ) where +import Control.Monad +import Data.Aeson.Types +import Data.Int +import Data.Maybe +import Data.Text as T hiding (drop) +import Data.Time +import qualified Data.Vector as V +import Database.PostgreSQL.Stakhanov.Internal +import Database.PostgreSQL.Stakhanov.Statements +import Database.PostgreSQL.Stakhanov.Types +import qualified Hasql.Connection as C +import qualified Hasql.Session as S +import Prelude hiding (drop, read) + +-- | Create a new `Queue`. +-- +-- > λ: create "MyQueue" conn +-- > Right (Queue {qName = "MyQueue", qPGConn = "a Hasql connection", qDetails = Nothing, qMetrics = Nothing}) +-- +create + :: T.Text -- ^ The name of the queue to create + -> C.Connection -- ^ The PostgreSQL connection to use + -> IO (Either S.SessionError Queue) +create t c = + S.run (S.statement t createQueue) c >>= + pureMap (\_ -> Queue t (HasqlConn c) Nothing Nothing) + +-- | Create an unlogged new `Queue`. This is useful +-- when write throughput is more important that durability. +-- See [PostgreSQL documentation about unlogged tables](https://www.postgresql.org/docs/current/sql-createtable.html#SQL-CREATETABLE-UNLOGGED). +-- +-- > λ: createUnlogged "MyQueue" conn +-- > Right (Queue {qName = "MyQueue", qPGConn = "a Hasql connection", qDetails = Nothing, qMetrics = Nothing}) +-- +createUnlogged + :: T.Text -- ^ The name of the queue to create + -> C.Connection -- ^ The PostgreSQL connection to use + -> IO (Either S.SessionError Queue) +createUnlogged t c = + S.run (S.statement t createUnloggedQueue) c >>= + pureMap (\_ -> Queue t (HasqlConn c) Nothing Nothing) + +-- | Declare an already existing `Queue`. +-- +-- > λ: declare "MyQueue" conn +-- > Queue {qName = "MyQueue", qPGConn = "a Hasql connection", qDetails = Nothing, qMetrics = Nothing} +-- +declare + :: T.Text -- ^ The name of the queue to declare + -> C.Connection -- ^ The PostgreSQL connection to use + -> Queue +declare t c = Queue t (HasqlConn c) Nothing Nothing + +-- | Permanently deletes all `Messages` in the given `Queue`. +-- Returns the number of `Messages` that were deleted. +purge + :: Queue -- ^ The queue to work with + -> IO (Either S.SessionError Int64) -- ^ Returns the number of messages that were deleted +purge Queue{..} = + S.run (S.statement qName purgeQueue) (unHasqlConn qPGConn) + +-- | Deletes a `Queue` and its archive. +drop + :: Queue -- ^ The queue to work with + -> IO (Either S.SessionError Bool) -- ^ Returns `True` in case of deletion +drop Queue{..} = + S.run (S.statement qName dropQueue) (unHasqlConn qPGConn) + +-- | Send a single `Message` to a `Queue`. +-- Returns the `MsgId` of the just created `Message`. +send + :: Queue -- ^ The queue to work with + -> Value -- ^ A JSON object to send as a message to the queue + -> IO (Either S.SessionError MsgId) -- ^ Returns the message ID of the just created message +send Queue{..} v@(Object _) = + S.run (S.statement (qName,v) sendMessage) (unHasqlConn qPGConn) +send _ _ = fail "The Aeson Value must be an Object" + +-- | Send a single `Message` to a `Queue` with optional metadata (a JSON object named headers) +-- and an optional `Delay`. Returns the `MsgId` of the just created `Message`. +send' + :: Queue -- ^ The queue to work with + -> Value -- ^ A JSON object sent as a message to the queue + -> Maybe Value -- ^ Maybe a JSON object sent as headers/metadata to the queue + -> Maybe Delay -- ^ Maybe a time before which the message becomes visible + -> IO (Either S.SessionError MsgId) -- ^ Returns the message ID of the just created message +send' Queue{..} v@(Object _) mv@(Just (Object _)) md = + S.run (S.statement () $ sendMessage' qName v mv md) (unHasqlConn qPGConn) +send' _ _ _ _ = fail "The Aeson Values must be Objects" + +-- | Send on or more `Messages` to a `Queue`. Returns the `MsgId` of just created `Message`. +batchSend + :: Queue -- ^ The queue to work with + -> V.Vector Value -- ^ A vector of JSON objects sent as messages to the queue + -> IO (Either S.SessionError MsgIds) -- ^ Returns a vector of message IDs of just created messages +batchSend Queue{..} v = + if allJSON v + then S.run (S.statement () $ sendMessages qName v) (unHasqlConn qPGConn) + else fail "All Aeson Values of the Vector must be Objects" + +-- | Send on or more `Messages` to a `Queue` with optional headers (a JSON object of metadata) +-- and an optional `Delay`. Returns `MsgId`s of just created `Messages`. +batchSend' + :: Queue -- ^ The queue to work with + -> V.Vector Value -- ^ A vector of JSON objects sent as messages to the queue + -> Maybe (V.Vector Value) -- ^ Maybe a vector of JSON objects sent as headers/metadata. Its length must be the same as the vector of messages + -> Maybe Delay -- ^ Maybe a time before which messages becomes visible + -> IO (Either S.SessionError MsgIds) -- ^ Returns a vector of message IDs of just created messages +batchSend' Queue{..} vv mvv md = + if allJSON vv && maybe True allJSON mvv + then + if isNothing mvv || isJust mvv && V.length vv == V.length (fromJust mvv) + then S.run (S.statement () $ sendMessages' qName vv mvv md) (unHasqlConn qPGConn) + else fail "The vector of headers must be equal to the vector of messages" + else fail "All Aeson Values of Vectors must be Objects" + +-- | Read one or more `Messages` from a `Queue`. The /Visibility Timeout/ (`VT`) specifies the amount of time +-- in seconds that the `Message` will be invisible to other consumers after reading. +read + :: Queue -- ^ The queue to work with + -> VT -- ^ The Visibility Timeout : the time in seconds that message(s) become invisible after reading + -> Qty -- ^ The number of messages to read from the queue + -> IO (Either S.SessionError (Maybe Messages)) +read Queue{..} v q = + S.run (S.statement (qName,v,q) readMessages) (unHasqlConn qPGConn) >>= pureMap maybeMessages + +-- | Same as `read`. Also provides convenient long-poll functionality. When there are no `Messages` in the `Queue`, +-- the function call will wait for /max_poll_seconds/ in duration before returning. If messages reach the queue +-- during that duration, they will be read and returned immediately. +readWithPoll + :: Queue -- ^ The queue to work with + -> VT -- ^ The Visibility Timeout : the time in seconds that message(s) become invisible after reading + -> Qty -- ^ The number of messages to read from the queue + -> Maybe Seconds -- ^ The max_poll_seconds : the time in seconds to wait for new messages to reach the queue. Defaults to 5 + -> Maybe Milliseconds -- ^ The milliseconds between the internal poll operations. Defaults to 100 + -> IO (Either S.SessionError (Maybe Messages)) +readWithPoll Queue{..} v q mmp mpi = + S.run (S.statement () $ readMessagesWithPoll qName v q mmp mpi) (unHasqlConn qPGConn) >>= pureMap maybeMessages + +-- | Reads one or more `Messages` from a `Queue` and /deletes them upon read/. +-- +-- > λ: pop MyQueue 2 +-- > Right (Just [Message {msgId = 2, readCount = 13, enqueuedAt = 2025-12-09 09:51:50.259464 UTC, visibilityTimeout = 2025-12-15 10:46:41.096843 UTC, message = Object (fromList [("Action",String "hug"),("Quantity",Number 3)]), headers = Nothing},Message {msgId = 3, readCount = 2, enqueuedAt = 2025-12-15 10:44:45.612983 UTC, visibilityTimeout = 2025-12-29 18:04:32.938332 UTC, message = Object (fromList [("Action",String "hug"),("Quantity",Number 5)]), headers = Object (fromList [("Reason",String empathy"")])}]) +-- +pop + :: Queue -- ^ The queue to work with + -> Qty -- ^ The number of messages to pop from the queue (defaults to 1) + -> IO (Either S.SessionError (Maybe Messages)) +pop Queue{..} y = + S.run (S.statement (qName,y) popMessages) (unHasqlConn qPGConn) >>= pureMap maybeMessages + +-- | Removes a single requested `Message` from the specified `Queue` +-- and inserts it into the `Queue`'s archive. +archive + :: Queue -- ^ The queue to work with + -> MsgId -- ^ The message ID of the message to archive + -> IO (Either S.SessionError Bool) +archive Queue{..} i = + S.run (S.statement (qName,i) archiveMessage) (unHasqlConn qPGConn) + +-- | Deletes a batch of requested `Messages` from the specified `Queue` and inserts them into the `Queue`'s archive. +-- Returns a `V.Vector` of `MsgId` that were successfully archived. +batchArchive + :: Queue -- ^ The queue to work with + -> V.Vector MsgId -- ^ A vector of message IDs to archive + -> IO (Either S.SessionError MsgIds) +batchArchive Queue{..} v = + S.run (S.statement () $ archiveMessages qName v) (unHasqlConn qPGConn) + +-- | Deletes a single `Message` from a `Queue`. +delete + :: Queue -- ^ The queue to work with + -> MsgId -- ^ The message ID to delete + -> IO (Either S.SessionError Bool) +delete Queue{..} i = S.run (S.statement (qName,i) deleteMessage) (unHasqlConn qPGConn) + +-- | Delete one or many `Messages` from a `Queue`. +batchDelete + :: Queue -- ^ The queue to work with + -> MsgIds -- ^ The vector of message IDs to delete + -> IO (Either S.SessionError MsgIds) +batchDelete Queue{..} v = + S.run (S.statement () $ deleteMessages qName v) (unHasqlConn qPGConn) + +-- | List all the `Queue`s that currently exist, with a raw Hasql `C.Connection` as parameter. +-- +-- > λ: listQueues conn +-- > Right [Queue {qName = "test", qPGConn = "a Hasql connection", qDetails = Just (Details {createdAt = 2025-12-18 14:33:41.563365 UTC, isPartitioned = False, isUnlogged = False}), qMetrics = Nothing},Queue {qName = "MyQueue", qPGConn = "a Hasql connection", qDetails = Just (Details {createdAt = 2026-01-09 19:05:24.976526 UTC, isPartitioned = False, isUnlogged = False}), qMetrics = Nothing}] +-- +listQueues + :: C.Connection -- ^ A Hasql connection + -> IO (Either S.SessionError Queues) +listQueues c = + S.run (S.statement () getQueuesDetails) c >>= pureMap (toQueue <$>) + where + toQueue r = + Queue + { qName = fst r + , qPGConn = HasqlConn c + , qDetails = Just (tupleToDetails $ snd r) + , qMetrics = Nothing } + +-- | Same as `listQueues`, with a `Queue` as parameter. +listQueues' + :: Queue -- ^ A queue to use its connection to reach PostgreSQL and add this connection to each queue collected + -> IO (Either S.SessionError Queues) +listQueues' Queue{..} = listQueues (unHasqlConn qPGConn) + +-- | Add `Details` information, collected with `listQueues`, to a `Queue` record. +-- +-- > λ: Right list <- listQueues conn +-- > λ: details MyQueue list +-- > Just (Queue {qName = "test", qPGConn = "a Hasql connection", qDetails = Just (Details {createdAt = 2025-12-18 14:33:41.563365 UTC, isPartitioned = False, isUnlogged = False}), qMetrics = Nothing}) +-- +details + :: Queue -- ^ The queue to get details for + -> Queues -- ^ A vector of queues obtained from one of the listQueues functions + -> Maybe Queue -- ^ The queue with details added +details q vq = + case get q vq of + Nothing -> Nothing + Just q' -> Just $ q { qDetails = qDetails q' } + where + get a b = do + let c = V.uncons b + case c of + Nothing -> Nothing + Just t -> + if qName a == qName (fst t) + then Just (fst t) + else get a (snd t) + +getQName :: Queue -> T.Text +getQName Queue{..} = qName + +getCreatedAt :: Queue -> Maybe UTCTime +getCreatedAt (Queue _ _ (Just Details{..}) _) = Just createdAt +getCreatedAt (Queue _ _ Nothing _) = Nothing + +getIsPartitioned :: Queue -> Maybe Bool +getIsPartitioned (Queue _ _ (Just Details{..}) _) = Just isPartitioned +getIsPartitioned (Queue _ _ Nothing _) = Nothing + +getIsUnlogged :: Queue -> Maybe Bool +getIsUnlogged (Queue _ _ (Just Details{..}) _) = Just isUnlogged +getIsUnlogged (Queue _ _ Nothing _) = Nothing + +-- | Sets the /Visibility Timeout/ (`VT`) of one or many `Messages` to a specified time duration +-- in the future. Returns the `Messages` that were updated. +batchSetVT + :: Queue -- ^ The queue to work with + -> MsgIds -- ^ The vector of message IDs to set visibility time + -> Seconds -- ^ Duration from now, in seconds, that the messages VT should be set to + -> IO (Either S.SessionError Messages) +batchSetVT Queue{..} v s = + S.run (S.statement () $ setMessagesVT qName v s) (unHasqlConn qPGConn) >>= pureMap (tupleToMessage <$>) + diff --git a/src/Database/PostgreSQL/Stakhanov/Connection.hs b/src/Database/PostgreSQL/Stakhanov/Connection.hs new file mode 100644 --- /dev/null +++ b/src/Database/PostgreSQL/Stakhanov/Connection.hs @@ -0,0 +1,18 @@ + +-- | This module is just here for a quick start with Stakhanov and Hasql, otherwise use the module [Hasql.Connection](https://hackage.haskell.org/package/hasql/docs/Hasql-Connection.html) and its subsequent modules to create full-blown Hasql connections. + +module Database.PostgreSQL.Stakhanov.Connection where + +import qualified Data.Text as T +import Hasql.Connection +import Hasql.Connection.Setting (connection) +import Hasql.Connection.Setting.Connection (string) + +-- | Get a local PostgreSQL connection to a Docker container of PGMQ, with the default PostgreSQL connection string "__postgres:\/\/postgres:postgres@0.0.0.0:5432/postgres__". +acquireLocalPGConn :: IO (Either ConnectionError Connection) +acquireLocalPGConn = acquire [connection $ string "postgres://postgres:postgres@0.0.0.0:5432/postgres"] + +-- | Get a PostgreSQL connection configured with the given PostgreSQL connection string. +acquirePGConn :: T.Text -> IO (Either ConnectionError Connection) +acquirePGConn t = acquire [connection $ string t] + diff --git a/src/Database/PostgreSQL/Stakhanov/Internal.hs b/src/Database/PostgreSQL/Stakhanov/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Database/PostgreSQL/Stakhanov/Internal.hs @@ -0,0 +1,95 @@ +module Database.PostgreSQL.Stakhanov.Internal where + +import Data.Aeson.Types +import Data.Int +import Data.List (intersperse) +import qualified Data.Monoid as M +import qualified Data.Text as T +import Data.Time +import Data.Vector as V +import Database.PostgreSQL.Stakhanov.Types +import qualified Hasql.Connection as C +import qualified Hasql.DynamicStatements.Snippet as S +import qualified Hasql.Encoders as E + +pureMap + :: (Applicative f1, Functor f2) + => (a -> b) -> f2 a -> f1 (f2 b) +pureMap f e = pure $ f <$> e + +isJSON :: Value -> Bool +isJSON (Object _) = True +isJSON _ = False + +allJSON :: Vector Value -> Bool +allJSON = V.all isJSON + +maybeMessages + :: Vector (Int64, Int32, UTCTime, UTCTime, Value, Maybe Value) + -> Maybe Messages +maybeMessages v = + if V.null v + then Nothing + else Just $ tupleToMessage <$> v + +tupleToDetails :: (UTCTime,Bool,Bool) -> Details +tupleToDetails (e1,e2,e3) = + Details + { createdAt = e1 + , isPartitioned = e2 + , isUnlogged = e3 + } + +tupleToQueueWithMetrics + :: C.Connection + -> (T.Text, Int64, Maybe Int32, Maybe Int32, Int64, UTCTime, Int64) + -> Queue +tupleToQueueWithMetrics c (e1,e2,e3,e4,e5,e6,e7) = + Queue + { qName = e1 + , qPGConn = HasqlConn c + , qDetails = Nothing + , qMetrics = Just $ tupleToMetrics (e2,e3,e4,e5,e6,e7) + } + +tupleToMetrics + :: (Int64, Maybe Int32, Maybe Int32, Int64, UTCTime, Int64) + -> Metrics +tupleToMetrics (e1,e2,e3,e4,e5,e6) = + Metrics + { queueLength = e1 + , newestMsgAge = e2 + , oldestMsgAge = e3 + , totalMessages = e4 + , scrapeTime = e5 + , queueVisibleLength = e6 } + +tupleToMessage + :: (Int64, Int32, UTCTime, UTCTime, Value, Maybe Value) + -> Message +tupleToMessage (e1,e2,e3,e4,e5,e6) = + Message + { msgId = e1 + , readCount = e2 + , enqueuedAt = e3 + , visibilityTimeout = e4 + , message = e5 + , headers = e6 } + +maybeHeaders :: Maybe Value -> S.Snippet +maybeHeaders (Just v) = "," <> S.encoderAndParam (E.nonNullable E.json) v <> "::jsonb" +maybeHeaders Nothing = mempty + +maybeDelay :: Maybe Delay -> S.Snippet +maybeDelay (Just (InSeconds s)) = "," <> S.encoderAndParam (E.nonNullable E.int4) s +maybeDelay (Just (WithTimestamp t)) = "," <> S.encoderAndParam (E.nonNullable E.timestamptz) t +maybeDelay Nothing = mempty + +jsonbArrayEncoder :: V.Vector Value -> S.Snippet +jsonbArrayEncoder v = + "ARRAY[" <> M.mconcat (intersperse (S.sql ",") $ V.toList $ S.encoderAndParam (E.nonNullable E.json) <$> v) <> "]::jsonb[]" + +bigintArrayEncoder :: V.Vector Int64 -> S.Snippet +bigintArrayEncoder v = + "ARRAY[" <> M.mconcat (intersperse (S.sql ",") $ V.toList $ S.encoderAndParam (E.nonNullable E.int8) <$> v) <> "]::bigint[]" + diff --git a/src/Database/PostgreSQL/Stakhanov/Metrics.hs b/src/Database/PostgreSQL/Stakhanov/Metrics.hs new file mode 100644 --- /dev/null +++ b/src/Database/PostgreSQL/Stakhanov/Metrics.hs @@ -0,0 +1,76 @@ +module Database.PostgreSQL.Stakhanov.Metrics + ( + + -- * Get queue(s) metrics + metrics + , allMetrics + + -- * Queue metrics getters + , getQueueLength + , getNewestMsgAge + , getOldestMsgAge + , getScrapeTime + , getTotalMessages + , getQueueVisibleLength + + ) where + +import Data.Int +import Data.Time +import qualified Data.Vector as V +import Database.PostgreSQL.Stakhanov.Internal +import Database.PostgreSQL.Stakhanov.Statements +import Database.PostgreSQL.Stakhanov.Types +import qualified Hasql.Session as S + +-- | Get `Queue`'s `Metrics`. +-- +-- > λ: metrics MyQueue +-- > Right (Queue {qName = "MyQueue", qPGConn = "a Hasql connection", qDetails = Nothing, qMetrics = Just (Metrics {queueLength = 24, newestMsgAge = Just 447278, oldestMsgAge = Just 2192954, totalMessages = 27, scrapeTime = 2026-01-09 19:53:59.503568 UTC, queueVisibleLength = 24})}) +-- +metrics + :: Queue -- ^ The name of the queue + -> IO (Either S.SessionError Queue) -- ^ The queue with metrics added +metrics q@Queue{..} = + S.run (S.statement qName getMetrics) (unHasqlConn qPGConn) >>= pureMap addMetrics + where + addMetrics m = q { qMetrics = Just $ tupleToMetrics m } + +-- | Get `Metrics` of all created `Queue`s. +allMetrics + :: Queue -- ^ The connection to PostgreSQL + -> IO (Either S.SessionError (V.Vector Queue)) +allMetrics Queue{..} = + let c = unHasqlConn qPGConn + in S.run (S.statement () getAllMetrics) c >>= pureMap (tupleToQueueWithMetrics c <$>) + +-- | Number of messages currently in the queue. +getQueueLength :: Queue -> Maybe Int64 +getQueueLength (Queue _ _ _ (Just Metrics{..})) = Just queueLength +getQueueLength (Queue _ _ _ Nothing) = Nothing + +-- | Age of the newest message in the queue, in seconds. +getNewestMsgAge :: Queue -> Maybe Seconds +getNewestMsgAge (Queue _ _ _ (Just Metrics{..})) = newestMsgAge +getNewestMsgAge (Queue _ _ _ Nothing) = Nothing + +-- | Age of the oldest message in the queue, in seconds. +getOldestMsgAge :: Queue -> Maybe Seconds +getOldestMsgAge (Queue _ _ _ (Just Metrics{..})) = oldestMsgAge +getOldestMsgAge (Queue _ _ _ Nothing) = Nothing + +-- | Total number of messages that have passed through the queue over all time. +getTotalMessages :: Queue -> Maybe Int64 +getTotalMessages (Queue _ _ _ (Just Metrics{..})) = Just totalMessages +getTotalMessages (Queue _ _ _ Nothing) = Nothing + +-- | The current timestamp. +getScrapeTime :: Queue -> Maybe UTCTime +getScrapeTime (Queue _ _ _ (Just Metrics{..})) = Just scrapeTime +getScrapeTime (Queue _ _ _ Nothing) = Nothing + +-- | Number of messages currently visible (vt <= now). +getQueueVisibleLength :: Queue -> Maybe Int64 +getQueueVisibleLength (Queue _ _ _ (Just Metrics{..})) = Just queueVisibleLength +getQueueVisibleLength (Queue _ _ _ Nothing) = Nothing + diff --git a/src/Database/PostgreSQL/Stakhanov/Statements.hs b/src/Database/PostgreSQL/Stakhanov/Statements.hs new file mode 100644 --- /dev/null +++ b/src/Database/PostgreSQL/Stakhanov/Statements.hs @@ -0,0 +1,182 @@ +module Database.PostgreSQL.Stakhanov.Statements where + +import Contravariant.Extras.Contrazip (contrazip2, contrazip3) +import Data.Aeson +import Data.ByteString +import Data.Int +import qualified Data.List as L +import Data.Maybe (fromMaybe) +import qualified Data.Text as T +import Data.Time +import qualified Data.Vector as V +import Database.PostgreSQL.Stakhanov.Internal +import Database.PostgreSQL.Stakhanov.Types +import qualified Hasql.Decoders as D +import qualified Hasql.DynamicStatements.Snippet as S +import Hasql.DynamicStatements.Statement +import qualified Hasql.Encoders as E +import Hasql.Statement +import qualified Hasql.TH as TH +import Prelude hiding (pi) + +createQueue :: Statement T.Text () +createQueue = [TH.resultlessStatement|select from pgmq.create($1::text)|] + +createUnloggedQueue :: Statement T.Text () +createUnloggedQueue = [TH.resultlessStatement|select from pgmq.create_unlogged($1::text)|] + +getQueuesDetails :: Statement () (V.Vector (T.Text, (UTCTime, Bool, Bool))) +getQueuesDetails = + Statement sql E.noParams decoder True + where + sql = "select queue_name,created_at,is_partitioned,is_unlogged from pgmq.list_queues()" + decoder = + D.rowVector $ + (,) <$> + D.column (D.nonNullable D.text) <*> + ((,,) <$> + D.column (D.nonNullable D.timestamptz) <*> + D.column (D.nonNullable D.bool) <*> + D.column (D.nonNullable D.bool)) + +purgeQueue :: Statement T.Text Int64 +purgeQueue = [TH.singletonStatement|select pgmq.purge_queue($1::text)::int8|] + +dropQueue :: Statement T.Text Bool +dropQueue = [TH.singletonStatement|select pgmq.drop_queue($1::text)::bool|] + +sendMessage :: Statement (T.Text,Value) Int64 +sendMessage = + Statement sql encoder decoder True + where + sql = "select * from pgmq.send($1::text,$2::jsonb)" + encoder = + contrazip2 + (E.param (E.nonNullable E.text)) + (E.param (E.nonNullable E.jsonb)) + decoder = D.singleRow $ D.column $ D.nonNullable D.int8 + +sendMessage' :: T.Text -> Value -> Maybe Value -> Maybe Delay -> Statement () Int64 +sendMessage' q v mv mi = + let snippet = + "select * from pgmq.send(" <> S.param q <> "," + <> S.param v <> maybeHeaders mv <> maybeDelay mi <> ")" + decoder = D.singleRow $ D.column $ D.nonNullable D.int8 + in dynamicallyParameterized snippet decoder True + +sendMessages :: T.Text -> (V.Vector Value) -> Statement () (V.Vector Int64) +sendMessages q msgs = + let snippet = "select * from pgmq.send_batch(" <> S.param q <> "," <> jsonbArrayEncoder msgs <> ")" + decoder = D.rowVector $ D.column $ D.nonNullable D.int8 + in dynamicallyParameterized snippet decoder True + +sendMessages' :: T.Text -> (V.Vector Value) -> Maybe (V.Vector Value) -> Maybe Delay -> Statement () (V.Vector Int64) +sendMessages' q vv mvv md = + let snippet = + "select * from pgmq.send_batch(" <> S.param q <> "," <> jsonbArrayEncoder vv + <> (fromMaybe mempty $ (mappend "," . jsonbArrayEncoder) <$> mvv) <> maybeDelay md <> ")" + decoder = D.rowVector $ D.column $ D.nonNullable D.int8 + in dynamicallyParameterized snippet decoder True + +readMessagesWithPoll :: T.Text -> Int32 -> Int32 -> Maybe Int32 -> Maybe Int32 -> Statement () (V.Vector (Int64, Int32, UTCTime, UTCTime, Value, Maybe Value)) +readMessagesWithPoll q vt qty mmp mpi = + let mp = maybe 5 id mmp + pi = maybe 100 id mpi + snippet = "select " <> S.sql columnsMessage <> " from pgmq.read_with_poll(" <> + mconcat (L.intersperse "," [S.param q, S.param vt, S.param qty, S.param mp, S.param pi]) <> ")" + in dynamicallyParameterized snippet tupleMessageDecoder True + +readMessages :: Statement (T.Text,Int32,Int32) (V.Vector (Int64, Int32, UTCTime, UTCTime, Value, Maybe Value)) +readMessages = + Statement sql encoder tupleMessageDecoder True + where + sql = "select " <> columnsMessage <> " from pgmq.read($1,$2,$3)" + encoder = + contrazip3 + (E.param (E.nonNullable E.text)) + (E.param (E.nonNullable E.int4)) + (E.param (E.nonNullable E.int4)) + +popMessages :: Statement (T.Text,Int32) (V.Vector (Int64, Int32, UTCTime, UTCTime, Value, Maybe Value)) +popMessages = + Statement sql encoder tupleMessageDecoder True + where + sql = "select " <> columnsMessage <> " from pgmq.pop($1,$2)" + encoder = + contrazip2 + (E.param (E.nonNullable E.text)) + (E.param (E.nonNullable E.int4)) + +getMetrics :: Statement T.Text (Int64, Maybe Int32, Maybe Int32, Int64, UTCTime, Int64) +getMetrics = + Statement sql encoder decoder True + where + sql = "select " <> columnsMetrics <> " from pgmq.metrics($1)" + encoder = E.param (E.nonNullable E.text) + decoder = + D.singleRow $ + (,,,,,) <$> + D.column (D.nonNullable D.int8) <*> + D.column (D.nullable D.int4) <*> + D.column (D.nullable D.int4) <*> + D.column (D.nonNullable D.int8) <*> + D.column (D.nonNullable D.timestamptz) <*> + D.column (D.nonNullable D.int8) + +getAllMetrics :: Statement () (V.Vector (T.Text, Int64, Maybe Int32, Maybe Int32, Int64, UTCTime, Int64)) +getAllMetrics = + Statement sql E.noParams decoder True + where + sql = "select queue_name," <> columnsMetrics <> " from pgmq.metrics_all()" + decoder = + D.rowVector $ + (,,,,,,) <$> + D.column (D.nonNullable D.text) <*> + D.column (D.nonNullable D.int8) <*> + D.column (D.nullable D.int4) <*> + D.column (D.nullable D.int4) <*> + D.column (D.nonNullable D.int8) <*> + D.column (D.nonNullable D.timestamptz) <*> + D.column (D.nonNullable D.int8) + +archiveMessage :: Statement (T.Text,Int64) Bool +archiveMessage = [TH.singletonStatement|select pgmq.archive($1::text,$2::int8)::bool|] + +archiveMessages :: T.Text -> (V.Vector Int64) -> Statement () (V.Vector Int64) +archiveMessages q v = + let snippet = "select * from pgmq.archive(" <> S.param q <> "," <> bigintArrayEncoder v <> ")" + decoder = D.rowVector (D.column (D.nonNullable D.int8)) + in dynamicallyParameterized snippet decoder True + +deleteMessage :: Statement (T.Text,Int64) Bool +deleteMessage = [TH.singletonStatement|select pgmq.delete($1::text,$2::int8)::bool|] + +deleteMessages :: T.Text -> (V.Vector Int64) -> Statement () (V.Vector Int64) +deleteMessages q v = + let snippet = "select * from pgmq.delete(" <> S.param q <> "," <> bigintArrayEncoder v <> ")" + decoder = D.rowVector (D.column (D.nonNullable D.int8)) + in dynamicallyParameterized snippet decoder True + +setMessagesVT :: T.Text -> V.Vector Int64 -> Int32 -> Statement () (V.Vector (Int64, Int32, UTCTime, UTCTime, Value, Maybe Value)) +setMessagesVT q v s = + let snippet = "select * from pgmq.set_vt(" <> S.param q <> "," + <> bigintArrayEncoder v <> "," <> S.param s <> ")" + in dynamicallyParameterized snippet tupleMessageDecoder True + +tupleMessageDecoder :: D.Result (V.Vector (Int64, Int32, UTCTime, UTCTime, Value, Maybe Value)) +tupleMessageDecoder = + D.rowVector $ + (,,,,,) <$> + D.column (D.nonNullable D.int8) <*> + D.column (D.nonNullable D.int4) <*> + D.column (D.nonNullable D.timestamptz) <*> + D.column (D.nonNullable D.timestamptz) <*> + D.column (D.nonNullable D.jsonb) <*> + D.column (D.nullable D.jsonb) + +columnsMessage :: ByteString +columnsMessage = "msg_id,read_ct,enqueued_at,vt,message,headers" + +columnsMetrics :: ByteString +columnsMetrics = "queue_length,newest_msg_age_sec,oldest_msg_age_sec,total_messages,scrape_time,queue_visible_length" + diff --git a/src/Database/PostgreSQL/Stakhanov/Types.hs b/src/Database/PostgreSQL/Stakhanov/Types.hs new file mode 100644 --- /dev/null +++ b/src/Database/PostgreSQL/Stakhanov/Types.hs @@ -0,0 +1,67 @@ +module Database.PostgreSQL.Stakhanov.Types where + +import Data.Aeson.Types +import Data.Int +import qualified Data.Text as T +import Data.Time +import Data.Vector +import qualified Hasql.Connection as C + +type VT = Int32 +type Qty = Int32 +type Seconds = Int32 +type Milliseconds = Int32 +type MsgId = Int64 +type MsgIds = Vector MsgId +type Messages = Vector Message +type Queues = Vector Queue + +newtype HasqlConn = HasqlConn { unHasqlConn :: C.Connection } + +instance Show HasqlConn where + show (HasqlConn _) = show ("a Hasql connection" :: String) + +data Queue = + Queue + { qName :: T.Text + , qPGConn :: !HasqlConn + , qDetails :: Maybe Details + , qMetrics :: Maybe Metrics + } deriving (Show) + +instance Eq Queue where + Queue n _ _ _ == Queue n' _ _ _ = n == n' + +data Details = + Details + { + createdAt :: UTCTime + , isPartitioned :: Bool + , isUnlogged :: Bool + } deriving (Show) + +data Message = + Message + { msgId :: MsgId + , readCount :: Int32 + , enqueuedAt :: UTCTime + , visibilityTimeout :: UTCTime + , message :: !Value + , headers :: !(Maybe Value) + } deriving (Show) + +instance Eq Message where + Message i _ _ _ _ _ == Message i' _ _ _ _ _ = i == i' + +data Metrics = + Metrics + { queueLength :: Int64 + , newestMsgAge :: Maybe Seconds + , oldestMsgAge :: Maybe Seconds + , totalMessages :: Int64 + , scrapeTime :: UTCTime + , queueVisibleLength :: Int64 + } deriving (Show) + +data Delay = InSeconds Int32 | WithTimestamp UTCTime + diff --git a/stakhanov.cabal b/stakhanov.cabal new file mode 100644 --- /dev/null +++ b/stakhanov.cabal @@ -0,0 +1,72 @@ +cabal-version: 3.0 +name: stakhanov +version: 0.0.1.0 +synopsis: A Haskell PGMQ client +description: + A fast Haskell PGMQ client for busy workers + +homepage: https://github.com/MichelBoucey/stakhanov +license: BSD-3-Clause +license-file: LICENSE +author: Michel Boucey +maintainer: michel.boucey@gmail.com +copyright: (c) 2026 - Michel Boucey +category: Database, PostgreSQL +build-type: Simple +extra-doc-files: CHANGELOG.md +extra-source-files: ReadMe.md +tested-with: + GHC ==9.6.7 || ==9.8.4 || ==9.10.3 || ==9.12.2 || ==9.14.1 + +source-repository head + type: git + location: https://github.com/MichelBoucey/stakhanov.git + +common warnings + ghc-options: -Wall + +library + import: warnings + exposed-modules: + Database.PostgreSQL.Stakhanov + Database.PostgreSQL.Stakhanov.Connection + Database.PostgreSQL.Stakhanov.Metrics + Database.PostgreSQL.Stakhanov.Types + + other-modules: + Database.PostgreSQL.Stakhanov.Internal + Database.PostgreSQL.Stakhanov.Statements + + default-extensions: + LambdaCase + OverloadedStrings + QuasiQuotes + RecordWildCards + + build-depends: + , aeson >=2.2.3 && <2.4 + , base >=4.8 && <5 + , bytestring >=0.11.5.4 && <0.13 + , contravariant-extras >=0.3.5 && <0.3.6 + , hasql >=1.9 && <1.10 + , hasql-dynamic-statements >=0.3 && <0.4 + , hasql-th >=0.4 && <0.5 + , text >=1.2.3 && <2.2 + , time >=1.9.3 && <1.16 + , vector >=0.12.2 && <0.14 + + hs-source-dirs: src + default-language: Haskell2010 + +test-suite stakhanov + type: exitcode-stdio-1.0 + hs-source-dirs: tests + main-is: hspec.hs + build-depends: + , aeson >=2.2.3 && <2.4 + , base >=4.8 && <5 + , hspec >=2.11 && <3 + , stakhanov + , vector >=0.12.2 && <0.14 + + default-language: Haskell2010 diff --git a/tests/hspec.hs b/tests/hspec.hs new file mode 100644 --- /dev/null +++ b/tests/hspec.hs @@ -0,0 +1,167 @@ +{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-} +{-# LANGUAGE OverloadedStrings #-} + +import Data.Aeson +import qualified Data.Aeson.KeyMap as K +import Data.Int +import qualified Data.Vector as V +import qualified Database.PostgreSQL.Stakhanov as S +import Database.PostgreSQL.Stakhanov.Connection +import Database.PostgreSQL.Stakhanov.Metrics +import Database.PostgreSQL.Stakhanov.Types +import Test.Hspec + +main :: IO () +main = hspec $ do + + describe "Create a Hspec test queue" $ + it "Return the record of the created queue" $ do + Right c <- acquireLocalPGConn + Right q <- S.create "HspecTestQueue" c + q `shouldBe` (q :: Queue) + + describe "Send a message" $ + it "Return the ID of message created" $ do + Right c <- acquireLocalPGConn + let q = S.declare "HspecTestQueue" c + j = Object (K.fromList [("Item", String "Banana"),("Qty", Number 3)]) + S.send q j `shouldReturn` Right 1 + + describe "Send a message with options (send')" $ + it "Return the ID of message created" $ do + Right c <- acquireLocalPGConn + let q = S.declare "HspecTestQueue" c + j1 = Object (K.fromList [("Item", String "Apple"),("Qty", Number 6)]) + j2 = Object (K.fromList [("Metadata", Object (K.fromList [("Checksum", Bool True), ("Lenght", Number 37)]))]) + S.send' q j1 (Just j2) Nothing `shouldReturn` Right 2 + + describe "Send a batch of messages" $ + it "Return the IDs of messages created" $ do + Right c <- acquireLocalPGConn + let q = S.declare "HspecTestQueue" c + o1 = Object (K.fromList [("Item", String "Pineapple"),("Qty", Number 4)]) + o2 = Object (K.fromList [("Item", String "Tomato"),("Qty", Number 9)]) + o3 = Object (K.fromList [("Item", String "strawnberry"),("Qty", Number 23)]) + v = V.fromList[o1,o2,o3] + S.batchSend q v `shouldReturn` Right (V.fromList[3,4,5]) + + describe "Send a batch of messages with a delay (batchSend')" $ + it "Return the IDs of messages created" $ do + Right c <- acquireLocalPGConn + let q = S.declare "HspecTestQueue" c + o1 = Object (K.fromList [("Item", String "Potato"),("Qty", Number 16)]) + o2 = Object (K.fromList [("Item", String "Carrot"),("Qty", Number 21)]) + o3 = Object (K.fromList [("Item", String "strawnberry"),("Qty", Number 23)]) + v = V.fromList[o1,o2,o3] + S.batchSend' q v Nothing (Just (InSeconds 1)) `shouldReturn` Right (V.fromList[6,7,8]) + + describe "Send a batch of messages with headers (batchSend')" $ + it "Return the IDs of messages created" $ do + Right c <- acquireLocalPGConn + let q = S.declare "HspecTestQueue" c + o1 = Object (K.fromList [("Item", String "Potato"),("Qty", Number 8)]) + o2 = Object (K.fromList [("Item", String "Carrot"),("Qty", Number 120)]) + o3 = Object (K.fromList [("Item", String "strawnberry"),("Qty", Number 1230)]) + o4 = Object (K.fromList [("Metadata", Object (K.fromList [("Checksum", Bool True), ("Lenght", Number 37)]))]) + o5 = Object (K.fromList [("Metadata", Object (K.fromList [("Checksum", Bool True), ("Lenght", Number 40)]))]) + o6 = Object (K.fromList [("Metadata", Object (K.fromList [("Checksum", Bool True), ("Lenght", Number 500)]))]) + v1 = V.fromList[o1,o2,o3] + v2 = V.fromList[o4,o5,o6] + S.batchSend' q v1 (Just v2) Nothing `shouldReturn` Right (V.fromList[9,10,11]) + + describe "Get the current length of the queue, extracted from its metrics" $ + it "Return the number of messages currently in the queue" $ do + Right c <- acquireLocalPGConn + let q = S.declare "HspecTestQueue" c + Right q' <- metrics q + pure (getQueueLength q') `shouldReturn` Just 11 + + describe "Get metrics of all queues" $ + it "Return a vector of queues with their metrics embbeded" $ do + Right c <- acquireLocalPGConn + let q = S.declare "HspecTestQueue" c + Right vq <- allMetrics q + vq `shouldBe` (vq :: V.Vector Queue) + + describe "Read a message" $ + it "Maybe returns a vector of messages" $ do + Right c <- acquireLocalPGConn + let q = S.declare "HspecTestQueue" c + Right vm <- S.read q 30 1 + vm `shouldBe` (vm :: Maybe Messages) + + describe "Archive a message" $ + it "Return True" $ do + Right c <- acquireLocalPGConn + let q = S.declare "HspecTestQueue" c + S.archive q 1 `shouldReturn` Right True + + describe "Archive messages" $ + it "Return the IDs of archived messages" $ do + Right c <- acquireLocalPGConn + let q = S.declare "HspecTestQueue" c + v = V.fromList[3,4,5] + S.batchArchive q v `shouldReturn` Right v + + describe "Delete a message" $ + it "Return True" $ do + Right c <- acquireLocalPGConn + let q = S.declare "HspecTestQueue" c + S.delete q 2 `shouldReturn` Right True + + describe "Set Visibility Timeout of the given message IDs (batchSetVT)" $ + it "Return the vector of updated messages" $ do + Right c <- acquireLocalPGConn + let q = S.declare "HspecTestQueue" c + v = V.fromList[6,7,8] + Right vm <- S.batchSetVT q v 30 + vm `shouldBe` (vm :: V.Vector Message) + + describe "Delete messages" $ + it "Return the vector of IDs of deleted messages" $ do + Right c <- acquireLocalPGConn + let q = S.declare "HspecTestQueue" c + vi = V.fromList[6,7,8] + S.batchDelete q vi `shouldReturn` Right vi + + describe "Pop messages from the queue" $ + it "Maybe return messages and delete them from the queue" $ do + Right c <- acquireLocalPGConn + let q = S.declare "HspecTestQueue" c + Right p <- S.pop q 2 + p `shouldBe` (p :: Maybe Messages) + + describe "Purge the Hspec test queue" $ + it "Return the number of deleted messages" $ do + Right c <- acquireLocalPGConn + let q = S.declare "HspecTestQueue" c + Right n <- S.purge q + n `shouldBe` (n :: Int64) + + describe "Get the list of queues with just a connection" $ + it "Return a vector of queues)" $ do + Right c <- acquireLocalPGConn + Right v <- S.listQueues c + v `shouldBe` (v :: V.Vector Queue) + + describe "Get the list of queues with the connection embbeded in a queue" $ + it "Return a vector of queues" $ do + Right c <- acquireLocalPGConn + let q = S.declare "HspecTestQueue" c + Right v <- S.listQueues' q + v `shouldBe` (v :: V.Vector Queue) + + describe "Get a detail of the queue details" $ + it "Return Just False" $ do + Right c <- acquireLocalPGConn + let q = S.declare "HspecTestQueue" c + Right v <- S.listQueues' q + let Just q' = S.details q v + pure (S.getIsUnlogged q') `shouldReturn` Just False + + describe "Delete the queue" $ + it "Return True" $ do + Right c <- acquireLocalPGConn + let q = S.declare "HspecTestQueue" c + S.drop q `shouldReturn` Right True +