diff --git a/app/Main.hs b/app/Main.hs --- a/app/Main.hs +++ b/app/Main.hs @@ -2,45 +2,49 @@ module Main where - +import Protolude import PostgREST.Config (AppConfig (..), + PgVersion (..), minimumPgVersion, prettyVersion, readOptions) +import PostgREST.Error (prettyUsageError) +import PostgREST.OpenAPI (isMalformedProxyUri) import PostgREST.DbStructure +import PostgREST.App (postgrest) import PostgRESTWS +import PostgRESTWS.HasqlBroadcast -import Control.Monad -import Control.Monad.IO.Class (liftIO) -import Data.Monoid ((<>)) -import Data.String.Conversions (cs) +import Control.AutoUpdate +import Data.ByteString.Base64 (decode) +import Data.String (IsString (..)) +import Data.Text (stripPrefix, pack, replace) +import Data.Text.IO (hPutStrLn, readFile) +import Data.Function (id) +import Data.Time.Clock.POSIX (getPOSIXTime) import qualified Hasql.Query as H import qualified Hasql.Session as H +import qualified Hasql.Connection as H import qualified Hasql.Decoders as HD import qualified Hasql.Encoders as HE import qualified Hasql.Pool as P import Network.Wai.Handler.Warp import System.IO (BufferMode (..), - hSetBuffering, stderr, - stdin, stdout) -import Web.JWT (secret) + hSetBuffering) +import Data.IORef #ifndef mingw32_HOST_OS import System.Posix.Signals -import Control.Concurrent (myThreadId) -import Data.IORef -import Control.Exception.Base (throwTo, AsyncException(..)) #endif -import qualified Database.PostgreSQL.LibPQ as PQ isServerVersionSupported :: H.Session Bool isServerVersionSupported = do ver <- H.query () pgVersion - return $ read (cs ver) >= minimumPgVersion + return $ ver >= pgvNum minimumPgVersion where pgVersion = - H.statement "SHOW server_version_num" - HE.unit (HD.singleRow $ HD.value HD.text) True + H.statement "SELECT current_setting('server_version_num')::integer" + HE.unit (HD.singleRow $ HD.value HD.int4) False main :: IO () main = do @@ -48,30 +52,39 @@ hSetBuffering stdin LineBuffering hSetBuffering stderr NoBuffering - conf <- readOptions - let port = configPort conf - pgSettings = cs (configDatabase conf) - appSettings = setPort port - . setServerName (cs $ "postgrest/" <> prettyVersion) + conf <- loadSecretFile =<< readOptions + let host = configHost conf + port = configPort conf + proxy = configProxyUri conf + pgSettings = toS (configDatabase conf) + appSettings = setHost ((fromString . toS) host) + . setPort port + . setServerName (toS $ "postgrest/" <> prettyVersion) $ defaultSettings - unless (secret "secret" /= configJwtSecret conf) $ - putStrLn "WARNING, running in insecure mode, JWT secret is the default value" - Prelude.putStrLn $ "Listening on port " ++ - (show $ configPort conf :: String) + when (isMalformedProxyUri $ toS <$> proxy) $ panic + "Malformed proxy uri, a correct example: https://example.com:8443/basePath" + putStrLn $ ("Listening on port " :: Text) <> show (configPort conf) + pool <- P.acquire (configPool conf, 10, pgSettings) + conOrError <- H.acquire pgSettings + let con = either (panic . show) id conOrError :: H.Connection + result <- P.use pool $ do supported <- isServerVersionSupported - unless supported $ error ( + unless supported $ panic ( "Cannot run in this PostgreSQL version, PostgREST needs at least " - <> show minimumPgVersion) - getDbStructure (cs $ configSchema conf) + <> pgvName minimumPgVersion) + getDbStructure (toS $ configSchema conf) - refDbStructure <- newIORef $ either (error.show) id result - notificationsCon <- PQ.connectdb pgSettings + forM_ (lefts [result]) $ \e -> do + hPutStrLn stderr (prettyUsageError e) + exitFailure + refDbStructure <- newIORef $ either (panic . show) id result + #ifndef mingw32_HOST_OS tid <- myThreadId forM_ [sigINT, sigTERM] $ \sig -> @@ -82,8 +95,42 @@ void $ installHandler sigHUP ( Catch . void . P.use pool $ do - s <- getDbStructure (cs $ configSchema conf) + s <- getDbStructure (toS $ configSchema conf) liftIO $ atomicWriteIORef refDbStructure s ) Nothing #endif - runSettings appSettings $ postgrestWsApp conf refDbStructure pool notificationsCon + + -- ask for the OS time at most once per second + getTime <- mkAutoUpdate + defaultUpdateSettings { updateAction = getPOSIXTime } + + multi <- newHasqlBroadcaster con + void $ relayMessagesForever multi + runSettings appSettings $ + postgrestWsMiddleware (configJwtSecret conf) getTime pool multi $ + postgrest conf refDbStructure pool getTime + +loadSecretFile :: AppConfig -> IO AppConfig +loadSecretFile conf = extractAndTransform mSecret + where + mSecret = decodeUtf8 <$> configJwtSecret conf + isB64 = configJwtSecretIsBase64 conf + + extractAndTransform :: Maybe Text -> IO AppConfig + extractAndTransform Nothing = return conf + extractAndTransform (Just s) = + fmap setSecret $ transformString isB64 =<< + case stripPrefix "@" s of + Nothing -> return s + Just filename -> readFile (toS filename) + + transformString :: Bool -> Text -> IO ByteString + transformString False t = return . encodeUtf8 $ t + transformString True t = + case decode (encodeUtf8 $ replaceUrlChars t) of + Left errMsg -> panic $ pack errMsg + Right bs -> return bs + + setSecret bs = conf { configJwtSecret = Just bs } + + replaceUrlChars = replace "_" "/" . replace "-" "+" . replace "." "=" diff --git a/postgrest-ws.cabal b/postgrest-ws.cabal --- a/postgrest-ws.cabal +++ b/postgrest-ws.cabal @@ -1,5 +1,5 @@ name: postgrest-ws -version: 0.1.0.1 +version: 0.1.0.2 synopsis: PostgREST extension to map LISTEN/NOTIFY messages to Websockets description: Please see README.md homepage: https://github.com/diogob/postgrest-ws#readme @@ -17,13 +17,19 @@ hs-source-dirs: src ghc-options: -Wall exposed-modules: PostgRESTWS + , PostgRESTWS.Broadcast + , PostgRESTWS.Database + , PostgRESTWS.HasqlBroadcast + + other-modules: PostgRESTWS.Claims + , PostgRESTWS.Types build-depends: base >= 4.7 && < 5 , hasql-pool >= 0.4 && < 0.5 , text >= 1.2 && < 2 , wai >= 3.2 && < 4 - , websockets >= 0.9 && < 0.10 + , websockets >= 0.9 && < 0.11 , wai-websockets >= 3.0 && < 4 - , postgrest >= 0.3 && < 0.4 + , postgrest >= 0.3 && < 0.5 , http-types >= 0.9 , bytestring >= 0.10 , postgresql-libpq @@ -31,36 +37,60 @@ , unordered-containers >= 0.2 , postgresql-libpq >= 0.9 && < 1.0 , aeson >= 0.11 - , string-conversions >= 0.4 + , protolude >= 0.1.6 && < 0.2 + , jwt >= 0.7.2 && < 0.8 + , hasql >= 0.19 + , either + , stm-containers + , stm default-language: Haskell2010 - default-extensions: OverloadedStrings + default-extensions: OverloadedStrings, NoImplicitPrelude executable postgrest-ws hs-source-dirs: app main-is: Main.hs ghc-options: -Wall -threaded -rtsopts -with-rtsopts=-N build-depends: base >= 4.7 && < 5 - , transformers >= 0.4 && < 0.5 - , string-conversions >= 0.4 && < 0.5 - , hasql >= 0.19 && < 0.20 - , hasql-pool >= 0.4 && < 0.5 + , transformers >= 0.4 && < 0.6 + , hasql >= 0.19 + , hasql-pool >= 0.4 , warp >= 3.2 && < 4 , unix >= 2.7 && < 3 , jwt >= 0.7 && < 1 , postgrest , postgrest-ws , postgresql-libpq >= 0.9 && < 1.0 + , protolude >= 0.1.6 && < 0.2 + , auto-update + , base64-bytestring + , text + , time default-language: Haskell2010 - default-extensions: OverloadedStrings + default-extensions: OverloadedStrings, NoImplicitPrelude test-suite postgrest-ws-test type: exitcode-stdio-1.0 hs-source-dirs: test main-is: Spec.hs + other-modules: BroadcastSpec + , DatabaseSpec + , HasqlBroadcastSpec build-depends: base + , protolude >= 0.1.6 && < 0.2 , postgrest-ws + , containers + , hspec + , hspec-wai + , hspec-wai-json + , aeson + , hasql + , hasql-pool + , http-types + , wai-extra + , stm ghc-options: -Wall -threaded -rtsopts -with-rtsopts=-N default-language: Haskell2010 + default-extensions: OverloadedStrings, NoImplicitPrelude source-repository head type: git diff --git a/src/PostgRESTWS.hs b/src/PostgRESTWS.hs --- a/src/PostgRESTWS.hs +++ b/src/PostgRESTWS.hs @@ -1,121 +1,87 @@ {-# LANGUAGE DeriveGeneric #-} module PostgRESTWS - ( postgrestWsApp + ( postgrestWsMiddleware ) where -import GHC.IORef -import qualified Hasql.Pool as H +import Protolude import qualified Network.Wai as Wai import qualified Network.Wai.Handler.WebSockets as WS import qualified Network.WebSockets as WS -import PostgREST.App as PGR -import PostgREST.Config as PGR -import PostgREST.Types as PGR - -import qualified Data.Text as T -import qualified Data.Text.Encoding as T +import qualified Hasql.Pool as H -import Control.Monad (forever, void, when) -import qualified Data.HashMap.Strict as M -import Data.String.Conversions (cs) -import Data.Time.Clock.POSIX (getPOSIXTime) -import qualified Database.PostgreSQL.LibPQ as PQ -import PostgREST.Auth (jwtClaims) +import qualified Data.Text.Encoding.Error as T +import Data.Time.Clock.POSIX (POSIXTime) import qualified Data.Aeson as A import qualified Data.ByteString as BS -import Data.ByteString.Lazy (toStrict) -import Data.Monoid - -import Control.Concurrent (forkIO, threadWaitReadSTM) -import GHC.Conc (atomically) +import qualified Data.ByteString.Lazy as BL -import GHC.Generics +import PostgRESTWS.Claims +import PostgRESTWS.Database +import PostgRESTWS.Broadcast (Multiplexer, onMessage, readTChan) +import qualified PostgRESTWS.Broadcast as B data Message = Message - { userClaims :: A.Object - , payload :: T.Text + { claims :: A.Object + , payload :: Text } deriving (Show, Eq, Generic) instance A.ToJSON Message -postgrestWsApp :: PGR.AppConfig - -> IORef PGR.DbStructure - -> H.Pool - -> PQ.Connection - -> Wai.Application -postgrestWsApp conf refDbStructure pool pqCon = - WS.websocketsOr WS.defaultConnectionOptions wsApp $ postgrest conf refDbStructure pool +postgrestWsMiddleware :: Maybe ByteString -> IO POSIXTime -> H.Pool -> Multiplexer -> Wai.Application -> Wai.Application +postgrestWsMiddleware = WS.websocketsOr WS.defaultConnectionOptions `compose` wsApp where - -- when the websocket is closed a ConnectionClosed Exception is triggered - -- this kills all children and frees resources for us - wsApp :: WS.ServerApp - wsApp pendingConn = do - time <- getPOSIXTime - let - claimsOrExpired = jwtClaims jwtSecret jwtToken time - case claimsOrExpired of - Left e -> rejectRequest e - Right claims -> do - -- role claim defaults to anon if not specified in jwt - -- We should accept only after verifying JWT - conn <- WS.acceptRequest pendingConn - -- each websocket needs its own listen connection to avoid - -- handling of multiple waiting threads in the same connection - when (hasRead claims) $ - void $ forkIO $ listenSession (channel claims) conf conn - -- all websockets share a single connection to NOTIFY - when (hasWrite claims) $ - forever $ notifySession (channel claims) claims pqCon conn - where - claimAsBS name cl = let A.String s = (cl M.! name) in T.encodeUtf8 s - channel = claimAsBS ("channel" :: T.Text) - mode = claimAsBS ("mode" :: T.Text) - hasRead cl = mode cl == "r" || mode cl == "rw" - hasWrite cl = mode cl == "w" || mode cl == "rw" - rejectRequest = WS.rejectRequest pendingConn . T.encodeUtf8 - jwtSecret = configJwtSecret conf - -- the first char in path is '/' the rest is the token - jwtToken = T.decodeUtf8 $ BS.drop 1 $ WS.requestPath $ WS.pendingRequest pendingConn + compose = (.) . (.) . (.) . (.) +-- when the websocket is closed a ConnectionClosed Exception is triggered +-- this kills all children and frees resources for us +wsApp :: Maybe ByteString -> IO POSIXTime -> H.Pool -> Multiplexer -> WS.ServerApp +wsApp mSecret getTime pqCon multi pendingConn = getTime >>= forkSessionsWhenTokenIsValid . validateClaims mSecret jwtToken + where + forkSessionsWhenTokenIsValid = either rejectRequest forkSessions + hasRead m = m == ("r" :: ByteString) || m == ("rw" :: ByteString) + hasWrite m = m == ("w" :: ByteString) || m == ("rw" :: ByteString) + rejectRequest = WS.rejectRequest pendingConn . encodeUtf8 + -- the first char in path is '/' the rest is the token + jwtToken = decodeUtf8 $ BS.drop 1 $ WS.requestPath $ WS.pendingRequest pendingConn + forkSessions (channel, mode, validClaims) = do + -- role claim defaults to anon if not specified in jwt + -- We should accept only after verifying JWT + conn <- WS.acceptRequest pendingConn + -- Fork a pinging thread to ensure browser connections stay alive + WS.forkPingThread conn 30 + -- each websocket needs its own listen connection to avoid + -- handling of multiple waiting threads in the same connection + + when (hasRead mode) $ + onMessage multi channel (\ch -> + forever $ atomically (readTChan ch) >>= WS.sendTextData conn . B.payload) + -- all websockets share a single connection to NOTIFY + notifySessionFinished <- if hasWrite mode + then forkAndWait $ forever $ notifySession channel validClaims pqCon conn + else newMVar () + takeMVar notifySessionFinished + -- private functions + -- Having both channel and claims as parameters seem redundant -- But it allows the function to ignore the claims structure and the source --- of the channel +-- of the channel, so all claims decoding can be coded in the caller notifySession :: BS.ByteString -> A.Object - -> PQ.Connection + -> H.Pool -> WS.Connection -> IO () -notifySession channel claims pqCon wsCon = - WS.receiveData wsCon >>= (notify . jsonMsg) +notifySession channel claimsToSend pool wsCon = + WS.receiveData wsCon >>= (void . send . jsonMsg) where - notify msg = void $ PQ.exec pqCon ("NOTIFY " <> channel <> ", '" <> msg <> "'") - jsonMsg = toStrict . A.encode . Message claims . T.decodeUtf8 + send = notifyPool pool channel + -- we need to decode the bytestring to re-encode valid JSON for the notification + jsonMsg = BL.toStrict . A.encode . Message claimsToSend . decodeUtf8With T.lenientDecode -listenSession :: BS.ByteString - -> PGR.AppConfig - -> WS.Connection - -> IO () -listenSession channel conf wsCon = do - pqCon <- PQ.connectdb pgSettings - listen pqCon - waitForNotifications pqCon - where - waitForNotifications = forever . fetch - listen con = void $ PQ.exec con $ "LISTEN " <> channel - pgSettings = cs $ configDatabase conf - fetch con = do - mNotification <- PQ.notifies con - case mNotification of - Nothing -> do - mfd <- PQ.socket con - case mfd of - Nothing -> error "Error checking for PostgreSQL notifications" - Just fd -> do - (waitRead, _) <- threadWaitReadSTM fd - atomically waitRead - void $ PQ.consumeInput con - Just notification -> - WS.sendTextData wsCon $ PQ.notifyExtra notification +forkAndWait :: IO () -> IO (MVar ()) +forkAndWait io = do + mvar <- newEmptyMVar + void $ forkFinally io (\_ -> putMVar mvar ()) + return mvar diff --git a/src/PostgRESTWS/Broadcast.hs b/src/PostgRESTWS/Broadcast.hs new file mode 100644 --- /dev/null +++ b/src/PostgRESTWS/Broadcast.hs @@ -0,0 +1,94 @@ +module PostgRESTWS.Broadcast ( Multiplexer (src) + , Message (..) + , SourceCommands (..) + , newMultiplexer + , onMessage + , relayMessages + , relayMessagesForever + , openChannelProducer + , closeChannelProducer + -- reexports + , readTQueue + , writeTQueue + , readTChan + ) where + +import Protolude +import qualified STMContainers.Map as M +import Control.Concurrent.STM.TChan +import Control.Concurrent.STM.TQueue + +data SourceCommands = Open ByteString | Close ByteString deriving (Show) +data Message = Message { channel :: ByteString + , payload :: ByteString + } deriving (Eq, Show) + +data Multiplexer = Multiplexer { channels :: M.Map ByteString Channel + , src :: ThreadId + , commands :: TQueue SourceCommands + , messages :: TQueue Message + } + +data Channel = Channel { broadcast :: TChan Message + , listeners :: Integer + , close :: STM () + } + +openChannelProducer :: Multiplexer -> ByteString -> STM () +openChannelProducer multi ch = writeTQueue (commands multi) (Open ch) + +closeChannelProducer :: Multiplexer -> ByteString -> STM () +closeChannelProducer multi chan = writeTQueue (commands multi) (Close chan) + +relayMessagesForever :: Multiplexer -> IO ThreadId +relayMessagesForever = forkIO . forever . relayMessages + +relayMessages :: Multiplexer -> IO () +relayMessages multi = + atomically $ do + m <- readTQueue (messages multi) + mChannel <- M.lookup (channel m) (channels multi) + case mChannel of + Nothing -> return () + Just c -> writeTChan (broadcast c) m + +newMultiplexer :: (TQueue SourceCommands -> TQueue Message -> IO a) + -> (Either SomeException a -> IO ()) + -> IO Multiplexer +newMultiplexer openProducer closeProducer = do + cmds <- newTQueueIO + msgs <- newTQueueIO + m <- liftA2 Multiplexer M.newIO (forkFinally (openProducer cmds msgs) closeProducer) + return $ m cmds msgs + +openChannel :: Multiplexer -> ByteString -> STM Channel +openChannel multi chan = do + c <- newBroadcastTChan + let newChannel = Channel{ broadcast = c + , listeners = 0 + , close = closeChannelProducer multi chan + } + M.insert newChannel chan (channels multi) + openChannelProducer multi chan + return newChannel + +onMessage :: Multiplexer -> ByteString -> (TChan Message -> IO()) -> IO () +onMessage multi chan action = do + listener <- atomically $ do + mC <- M.lookup chan (channels multi) + c <- case mC of + Nothing -> openChannel multi chan + Just ch -> return ch + M.delete chan (channels multi) + let newChannel = Channel{ broadcast = broadcast c, listeners = listeners c + 1, close = close c} + M.insert newChannel chan (channels multi) + dupTChan $ broadcast newChannel + void $ forkFinally (action listener) (\_ -> atomically $ do + mC <- M.lookup chan (channels multi) + let c = fromMaybe (panic $ "trying to remove listener from non existing channel: " <> toS chan) mC + M.delete chan (channels multi) + when (listeners c - 1 == 0) $ closeChannelProducer multi chan + when (listeners c - 1 > 0) $ do + let newChannel = Channel{ broadcast = broadcast c, listeners = listeners c - 1, close = close c} + M.insert newChannel chan (channels multi) + ) diff --git a/src/PostgRESTWS/Claims.hs b/src/PostgRESTWS/Claims.hs new file mode 100644 --- /dev/null +++ b/src/PostgRESTWS/Claims.hs @@ -0,0 +1,33 @@ +module PostgRESTWS.Claims + ( validateClaims + ) where + +import Protolude +import qualified Data.HashMap.Strict as M +import PostgREST.Auth as PGR +import Web.JWT (binarySecret) +import Data.Aeson (Value (..)) +import Data.Time.Clock.POSIX (POSIXTime) + +type Claims = M.HashMap Text Value +type ConnectionInfo = (ByteString, ByteString, Claims) + +validateClaims :: Maybe ByteString -> Text -> POSIXTime -> Either Text ConnectionInfo +validateClaims secret jwtToken time = do + cl <- case jwtClaims jwtSecret jwtToken time of + PGR.JWTClaims c -> Right c + _ -> Left "Error" + jChannel <- claimAsJSON "channel" cl + jMode <- claimAsJSON "mode" cl + channel <- value2BS jChannel + mode <- value2BS jMode + Right (channel, mode, cl) + where + jwtSecret = binarySecret <$> secret + value2BS val = case val of + String s -> Right $ encodeUtf8 s + _ -> Left "claim is not string value" + claimAsJSON :: Text -> Claims -> Either Text Value + claimAsJSON name cl = case M.lookup name cl of + Just el -> Right el + Nothing -> Left (name <> " not in claims") diff --git a/src/PostgRESTWS/Database.hs b/src/PostgRESTWS/Database.hs new file mode 100644 --- /dev/null +++ b/src/PostgRESTWS/Database.hs @@ -0,0 +1,61 @@ +module PostgRESTWS.Database + ( notifyPool + , notify + , listen + , unlisten + , waitForNotifications + ) where + +import Protolude +import Hasql.Pool (Pool, UsageError, use) +import Hasql.Session (sql, run) +import qualified Hasql.Session as S +import Hasql.Connection (Connection, withLibPQConnection) +import qualified Database.PostgreSQL.LibPQ as PQ +import Data.Either.Combinators + +import PostgRESTWS.Types + +notifyPool :: Pool -> ByteString -> ByteString -> IO (Either Error ()) +notifyPool pool channel mesg = + mapError <$> use pool (sql ("NOTIFY " <> channel <> ", '" <> mesg <> "'")) + where + mapError :: Either UsageError () -> Either Error () + mapError = mapLeft (NotifyError . show) + +notify :: Connection -> ByteString -> ByteString -> IO (Either Error ()) +notify con channel mesg = + mapError <$> run (sql ("NOTIFY " <> channel <> ", '" <> mesg <> "'")) con + where + mapError :: Either S.Error () -> Either Error () + mapError = mapLeft (NotifyError . show) + +listen :: Connection -> ByteString -> IO () +listen con channel = + void $ withLibPQConnection con execListen + where + execListen pqCon = void $ PQ.exec pqCon $ "LISTEN " <> channel + +unlisten :: Connection -> ByteString -> IO () +unlisten con channel = + void $ withLibPQConnection con execListen + where + execListen pqCon = void $ PQ.exec pqCon $ "UNLISTEN " <> channel + + +waitForNotifications :: (ByteString -> ByteString -> IO()) -> Connection -> IO () +waitForNotifications sendNotification con = + withLibPQConnection con $ void . forkIO . forever . pqFetch + where + pqFetch pqCon = do + mNotification <- PQ.notifies pqCon + case mNotification of + Nothing -> do + mfd <- PQ.socket pqCon + case mfd of + Nothing -> panic "Error checking for PostgreSQL notifications" + Just fd -> do + void $ threadWaitRead fd + void $ PQ.consumeInput pqCon + Just notification -> + sendNotification (PQ.notifyRelname notification) (PQ.notifyExtra notification) diff --git a/src/PostgRESTWS/HasqlBroadcast.hs b/src/PostgRESTWS/HasqlBroadcast.hs new file mode 100644 --- /dev/null +++ b/src/PostgRESTWS/HasqlBroadcast.hs @@ -0,0 +1,26 @@ +module PostgRESTWS.HasqlBroadcast + ( newHasqlBroadcaster + -- re-export + , acquire + , relayMessages + , relayMessagesForever + ) where + +import Protolude + +import Hasql.Connection + +import PostgRESTWS.Database +import PostgRESTWS.Broadcast + +newHasqlBroadcaster :: Connection -> IO Multiplexer +newHasqlBroadcaster con = newMultiplexer (\cmds msgs-> do + waitForNotifications + (\c m-> atomically $ writeTQueue msgs $ Message c m) + con + forever $ do + cmd <- atomically $ readTQueue cmds + case cmd of + Open ch -> listen con ch + Close ch -> unlisten con ch + ) (\_ -> return ()) diff --git a/src/PostgRESTWS/Types.hs b/src/PostgRESTWS/Types.hs new file mode 100644 --- /dev/null +++ b/src/PostgRESTWS/Types.hs @@ -0,0 +1,5 @@ +module PostgRESTWS.Types where + +import Protolude + +newtype Error = NotifyError Text diff --git a/test/BroadcastSpec.hs b/test/BroadcastSpec.hs new file mode 100644 --- /dev/null +++ b/test/BroadcastSpec.hs @@ -0,0 +1,48 @@ +module BroadcastSpec (spec) where + +import Protolude +import Control.Concurrent.STM.TChan +import Control.Concurrent.STM.TQueue + +import Test.Hspec + +import PostgRESTWS.Broadcast + +spec :: Spec +spec = do + describe "newMultiplexer" $ + it "opens a separate thread for a producer function" $ do + output <- newTQueueIO :: IO (TQueue ThreadId) + + void $ liftIO $ newMultiplexer (\_ _-> do + tid <- myThreadId + atomically $ writeTQueue output tid + ) (\_ -> return ()) + + outMsg <- atomically $ readTQueue output + myThreadId `shouldNotReturn` outMsg + describe "relayMessages" $ + it "relays a single message from producer to 1 listener on 1 test channel" $ do + output <- newTQueueIO :: IO (TQueue Message) + multi <- liftIO $ newMultiplexer (\_ msgs-> + atomically $ writeTQueue msgs (Message "test" "payload")) (\_ -> return ()) + void $ onMessage multi "test" (\ch -> + atomically $ readTChan ch >>= writeTQueue output) + + liftIO $ relayMessages multi + + outMsg <- atomically $ readTQueue output + outMsg `shouldBe` Message "test" "payload" + describe "onMessage" $ + it "sends an open command to the producer" $ do + output <- newTQueueIO :: IO (TQueue Text) + multi <- liftIO $ newMultiplexer (\cmds _-> + atomically $ do + cmd <- readTQueue cmds + writeTQueue output (show cmd :: Text) + ) (\_ -> return ()) + + void $ onMessage multi "new channel" (\_ -> return ()) + + outMsg <- atomically $ readTQueue output + outMsg `shouldBe` "Open \"new channel\"" diff --git a/test/DatabaseSpec.hs b/test/DatabaseSpec.hs new file mode 100644 --- /dev/null +++ b/test/DatabaseSpec.hs @@ -0,0 +1,26 @@ +module DatabaseSpec (spec) where + +import Protolude + +import Data.Function (id) +import qualified Hasql.Connection as H +import Test.Hspec + +import PostgRESTWS.Database + +spec :: Spec +spec = + describe "waitForNotifications" $ + it "does not block the connection ann trigger action upon notification" $ do + conOrError <- H.acquire "postgres://localhost/postgrest_test" + let con = either (panic . show) id conOrError :: H.Connection + notification <- liftIO newEmptyMVar + + waitForNotifications (curry $ putMVar notification) con + listen con "test" + + conOrError2 <- H.acquire "postgres://localhost/postgrest_test" + let con2 = either (panic . show) id conOrError2 :: H.Connection + void $ notify con2 "test" "hello there" + + readMVar notification `shouldReturn` ("test", "hello there") diff --git a/test/HasqlBroadcastSpec.hs b/test/HasqlBroadcastSpec.hs new file mode 100644 --- /dev/null +++ b/test/HasqlBroadcastSpec.hs @@ -0,0 +1,46 @@ +module HasqlBroadcastSpec (spec) where + +import Protolude + +import Data.Function (id) +import qualified Hasql.Query as H +import qualified Hasql.Session as H +import qualified Hasql.Decoders as HD +import qualified Hasql.Encoders as HE + +import Test.Hspec + +import PostgRESTWS.Broadcast +import PostgRESTWS.HasqlBroadcast + +spec :: Spec +spec = describe "newHasqlBroadcaster" $ do + let booleanQueryShouldReturn con query expected = + either (panic . show) id + <$> H.run query con + `shouldReturn` expected + newConnection connStr = + either (panic . show) id + <$> acquire connStr + + it "start listening on a database connection as we send an Open command" $ do + con <- newConnection "postgres://localhost/postgrest_test" + multi <- liftIO $ newHasqlBroadcaster con + + atomically $ openChannelProducer multi "test" + + let statement = H.statement "SELECT EXISTS (SELECT 1 FROM pg_stat_activity WHERE query ~* 'LISTEN \"test\"')" + HE.unit (HD.singleRow $ HD.value HD.bool) False + query = H.query () statement + booleanQueryShouldReturn con query True + + it "stops listening on a database connection as we send a Close command" $ do + con <- newConnection "postgres://localhost/postgrest_test" + multi <- liftIO $ newHasqlBroadcaster con + + atomically $ closeChannelProducer multi "test" + + let statement = H.statement "SELECT EXISTS (SELECT 1 FROM pg_stat_activity WHERE query ~* 'UNLISTEN \"test\"')" + HE.unit (HD.singleRow $ HD.value HD.bool) False + query = H.query () statement + booleanQueryShouldReturn con query True diff --git a/test/Spec.hs b/test/Spec.hs --- a/test/Spec.hs +++ b/test/Spec.hs @@ -1,2 +1,1 @@ -main :: IO () -main = putStrLn "Test suite not yet implemented" +{-# OPTIONS_GHC -F -pgmF hspec-discover #-}