{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE OverloadedStrings #-}

-- | Resolving which queues a worker process should run from
-- @ARBITER_ENABLED_QUEUES@.
module Arbiter.Worker.EnabledQueues
  ( enabledQueuesEnvVar
  , getEnabledQueues
  , enabledQueuesForMonad
  , requestedQueues
  , requestedQueuesForMonad
  ) where

import Arbiter.Core.Exceptions (throwInternal)
import Arbiter.Core.MonadArbiter (RegistryOf)
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"

-- | The comma-separated queue names an environment variable asks for, every one of them
-- checked against the registry. Unset or blank gives the registry's whole queue set, and
-- a name outside it throws.
getEnabledQueues
  :: (RegistryTables registry)
  => String
  -- ^ Environment variable name
  -> Proxy registry
  -- ^ Registry proxy
  -> IO [Text]
getEnabledQueues :: forall (registry :: JobPayloadRegistry).
RegistryTables registry =>
String -> Proxy registry -> IO [Text]
getEnabledQueues String
envVar 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
<$> String -> Proxy registry -> IO (Maybe [Text])
forall (registry :: JobPayloadRegistry).
RegistryTables registry =>
String -> Proxy registry -> IO (Maybe [Text])
requestedQueues String
envVar Proxy registry
registry

-- | The queue names an environment variable asks for, or 'Nothing' when it is
-- unset or blank. Names are validated against the registry.
requestedQueues
  :: (RegistryTables registry)
  => String
  -- ^ Environment variable name
  -> Proxy registry
  -- ^ Registry proxy
  -> IO (Maybe [Text])
requestedQueues :: forall (registry :: JobPayloadRegistry).
RegistryTables registry =>
String -> Proxy registry -> IO (Maybe [Text])
requestedQueues String
envVar Proxy registry
registry = do
  rawValue <- String -> IO (Maybe String)
lookupEnv String
envVar
  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
envVar 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

-- | 'getEnabledQueues' for @ARBITER_ENABLED_QUEUES@. Resolve the registry from
-- the monad through 'RegistryOf'.
enabledQueuesForMonad
  :: forall m
   . (RegistryTables (RegistryOf m))
  => IO [Text]
enabledQueuesForMonad :: forall (m :: * -> *). RegistryTables (RegistryOf m) => IO [Text]
enabledQueuesForMonad = String -> Proxy (RegistryOf m) -> IO [Text]
forall (registry :: JobPayloadRegistry).
RegistryTables registry =>
String -> Proxy registry -> IO [Text]
getEnabledQueues String
enabledQueuesEnvVar (forall (t :: JobPayloadRegistry). Proxy t
forall {k} (t :: k). Proxy t
Proxy @(RegistryOf m))

-- | 'requestedQueues' for @ARBITER_ENABLED_QUEUES@. Resolve the registry from
-- the monad through 'RegistryOf'.
requestedQueuesForMonad
  :: forall m
   . (RegistryTables (RegistryOf m))
  => IO (Maybe [Text])
requestedQueuesForMonad :: forall (m :: * -> *).
RegistryTables (RegistryOf m) =>
IO (Maybe [Text])
requestedQueuesForMonad = String -> Proxy (RegistryOf m) -> IO (Maybe [Text])
forall (registry :: JobPayloadRegistry).
RegistryTables registry =>
String -> Proxy registry -> IO (Maybe [Text])
requestedQueues String
enabledQueuesEnvVar (forall (t :: JobPayloadRegistry). Proxy t
forall {k} (t :: k). Proxy t
Proxy @(RegistryOf m))