{-# LANGUAGE OverloadedStrings #-}

-- | Worker queue selection from @ARBITER_ENABLED_QUEUES@.
module Arbiter.Worker.EnabledQueues
  ( getEnabledQueues
  , requestedQueues
  ) where

import Arbiter.Core.Exceptions (throwInternal)
import Arbiter.Core.QueueRegistry (RegistryTables (..))
import Data.Maybe (fromMaybe)
import Data.Proxy (Proxy)
import Data.Text (Text)
import Data.Text qualified as T
import System.Environment (lookupEnv)

-- | The environment variable naming the queues a worker process should run.
enabledQueuesEnvVar :: String
enabledQueuesEnvVar :: String
enabledQueuesEnvVar = String
"ARBITER_ENABLED_QUEUES"

-- | Validated, comma-separated queue names from @ARBITER_ENABLED_QUEUES@.
-- Unset or blank selects all registered queues. Unknown names, or a value that names
-- no queues, throw 'Arbiter.Core.Exceptions.InternalException'.
getEnabledQueues :: (RegistryTables registry) => Proxy registry -> IO [Text]
getEnabledQueues :: forall (registry :: JobPayloadRegistry).
RegistryTables registry =>
Proxy registry -> IO [Text]
getEnabledQueues Proxy registry
registry =
  [Text] -> Maybe [Text] -> [Text]
forall a. a -> Maybe a -> a
fromMaybe (Proxy registry -> [Text]
forall (registry :: JobPayloadRegistry).
RegistryTables registry =>
Proxy registry -> [Text]
registryTableNames Proxy registry
registry) (Maybe [Text] -> [Text]) -> IO (Maybe [Text]) -> IO [Text]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Proxy registry -> IO (Maybe [Text])
forall (registry :: JobPayloadRegistry).
RegistryTables registry =>
Proxy registry -> IO (Maybe [Text])
requestedQueues Proxy registry
registry

-- | Validated names from @ARBITER_ENABLED_QUEUES@, or 'Nothing' when unset or blank.
requestedQueues :: (RegistryTables registry) => Proxy registry -> IO (Maybe [Text])
requestedQueues :: forall (registry :: JobPayloadRegistry).
RegistryTables registry =>
Proxy registry -> IO (Maybe [Text])
requestedQueues Proxy registry
registry = do
  rawValue <- String -> IO (Maybe String)
lookupEnv String
enabledQueuesEnvVar
  case T.strip . T.pack <$> rawValue of
    Just Text
trimmed | Bool -> Bool
not (Text -> Bool
T.null Text
trimmed) -> [Text] -> Maybe [Text]
forall a. a -> Maybe a
Just ([Text] -> Maybe [Text]) -> IO [Text] -> IO (Maybe [Text])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> IO [Text]
validate Text
trimmed
    Maybe Text
_ -> Maybe [Text] -> IO (Maybe [Text])
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe [Text]
forall a. Maybe a
Nothing
  where
    validate :: Text -> IO [Text]
validate Text
trimmed =
      let allQueues :: [Text]
allQueues = Proxy registry -> [Text]
forall (registry :: JobPayloadRegistry).
RegistryTables registry =>
Proxy registry -> [Text]
registryTableNames Proxy registry
registry
          requested :: [Text]
requested = (Text -> Bool) -> [Text] -> [Text]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Text -> Bool) -> Text -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Bool
T.null) ([Text] -> [Text]) -> ([Text] -> [Text]) -> [Text] -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text -> Text) -> [Text] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map Text -> Text
T.strip ([Text] -> [Text]) -> [Text] -> [Text]
forall a b. (a -> b) -> a -> b
$ HasCallStack => Text -> Text -> [Text]
Text -> Text -> [Text]
T.splitOn Text
"," Text
trimmed
          invalid :: [Text]
invalid = (Text -> Bool) -> [Text] -> [Text]
forall a. (a -> Bool) -> [a] -> [a]
filter (Text -> [Text] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Text]
allQueues) [Text]
requested
       in case ([Text]
requested, [Text]
invalid) of
            ([], [Text]
_) -> Text -> IO [Text]
forall (m :: * -> *) a. MonadIO m => Text -> m a
throwInternal (Text -> IO [Text]) -> Text -> IO [Text]
forall a b. (a -> b) -> a -> b
$ String -> Text
T.pack String
enabledQueuesEnvVar Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" is set but names no queues"
            ([Text]
_, []) -> [Text] -> IO [Text]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Text]
requested
            ([Text], [Text])
_ -> Text -> IO [Text]
forall (m :: * -> *) a. MonadIO m => Text -> m a
throwInternal (Text -> IO [Text]) -> Text -> IO [Text]
forall a b. (a -> b) -> a -> b
$ Text
"Unknown queue names: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> [Text] -> Text
T.intercalate Text
", " [Text]
invalid