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.3.0.0 +version: 0.4.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. @@ -16,7 +16,7 @@ . > forever $ do > -- Attempt get a payload or block until one is available - > payload <- lock "queue_schema" conn + > payload <- lock "queue" conn > > -- Perform application specifc parsing of the payload value > case fromJSON $ pValue payload of @@ -24,7 +24,7 @@ > Error err -> logErr err > > -- Remove the payload from future processing - > dequeue "queue_schema" conn $ pId payload + > dequeue "queue" conn $ pId payload > > To support multiple queues in the same database, the API expects a table name string > to determine which queue tables to use. 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 @@ -143,8 +143,11 @@ ------------------------------------------------------------------------------- --- DB API ------------------------------------------------------------------------------- +withSchema :: String -> Simple.Query -> Simple.Query +withSchema schemaName q = "SET search_path TO " <> fromString schemaName <> "; " <> q + notifyName :: IsString s => String -> s -notifyName tableName = fromString $ tableName <> "_enqueue" +notifyName schemaName = fromString $ schemaName <> "_enqueue" {-| Enqueue a new JSON value into the queue. This particularly function can be composed as part of a larger database transaction. For instance, @@ -158,10 +161,10 @@ @ -} enqueueDB :: String -> Value -> DB PayloadId -enqueueDB tableName value = - fmap head $ query ([sql| - NOTIFY |] <> " " <> notifyName tableName <> ";" <> [sql| - INSERT INTO|] <> " " <> fromString tableName <> " " <> [sql|(value) +enqueueDB schemaName value = + fmap head $ query (withSchema schemaName $ [sql| + NOTIFY |] <> " " <> notifyName schemaName <> ";" <> [sql| + INSERT INTO payloads (value) VALUES (?) RETURNING id;|] ) @@ -174,12 +177,12 @@ blocking version utilizing PostgreSQL's NOTIFY and LISTEN, see 'lock' -} tryLockDB :: String -> DB (Maybe Payload) -tryLockDB tableName = fmap listToMaybe $ query_ $ - [sql| UPDATE|] <> " " <> fromString tableName <> " " <> [sql| +tryLockDB schemaName = fmap listToMaybe $ query_ $ withSchema schemaName + [sql| UPDATE payloads SET state='locked' WHERE id in ( SELECT id - FROM|] <> " " <> fromString tableName <> " " <> [sql| + FROM payloads WHERE state='enqueued' ORDER BY created_at ASC LIMIT 1 @@ -193,8 +196,8 @@ useful. The DB version is provided for completeness. -} unlockDB :: String -> PayloadId -> DB () -unlockDB tableName payloadId = void $ execute ( - [sql| UPDATE|] <> " " <> fromString tableName <> " " <> [sql| +unlockDB schemaName payloadId = void $ execute (withSchema schemaName + [sql| UPDATE payloads SET state='enqueued' WHERE id=? AND state='locked' |]) @@ -202,8 +205,8 @@ -- | Transition a 'Payload' to the 'Dequeued' state. dequeueDB :: String -> PayloadId -> DB () -dequeueDB tableName payloadId = void $ execute ( - [sql| UPDATE|] <> " " <> fromString tableName <> " " <> [sql| +dequeueDB schemaName payloadId = void $ execute (withSchema schemaName + [sql| UPDATE payloads SET state='dequeued' WHERE id=? |]) @@ -211,9 +214,9 @@ -- | Get the number of rows in the 'Enqueued' state. getCountDB :: String -> DB Int64 -getCountDB tableName = fmap (fromOnly . head) $ query_ $ +getCountDB schemaName = fmap (fromOnly . head) $ query_ $ withSchema schemaName [sql| SELECT count(*) - FROM|] <> " " <> fromString tableName <> " " <> [sql| + FROM payloads WHERE state='enqueued' |] ------------------------------------------------------------------------------- @@ -223,7 +226,7 @@ which can be composed with other queries in a single transaction. -} enqueue :: String -> Connection -> Value -> IO PayloadId -enqueue tableName conn value = runDBT (enqueueDB tableName value) ReadCommitted conn +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 @@ -231,12 +234,12 @@ 'Serializable' transaction. -} tryLock :: String -> Connection -> IO (Maybe Payload) -tryLock tableName conn = runDBTSerializable (tryLockDB tableName) conn +tryLock schemaName conn = runDBTSerializable (tryLockDB schemaName) conn notifyPayload :: String -> Connection -> IO () -notifyPayload tableName conn = do +notifyPayload schemaName conn = do Notification {..} <- getNotification conn - unless (notificationChannel == notifyName tableName) $ notifyPayload tableName 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 @@ -244,14 +247,14 @@ waiting for new payloads, without scarficing promptness. -} lock :: String -> Connection -> IO Payload -lock tableName conn = bracket_ - (Simple.execute_ conn $ "LISTEN " <> notifyName tableName) - (Simple.execute_ conn $ "UNLISTEN " <> notifyName tableName) +lock schemaName conn = bracket_ + (Simple.execute_ conn $ "LISTEN " <> notifyName schemaName) + (Simple.execute_ conn $ "UNLISTEN " <> notifyName schemaName) $ fix $ \continue -> do - m <- tryLock tableName conn + m <- tryLock schemaName conn case m of Nothing -> do - notifyPayload tableName conn + notifyPayload schemaName conn continue Just x -> return x @@ -260,15 +263,15 @@ shutdown. For a DB API version see 'unlockDB' -} unlock :: String -> Connection -> PayloadId -> IO () -unlock tableName conn x = runDBTSerializable (unlockDB tableName x) conn +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 :: String -> Connection -> PayloadId -> IO () -dequeue tableName conn x = runDBTSerializable (dequeueDB tableName x) conn +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 :: String -> Connection -> IO Int64 -getCount tableName = runDBT (getCountDB tableName) ReadCommitted +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,7 +73,7 @@ data PartialOptions = PartialOptions { threadCount :: Last Int , dbOptions :: PostgreSQL.PartialOptions - , tableName :: Last String + , schemaName :: Last String } deriving (Show, Eq, Typeable) instance Monoid PartialOptions where @@ -82,7 +82,7 @@ PartialOptions (threadCount x <> threadCount y) (dbOptions x <> dbOptions y) - (tableName x <> tableName y) + (schemaName x <> schemaName y) -- | The default 'threadCount' is 1. -- The default db options are specified in @@ -107,7 +107,7 @@ -- | Convert a 'PartialOptions' to a final 'Options' completeOptions :: PartialOptions -> Either [String] Options completeOptions = \case - PartialOptions { threadCount = Last (Just oThreadCount), dbOptions, tableName = Last (Just oSchemaName) } -> + PartialOptions { threadCount = Last (Just oThreadCount), dbOptions, schemaName = Last (Just oSchemaName) } -> Options oThreadCount <$> PostgreSQL.completeOptions dbOptions <*> pure oSchemaName _ -> Left ["Missing threadCount"] 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 @@ -12,7 +12,7 @@ @ CREATE TYPE state_t AS ENUM ('enqueued', 'locked', 'dequeued'); - CREATE TABLE |] <> " " <> fromString tableName <> [sql| + CREATE TABLE payloads ( id uuid PRIMARY KEY , value jsonb NOT NULL , state state_t NOT NULL DEFAULT 'enqueued' @@ -20,7 +20,7 @@ , modified_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT clock_timestamp() ); - CREATE INDEX state_idx ON|] <> " " <> fromString tableName <> " " <> [sql|(state); + CREATE INDEX state_idx ON payloads (state); CREATE OR REPLACE FUNCTION update_row_modified_function_() RETURNS TRIGGER @@ -38,10 +38,12 @@ @ -} migrate :: String -> Connection -> IO () -migrate tableName conn = void $ execute_ conn $ [sql| +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 |] <> " " <> fromString tableName <> [sql| + CREATE TABLE payloads ( id SERIAL PRIMARY KEY , value jsonb NOT NULL , state state_t NOT NULL DEFAULT 'enqueued' @@ -49,7 +51,7 @@ , modified_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT clock_timestamp() ); - CREATE INDEX state_idx ON |] <> " " <> fromString tableName <> [sql|(state); + CREATE INDEX state_idx ON payloads (state); CREATE OR REPLACE FUNCTION update_row_modified_function_() RETURNS TRIGGER 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,37 +17,37 @@ main :: IO () main = hspec spec -tableName :: String -tableName = "queue" +schemaName :: String +schemaName = "queue" spec :: Spec -spec = describeDB (migrate tableName) "Database.Queue" $ do +spec = describeDB (migrate schemaName) "Database.Queue" $ do itDB "empty locks nothing" $ - tryLockDB tableName `shouldReturn` Nothing + tryLockDB schemaName `shouldReturn` Nothing itDB "empty gives count 0" $ - getCountDB tableName `shouldReturn` 0 + getCountDB schemaName `shouldReturn` 0 itDB "enqueues/tryLocks/unlocks" $ do - payloadId <- enqueueDB tableName $ String "Hello" - getCountDB tableName `shouldReturn` 1 + payloadId <- enqueueDB schemaName $ String "Hello" + getCountDB schemaName `shouldReturn` 1 - Just Payload {..} <- tryLockDB tableName - getCountDB tableName `shouldReturn` 0 + Just Payload {..} <- tryLockDB schemaName + getCountDB schemaName `shouldReturn` 0 pId `shouldBe` payloadId pValue `shouldBe` String "Hello" - tryLockDB tableName `shouldReturn` Nothing + tryLockDB schemaName `shouldReturn` Nothing - unlockDB tableName pId - getCountDB tableName `shouldReturn` 1 + unlockDB schemaName pId + getCountDB schemaName `shouldReturn` 1 itDB "tryLocks/dequeues" $ do - Just Payload {..} <- tryLockDB tableName - getCountDB tableName `shouldReturn` 0 + Just Payload {..} <- tryLockDB schemaName + getCountDB schemaName `shouldReturn` 0 - dequeueDB tableName pId `shouldReturn` () - tryLockDB tableName `shouldReturn` Nothing + dequeueDB schemaName pId `shouldReturn` () + tryLockDB schemaName `shouldReturn` Nothing it "enqueues and dequeues concurrently tryLock" $ \testDB -> do let withPool' = withPool testDB @@ -57,18 +57,18 @@ ref <- newIORef [] loopThreads <- replicateM 10 $ async $ fix $ \next -> do - mpayload <- withPool' $ tryLock tableName + mpayload <- withPool' $ tryLock schemaName case mpayload of Nothing -> next Just Payload {..} -> do lastCount <- atomicModifyIORef ref $ \xs -> (pValue : xs, length xs + 1) - withPool' $ \c -> dequeue tableName c 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' $ \c -> enqueue tableName c $ toJSON i + forkIO $ void $ withPool' $ \c -> enqueue schemaName c $ toJSON i waitAnyCancel loopThreads Just decoded <- mapM (decode . encode) <$> readIORef ref @@ -82,15 +82,15 @@ ref <- newIORef [] loopThreads <- replicateM 10 $ async $ fix $ \next -> do - Payload {..} <- withPool' $ lock tableName + Payload {..} <- withPool' $ lock schemaName lastCount <- atomicModifyIORef ref $ \xs -> (pValue : xs, length xs + 1) - withPool' $ \c -> dequeue tableName c 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' $ \c -> - enqueue tableName c $ toJSON i + enqueue schemaName c $ toJSON i waitAnyCancel loopThreads Just decoded <- mapM (decode . encode) <$> readIORef ref