postgresql-simple-queue 0.1.0.1 → 0.2.0.0
raw patch · 5 files changed
+99/−73 lines, 5 files
Files
- postgresql-simple-queue.cabal +8/−5
- src/Database/PostgreSQL/Simple/Queue.hs +48/−37
- src/Database/PostgreSQL/Simple/Queue/Main.hs +12/−7
- src/Database/PostgreSQL/Simple/Queue/Migrate.hs +7/−3
- test/Database/PostgreSQL/Simple/QueueSpec.hs +24/−21
postgresql-simple-queue.cabal view
@@ -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
src/Database/PostgreSQL/Simple/Queue.hs view
@@ -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
src/Database/PostgreSQL/Simple/Queue/Main.hs view
@@ -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 ()
src/Database/PostgreSQL/Simple/Queue/Migrate.hs view
@@ -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
test/Database/PostgreSQL/Simple/QueueSpec.hs view
@@ -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