diff --git a/postgresql-simple-queue.cabal b/postgresql-simple-queue.cabal --- a/postgresql-simple-queue.cabal +++ b/postgresql-simple-queue.cabal @@ -1,5 +1,5 @@ name: postgresql-simple-queue -version: 0.1.0.1 +version: 0.2.0.0 synopsis: A PostgreSQL backed queue description: This module utilize PostgreSQL to implement a durable queue for efficently processing arbitrary payloads which can be represented as JSON. @@ -10,13 +10,13 @@ > createAccount userRecord = do > 'runDBTSerializable' $ do > createUserDB userRecord - > 'enqueueDB' $ makeVerificationEmail userRecord + > 'enqueueDB' "queue_schema" $ makeVerificationEmail userRecord . In another thread or process, the consumer would drain the queue. . > forever $ do > -- Attempt get a payload or block until one is available - > payload <- lock conn + > payload <- lock "queue_schema" conn > > -- Perform application specifc parsing of the payload value > case fromJSON $ pValue payload of @@ -24,7 +24,10 @@ > Error err -> logErr err > > -- Remove the payload from future processing - > dequeue conn $ pId payload + > dequeue "queue_schema" conn $ pId payload + > + > To support multiple queues in the same database, the API expects a schema name string + > to determine which queue tables to use. homepage: https://github.com/jfischoff/postgresql-queue#readme license: BSD3 license-file: LICENSE @@ -78,7 +81,7 @@ , lifted-base default-language: Haskell2010 -test-suite db-testing-example-test +test-suite unit-tests type: exitcode-stdio-1.0 hs-source-dirs: test main-is: Spec.hs diff --git a/src/Database/PostgreSQL/Simple/Queue.hs b/src/Database/PostgreSQL/Simple/Queue.hs --- a/src/Database/PostgreSQL/Simple/Queue.hs +++ b/src/Database/PostgreSQL/Simple/Queue.hs @@ -46,6 +46,7 @@ {-# LANGUAGE QuasiQuotes #-} {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE OverloadedStrings #-} module Database.PostgreSQL.Simple.Queue ( -- * Types PayloadId (..) @@ -75,7 +76,7 @@ import Data.Maybe import Data.Text (Text) import Data.Time -import Data.UUID +import Data.UUID (UUID) import Database.PostgreSQL.Simple (Connection, Only (..)) import qualified Database.PostgreSQL.Simple as Simple import Database.PostgreSQL.Simple.FromField @@ -87,6 +88,8 @@ import Database.PostgreSQL.Simple.Transaction import Database.PostgreSQL.Transact import System.Random +import Data.Monoid +import Data.String ------------------------------------------------------------------------------- --- Types ------------------------------------------------------------------------------- @@ -144,6 +147,12 @@ ------------------------------------------------------------------------------- --- DB API ------------------------------------------------------------------------------- +withSchema :: String -> Simple.Query -> Simple.Query +withSchema schemaName q = "SET search_path TO " <> fromString schemaName <> "; " <> q + +notifyName :: IsString s => String -> s +notifyName schemaName = fromString $ schemaName <> "_enqueue" + retryOnUniqueViolation :: MonadCatch m => m a -> m a retryOnUniqueViolation act = try act >>= \case Right x -> return x @@ -165,14 +174,16 @@ 'enqueueDB' $ makeVerificationEmail userRecord @ -} -enqueueDB :: Value -> DB PayloadId -enqueueDB value = retryOnUniqueViolation $ do +enqueueDB :: String -> Value -> DB PayloadId +enqueueDB schemaName value = retryOnUniqueViolation $ do pid <- liftIO randomIO - execute [sql| INSERT INTO payloads (id, value) - VALUES (?, ?); - NOTIFY enqueue; - |] - (pid, value) + execute (withSchema schemaName $ [sql| + INSERT INTO payloads (id, value) + VALUES (?, ?); + NOTIFY |] <> " " <> notifyName schemaName <> + ";" + ) + (pid, value) return $ PayloadId pid {-| Return a the oldest 'Payload' in the 'Enqueued' state, or 'Nothing' @@ -181,8 +192,8 @@ with other transactions. See 'tryLock' the IO API version, or for a blocking version utilizing PostgreSQL's NOTIFY and LISTEN, see 'lock' -} -tryLockDB :: DB (Maybe Payload) -tryLockDB = listToMaybe <$> query_ +tryLockDB :: String -> DB (Maybe Payload) +tryLockDB schemaName = fmap listToMaybe $ query_ $ withSchema schemaName [sql| UPDATE payloads SET state='locked' WHERE id in @@ -200,26 +211,26 @@ shutdown. In general the IO API version, 'unlock', is probably more useful. The DB version is provided for completeness. -} -unlockDB :: PayloadId -> DB () -unlockDB payloadId = void $ execute +unlockDB :: String -> PayloadId -> DB () +unlockDB schemaName payloadId = void $ execute (withSchema schemaName [sql| UPDATE payloads SET state='enqueued' WHERE id=? AND state='locked' - |] + |]) payloadId -- | Transition a 'Payload' to the 'Dequeued' state. -dequeueDB :: PayloadId -> DB () -dequeueDB payloadId = void $ execute +dequeueDB :: String -> PayloadId -> DB () +dequeueDB schemaName payloadId = void $ execute (withSchema schemaName [sql| UPDATE payloads SET state='dequeued' WHERE id=? - |] + |]) payloadId -- | Get the number of rows in the 'Enqueued' state. -getCountDB :: DB Int64 -getCountDB = fromOnly . head <$> query_ +getCountDB :: String -> DB Int64 +getCountDB schemaName = fmap (fromOnly . head) $ query_ $ withSchema schemaName [sql| SELECT count(*) FROM payloads WHERE state='enqueued' @@ -230,36 +241,36 @@ {-| Enqueue a new JSON value into the queue. See 'enqueueDB' for a version which can be composed with other queries in a single transaction. -} -enqueue :: Connection -> Value -> IO PayloadId -enqueue conn value = runDBT (enqueueDB value) ReadCommitted conn +enqueue :: String -> Connection -> Value -> IO PayloadId +enqueue schemaName conn value = runDBT (enqueueDB schemaName value) ReadCommitted conn {-| Return a the oldest 'Payload' in the 'Enqueued' state or 'Nothing' if there are no payloads. For a blocking version utilizing PostgreSQL's NOTIFY and LISTEN, see 'lock'. This functions runs 'tryLockDB' as a 'Serializable' transaction. -} -tryLock :: Connection -> IO (Maybe Payload) -tryLock = runDBTSerializable tryLockDB +tryLock :: String -> Connection -> IO (Maybe Payload) +tryLock schemaName conn = runDBTSerializable (tryLockDB schemaName) conn -notifyPayload :: Connection -> IO () -notifyPayload conn = do +notifyPayload :: String -> Connection -> IO () +notifyPayload schemaName conn = do Notification {..} <- getNotification conn - unless (notificationChannel == "enqueue") $ notifyPayload conn + unless (notificationChannel == notifyName schemaName) $ notifyPayload schemaName conn {-| Return the oldest 'Payload' in the 'Enqueued' state or block until a payload arrives. This function utilizes PostgreSQL's LISTEN and NOTIFY functionality to avoid excessively polling of the DB while waiting for new payloads, without scarficing promptness. -} -lock :: Connection -> IO Payload -lock conn = bracket_ - (Simple.execute_ conn "LISTEN enqueue") - (Simple.execute_ conn "UNLISTEN enqueue") +lock :: String -> Connection -> IO Payload +lock schemaName conn = bracket_ + (Simple.execute_ conn $ "LISTEN " <> notifyName schemaName) + (Simple.execute_ conn $ "UNLISTEN " <> notifyName schemaName) $ fix $ \continue -> do - m <- tryLock conn + m <- tryLock schemaName conn case m of Nothing -> do - notifyPayload conn + notifyPayload schemaName conn continue Just x -> return x @@ -267,16 +278,16 @@ Useful for responding to asynchronous exceptions during a unexpected shutdown. For a DB API version see 'unlockDB' -} -unlock :: Connection -> PayloadId -> IO () -unlock conn x = runDBTSerializable (unlockDB x) conn +unlock :: String -> Connection -> PayloadId -> IO () +unlock schemaName conn x = runDBTSerializable (unlockDB schemaName x) conn -- | Transition a 'Payload' to the 'Dequeued' state. his functions runs -- 'dequeueDB' as a 'Serializable' transaction. -dequeue :: Connection -> PayloadId -> IO () -dequeue conn x = runDBTSerializable (dequeueDB x) conn +dequeue :: String -> Connection -> PayloadId -> IO () +dequeue schemaName conn x = runDBTSerializable (dequeueDB schemaName x) conn {-| Get the number of rows in the 'Enqueued' state. This function runs 'getCountDB' in a 'ReadCommitted' transaction. -} -getCount :: Connection -> IO Int64 -getCount = runDBT getCountDB ReadCommitted +getCount :: String -> Connection -> IO Int64 +getCount schemaName = runDBT (getCountDB schemaName) ReadCommitted diff --git a/src/Database/PostgreSQL/Simple/Queue/Main.hs b/src/Database/PostgreSQL/Simple/Queue/Main.hs --- a/src/Database/PostgreSQL/Simple/Queue/Main.hs +++ b/src/Database/PostgreSQL/Simple/Queue/Main.hs @@ -73,38 +73,43 @@ data PartialOptions = PartialOptions { threadCount :: Last Int , dbOptions :: PostgreSQL.PartialOptions + , schemaName :: Last String } deriving (Show, Eq, Typeable) instance Monoid PartialOptions where - mempty = PartialOptions mempty mempty + mempty = PartialOptions mempty mempty mempty mappend x y = PartialOptions (threadCount x <> threadCount y) (dbOptions x <> dbOptions y) + (schemaName x <> schemaName y) -- | The default 'threadCount' is 1. -- The default db options are specified in -- 'Database.PostgreSQL.Simple.Options.PartialOptions' instance Default PartialOptions where - def = PartialOptions (return 1) def + def = PartialOptions (return 1) def $ return "queue" instance ParseRecord PartialOptions where parseRecord = PartialOptions <$> parseFields Nothing (Just "thread-count") (Just 't') <*> parseRecord + <*> parseFields Nothing (Just "schema-name") (Just 's') -- | Final Options used by 'run'. data Options = Options { oThreadCount :: Int , oDBOptions :: PostgreSQL.Options + , oSchemaName :: String } deriving (Show, Eq) -- | Convert a 'PartialOptions' to a final 'Options' completeOptions :: PartialOptions -> Either [String] Options completeOptions = \case - PartialOptions { threadCount = Last (Just oThreadCount), dbOptions } -> - Options oThreadCount <$> PostgreSQL.completeOptions dbOptions + PartialOptions { threadCount = Last (Just oThreadCount), dbOptions, schemaName = Last (Just oSchemaName) } -> + Options oThreadCount <$> PostgreSQL.completeOptions dbOptions <*> pure oSchemaName + _ -> Left ["Missing threadCount"] {-| This function is a helper for creating queue consumer executables. @@ -166,10 +171,10 @@ threads :: [Async (StM m ())] <- replicateM oThreadCount $ async $ void $ forever $ do - payload <- liftIO $ withResource connectionPool lock - count <- liftIO $ withResource connectionPool getCount + payload <- liftIO $ withResource connectionPool $ lock oSchemaName + count <- liftIO $ withResource connectionPool $ getCount oSchemaName f payload count - liftIO $ withResource connectionPool $ flip dequeue (pId payload) + liftIO $ withResource connectionPool $ \c -> dequeue oSchemaName c (pId payload) _ :: (Async (StM m ()), ()) <- waitAnyCancel threads return () diff --git a/src/Database/PostgreSQL/Simple/Queue/Migrate.hs b/src/Database/PostgreSQL/Simple/Queue/Migrate.hs --- a/src/Database/PostgreSQL/Simple/Queue/Migrate.hs +++ b/src/Database/PostgreSQL/Simple/Queue/Migrate.hs @@ -1,8 +1,10 @@ -{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE QuasiQuotes, OverloadedStrings #-} module Database.PostgreSQL.Simple.Queue.Migrate where import Control.Monad import Database.PostgreSQL.Simple import Database.PostgreSQL.Simple.SqlQQ +import Data.Monoid +import Data.String {-| This function creates a table and enumeration type that is appriopiate for the queue. The following sql is used. @@ -35,8 +37,10 @@ language 'plpgsql'; @ -} -migrate :: Connection -> IO () -migrate conn = void $ execute_ conn [sql| +migrate :: String -> Connection -> IO () +migrate schemaName conn = void $ execute_ conn $ + "CREATE SCHEMA IF NOT EXISTS " <> fromString schemaName <> ";" <> + "SET search_path TO " <> fromString schemaName <> ";" <> [sql| CREATE TYPE state_t AS ENUM ('enqueued', 'locked', 'dequeued'); CREATE TABLE payloads diff --git a/test/Database/PostgreSQL/Simple/QueueSpec.hs b/test/Database/PostgreSQL/Simple/QueueSpec.hs --- a/test/Database/PostgreSQL/Simple/QueueSpec.hs +++ b/test/Database/PostgreSQL/Simple/QueueSpec.hs @@ -17,34 +17,37 @@ main :: IO () main = hspec spec +schemaName :: String +schemaName = "queue" + spec :: Spec -spec = describeDB migrate "Database.Queue" $ do +spec = describeDB (migrate schemaName) "Database.Queue" $ do itDB "empty locks nothing" $ - tryLockDB `shouldReturn` Nothing + tryLockDB schemaName `shouldReturn` Nothing itDB "empty gives count 0" $ - getCountDB `shouldReturn` 0 + getCountDB schemaName `shouldReturn` 0 itDB "enqueues/tryLocks/unlocks" $ do - payloadId <- enqueueDB $ String "Hello" - getCountDB `shouldReturn` 1 + payloadId <- enqueueDB schemaName $ String "Hello" + getCountDB schemaName `shouldReturn` 1 - Just Payload {..} <- tryLockDB - getCountDB `shouldReturn` 0 + Just Payload {..} <- tryLockDB schemaName + getCountDB schemaName `shouldReturn` 0 pId `shouldBe` payloadId pValue `shouldBe` String "Hello" - tryLockDB `shouldReturn` Nothing + tryLockDB schemaName `shouldReturn` Nothing - unlockDB pId - getCountDB `shouldReturn` 1 + unlockDB schemaName pId + getCountDB schemaName `shouldReturn` 1 itDB "tryLocks/dequeues" $ do - Just Payload {..} <- tryLockDB - getCountDB `shouldReturn` 0 + Just Payload {..} <- tryLockDB schemaName + getCountDB schemaName `shouldReturn` 0 - dequeueDB pId `shouldReturn` () - tryLockDB `shouldReturn` Nothing + dequeueDB schemaName pId `shouldReturn` () + tryLockDB schemaName `shouldReturn` Nothing it "enqueues and dequeues concurrently tryLock" $ \testDB -> do let withPool' = withPool testDB @@ -54,18 +57,18 @@ ref <- newIORef [] loopThreads <- replicateM 10 $ async $ fix $ \next -> do - mpayload <- withPool' tryLock + mpayload <- withPool' $ tryLock schemaName case mpayload of Nothing -> next Just Payload {..} -> do lastCount <- atomicModifyIORef ref $ \xs -> (pValue : xs, length xs + 1) - withPool' $ flip dequeue pId + withPool' $ \c -> dequeue schemaName c pId when (lastCount < elementCount) next -- Fork a hundred threads and enqueue an index forM_ [0 .. elementCount - 1] $ \i -> - forkIO $ void $ withPool' $ flip enqueue $ toJSON i + forkIO $ void $ withPool' $ \c -> enqueue schemaName c $ toJSON i waitAnyCancel loopThreads Just decoded <- mapM (decode . encode) <$> readIORef ref @@ -79,15 +82,15 @@ ref <- newIORef [] loopThreads <- replicateM 10 $ async $ fix $ \next -> do - Payload {..} <- withPool' lock + Payload {..} <- withPool' $ lock schemaName lastCount <- atomicModifyIORef ref $ \xs -> (pValue : xs, length xs + 1) - withPool' $ flip dequeue pId + withPool' $ \c -> dequeue schemaName c pId when (lastCount < elementCount) next -- Fork a hundred threads and enqueue an index - forM_ [0 .. elementCount - 1] $ \i -> forkIO $ void $ withPool' $ - flip enqueue $ toJSON i + forM_ [0 .. elementCount - 1] $ \i -> forkIO $ void $ withPool' $ \c -> + enqueue schemaName c $ toJSON i waitAnyCancel loopThreads Just decoded <- mapM (decode . encode) <$> readIORef ref