{-# LANGUAGE OverloadedStrings, RecordWildCards #-}
module Database.Redis.Cluster.Command where

import Data.Char(toLower)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as Char8
import qualified Data.HashMap.Strict as HM
import Database.Redis.Types(RedisResult(decode))
import Database.Redis.Protocol(Reply(..))

data Flag
    = Write
    | ReadOnly
    | DenyOOM
    | Admin
    | PubSub
    | NoScript
    | Random
    | SortForScript
    | Loading
    | Stale
    | SkipMonitor
    | Asking
    | Fast
    | MovableKeys
    | Other BS.ByteString deriving (Int -> Flag -> ShowS
[Flag] -> ShowS
Flag -> [Char]
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
showList :: [Flag] -> ShowS
$cshowList :: [Flag] -> ShowS
show :: Flag -> [Char]
$cshow :: Flag -> [Char]
showsPrec :: Int -> Flag -> ShowS
$cshowsPrec :: Int -> Flag -> ShowS
Show, Flag -> Flag -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Flag -> Flag -> Bool
$c/= :: Flag -> Flag -> Bool
== :: Flag -> Flag -> Bool
$c== :: Flag -> Flag -> Bool
Eq)


data AritySpec = Required Integer | MinimumRequired Integer deriving (Int -> AritySpec -> ShowS
[AritySpec] -> ShowS
AritySpec -> [Char]
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
showList :: [AritySpec] -> ShowS
$cshowList :: [AritySpec] -> ShowS
show :: AritySpec -> [Char]
$cshow :: AritySpec -> [Char]
showsPrec :: Int -> AritySpec -> ShowS
$cshowsPrec :: Int -> AritySpec -> ShowS
Show)

data LastKeyPositionSpec = LastKeyPosition Integer | UnlimitedKeys Integer deriving (Int -> LastKeyPositionSpec -> ShowS
[LastKeyPositionSpec] -> ShowS
LastKeyPositionSpec -> [Char]
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
showList :: [LastKeyPositionSpec] -> ShowS
$cshowList :: [LastKeyPositionSpec] -> ShowS
show :: LastKeyPositionSpec -> [Char]
$cshow :: LastKeyPositionSpec -> [Char]
showsPrec :: Int -> LastKeyPositionSpec -> ShowS
$cshowsPrec :: Int -> LastKeyPositionSpec -> ShowS
Show)

newtype InfoMap = InfoMap (HM.HashMap String CommandInfo)

-- Represents the result of the COMMAND command, which returns information
-- about the position of keys in a request
data CommandInfo = CommandInfo
    { CommandInfo -> ByteString
name :: BS.ByteString
    , CommandInfo -> AritySpec
arity :: AritySpec
    , CommandInfo -> [Flag]
flags :: [Flag]
    , CommandInfo -> Integer
firstKeyPosition :: Integer
    , CommandInfo -> LastKeyPositionSpec
lastKeyPosition :: LastKeyPositionSpec
    , CommandInfo -> Integer
stepCount :: Integer
    } deriving (Int -> CommandInfo -> ShowS
[CommandInfo] -> ShowS
CommandInfo -> [Char]
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
showList :: [CommandInfo] -> ShowS
$cshowList :: [CommandInfo] -> ShowS
show :: CommandInfo -> [Char]
$cshow :: CommandInfo -> [Char]
showsPrec :: Int -> CommandInfo -> ShowS
$cshowsPrec :: Int -> CommandInfo -> ShowS
Show)

instance RedisResult CommandInfo where
    decode :: Reply -> Either Reply CommandInfo
decode (MultiBulk (Just
        [ Bulk (Just ByteString
commandName)
        , Integer Integer
aritySpec
        , MultiBulk (Just [Reply]
replyFlags)
        , Integer Integer
firstKeyPos
        , Integer Integer
lastKeyPos
        , Integer Integer
replyStepCount])) = do
            [Flag]
parsedFlags <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Reply -> Either Reply Flag
parseFlag [Reply]
replyFlags
            LastKeyPositionSpec
lastKey <- Either Reply LastKeyPositionSpec
parseLastKeyPos
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ CommandInfo
                { name :: ByteString
name = ByteString
commandName
                , arity :: AritySpec
arity = Integer -> AritySpec
parseArity Integer
aritySpec
                , flags :: [Flag]
flags = [Flag]
parsedFlags
                , firstKeyPosition :: Integer
firstKeyPosition = Integer
firstKeyPos
                , lastKeyPosition :: LastKeyPositionSpec
lastKeyPosition = LastKeyPositionSpec
lastKey
                , stepCount :: Integer
stepCount = Integer
replyStepCount
                } where
        parseArity :: Integer -> AritySpec
parseArity Integer
int = case Integer
int of
            Integer
i | Integer
i forall a. Ord a => a -> a -> Bool
>= Integer
0 -> Integer -> AritySpec
Required Integer
i
            Integer
i -> Integer -> AritySpec
MinimumRequired forall a b. (a -> b) -> a -> b
$ forall a. Num a => a -> a
abs Integer
i
        parseFlag :: Reply -> Either Reply Flag
        parseFlag :: Reply -> Either Reply Flag
parseFlag (SingleLine ByteString
flag) = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case ByteString
flag of
            ByteString
"write" -> Flag
Write
            ByteString
"readonly" -> Flag
ReadOnly
            ByteString
"denyoom" -> Flag
DenyOOM
            ByteString
"admin" -> Flag
Admin
            ByteString
"pubsub" -> Flag
PubSub
            ByteString
"noscript" -> Flag
NoScript
            ByteString
"random" -> Flag
Random
            ByteString
"sort_for_script" -> Flag
SortForScript
            ByteString
"loading" -> Flag
Loading
            ByteString
"stale" -> Flag
Stale
            ByteString
"skip_monitor" -> Flag
SkipMonitor
            ByteString
"asking" -> Flag
Asking
            ByteString
"fast" -> Flag
Fast
            ByteString
"movablekeys" -> Flag
MovableKeys
            ByteString
other -> ByteString -> Flag
Other ByteString
other
        parseFlag Reply
bad = forall a b. a -> Either a b
Left Reply
bad
        parseLastKeyPos :: Either Reply LastKeyPositionSpec
        parseLastKeyPos :: Either Reply LastKeyPositionSpec
parseLastKeyPos = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Integer
lastKeyPos of
            Integer
i | Integer
i forall a. Ord a => a -> a -> Bool
< Integer
0 -> Integer -> LastKeyPositionSpec
UnlimitedKeys (-Integer
i forall a. Num a => a -> a -> a
- Integer
1)
            Integer
i -> Integer -> LastKeyPositionSpec
LastKeyPosition Integer
i
    -- since redis 6.0
    decode (MultiBulk (Just
        [ name :: Reply
name@(Bulk (Just ByteString
_))
        , arity :: Reply
arity@(Integer Integer
_)
        , flags :: Reply
flags@(MultiBulk (Just [Reply]
_))
        , firstPos :: Reply
firstPos@(Integer Integer
_)
        , lastPos :: Reply
lastPos@(Integer Integer
_)
        , step :: Reply
step@(Integer Integer
_)
        , MultiBulk Maybe [Reply]
_  -- ACL categories
        ])) =
        forall a. RedisResult a => Reply -> Either Reply a
decode (Maybe [Reply] -> Reply
MultiBulk (forall a. a -> Maybe a
Just [Reply
name, Reply
arity, Reply
flags, Reply
firstPos, Reply
lastPos, Reply
step]))

    decode Reply
e = forall a b. a -> Either a b
Left Reply
e

newInfoMap :: [CommandInfo] -> InfoMap
newInfoMap :: [CommandInfo] -> InfoMap
newInfoMap = HashMap [Char] CommandInfo -> InfoMap
InfoMap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
HM.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (\CommandInfo
c -> (ByteString -> [Char]
Char8.unpack forall a b. (a -> b) -> a -> b
$ CommandInfo -> ByteString
name CommandInfo
c, CommandInfo
c))

keysForRequest :: InfoMap -> [BS.ByteString] -> Maybe [BS.ByteString]
keysForRequest :: InfoMap -> [ByteString] -> Maybe [ByteString]
keysForRequest InfoMap
_ [ByteString
"DEBUG", ByteString
"OBJECT", ByteString
key] =
    -- `COMMAND` output for `DEBUG` would let us believe it doesn't have any
    -- keys, but the `DEBUG OBJECT` subcommand does.
    forall a. a -> Maybe a
Just [ByteString
key]
keysForRequest InfoMap
_ [ByteString
"QUIT"] =
    -- The `QUIT` command is not listed in the `COMMAND` output.
    forall a. a -> Maybe a
Just []
keysForRequest (InfoMap HashMap [Char] CommandInfo
infoMap) request :: [ByteString]
request@(ByteString
command:[ByteString]
_) = do
    CommandInfo
info <- forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HM.lookup (forall a b. (a -> b) -> [a] -> [b]
map Char -> Char
toLower forall a b. (a -> b) -> a -> b
$ ByteString -> [Char]
Char8.unpack ByteString
command) HashMap [Char] CommandInfo
infoMap
    CommandInfo -> [ByteString] -> Maybe [ByteString]
keysForRequest' CommandInfo
info [ByteString]
request
keysForRequest InfoMap
_ [] = forall a. Maybe a
Nothing

keysForRequest' :: CommandInfo -> [BS.ByteString] -> Maybe [BS.ByteString]
keysForRequest' :: CommandInfo -> [ByteString] -> Maybe [ByteString]
keysForRequest' CommandInfo
info [ByteString]
request
    | CommandInfo -> Bool
isMovable CommandInfo
info =
        [ByteString] -> Maybe [ByteString]
parseMovable [ByteString]
request
    | CommandInfo -> Integer
stepCount CommandInfo
info forall a. Eq a => a -> a -> Bool
== Integer
0 =
        forall a. a -> Maybe a
Just []
    | Bool
otherwise = do
        let possibleKeys :: [ByteString]
possibleKeys = case CommandInfo -> LastKeyPositionSpec
lastKeyPosition CommandInfo
info of
                LastKeyPosition Integer
end -> forall a. Int -> [a] -> [a]
take (forall a. Enum a => a -> Int
fromEnum forall a b. (a -> b) -> a -> b
$ Integer
1 forall a. Num a => a -> a -> a
+ Integer
end forall a. Num a => a -> a -> a
- CommandInfo -> Integer
firstKeyPosition CommandInfo
info) forall a b. (a -> b) -> a -> b
$ forall a. Int -> [a] -> [a]
drop (forall a. Enum a => a -> Int
fromEnum forall a b. (a -> b) -> a -> b
$ CommandInfo -> Integer
firstKeyPosition CommandInfo
info) [ByteString]
request
                UnlimitedKeys Integer
end ->
                    forall a. Int -> [a] -> [a]
drop (forall a. Enum a => a -> Int
fromEnum forall a b. (a -> b) -> a -> b
$ CommandInfo -> Integer
firstKeyPosition CommandInfo
info) forall a b. (a -> b) -> a -> b
$
                       forall a. Int -> [a] -> [a]
take (forall (t :: * -> *) a. Foldable t => t a -> Int
length [ByteString]
request forall a. Num a => a -> a -> a
- forall a. Enum a => a -> Int
fromEnum Integer
end) [ByteString]
request
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Int -> [a] -> [a]
takeEvery (forall a. Enum a => a -> Int
fromEnum forall a b. (a -> b) -> a -> b
$ CommandInfo -> Integer
stepCount CommandInfo
info) [ByteString]
possibleKeys

isMovable :: CommandInfo -> Bool
isMovable :: CommandInfo -> Bool
isMovable CommandInfo{Integer
[Flag]
ByteString
LastKeyPositionSpec
AritySpec
stepCount :: Integer
lastKeyPosition :: LastKeyPositionSpec
firstKeyPosition :: Integer
flags :: [Flag]
arity :: AritySpec
name :: ByteString
stepCount :: CommandInfo -> Integer
lastKeyPosition :: CommandInfo -> LastKeyPositionSpec
firstKeyPosition :: CommandInfo -> Integer
flags :: CommandInfo -> [Flag]
arity :: CommandInfo -> AritySpec
name :: CommandInfo -> ByteString
..} = Flag
MovableKeys forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Flag]
flags

parseMovable :: [BS.ByteString] -> Maybe [BS.ByteString]
parseMovable :: [ByteString] -> Maybe [ByteString]
parseMovable (ByteString
"SORT":ByteString
key:[ByteString]
_) = forall a. a -> Maybe a
Just [ByteString
key]
parseMovable (ByteString
"EVAL":ByteString
_:[ByteString]
rest) = [ByteString] -> Maybe [ByteString]
readNumKeys [ByteString]
rest
parseMovable (ByteString
"EVALSHA":ByteString
_:[ByteString]
rest) = [ByteString] -> Maybe [ByteString]
readNumKeys [ByteString]
rest
parseMovable (ByteString
"ZUNIONSTORE":ByteString
_:[ByteString]
rest) = [ByteString] -> Maybe [ByteString]
readNumKeys [ByteString]
rest
parseMovable (ByteString
"ZINTERSTORE":ByteString
_:[ByteString]
rest) = [ByteString] -> Maybe [ByteString]
readNumKeys [ByteString]
rest
parseMovable (ByteString
"XREAD":[ByteString]
rest) = [ByteString] -> Maybe [ByteString]
readXreadKeys [ByteString]
rest
parseMovable (ByteString
"XREADGROUP":ByteString
"GROUP":ByteString
_:ByteString
_:[ByteString]
rest) = [ByteString] -> Maybe [ByteString]
readXreadgroupKeys [ByteString]
rest
parseMovable [ByteString]
_ = forall a. Maybe a
Nothing

readXreadKeys :: [BS.ByteString] -> Maybe [BS.ByteString]
readXreadKeys :: [ByteString] -> Maybe [ByteString]
readXreadKeys (ByteString
"COUNT":ByteString
_:[ByteString]
rest) = [ByteString] -> Maybe [ByteString]
readXreadKeys [ByteString]
rest
readXreadKeys (ByteString
"BLOCK":ByteString
_:[ByteString]
rest) = [ByteString] -> Maybe [ByteString]
readXreadKeys [ByteString]
rest
readXreadKeys (ByteString
"STREAMS":[ByteString]
rest) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a. Int -> [a] -> [a]
take (forall (t :: * -> *) a. Foldable t => t a -> Int
length [ByteString]
rest forall a. Integral a => a -> a -> a
`div` Int
2) [ByteString]
rest
readXreadKeys [ByteString]
_ = forall a. Maybe a
Nothing

readXreadgroupKeys :: [BS.ByteString] -> Maybe [BS.ByteString]
readXreadgroupKeys :: [ByteString] -> Maybe [ByteString]
readXreadgroupKeys (ByteString
"COUNT":ByteString
_:[ByteString]
rest) = [ByteString] -> Maybe [ByteString]
readXreadKeys [ByteString]
rest
readXreadgroupKeys (ByteString
"BLOCK":ByteString
_:[ByteString]
rest) = [ByteString] -> Maybe [ByteString]
readXreadKeys [ByteString]
rest
readXreadgroupKeys (ByteString
"NOACK":[ByteString]
rest) = [ByteString] -> Maybe [ByteString]
readXreadKeys [ByteString]
rest
readXreadgroupKeys (ByteString
"STREAMS":[ByteString]
rest) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a. Int -> [a] -> [a]
take (forall (t :: * -> *) a. Foldable t => t a -> Int
length [ByteString]
rest forall a. Integral a => a -> a -> a
`div` Int
2) [ByteString]
rest
readXreadgroupKeys [ByteString]
_ = forall a. Maybe a
Nothing

readNumKeys :: [BS.ByteString] -> Maybe [BS.ByteString]
readNumKeys :: [ByteString] -> Maybe [ByteString]
readNumKeys (ByteString
rawNumKeys:[ByteString]
rest) = do
    Int
numKeys <- forall a. Read a => [Char] -> Maybe a
readMaybe (ByteString -> [Char]
Char8.unpack ByteString
rawNumKeys)
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Int -> [a] -> [a]
take Int
numKeys [ByteString]
rest
readNumKeys [ByteString]
_ = forall a. Maybe a
Nothing
-- takeEvery 1 [1,2,3,4,5] ->[1,2,3,4,5]
-- takeEvery 2 [1,2,3,4,5] ->[1,3,5]
-- takeEvery 3 [1,2,3,4,5] ->[1,4]
takeEvery :: Int -> [a] -> [a]
takeEvery :: forall a. Int -> [a] -> [a]
takeEvery Int
_ [] = []
takeEvery Int
n (a
x:[a]
xs) = a
x forall a. a -> [a] -> [a]
: forall a. Int -> [a] -> [a]
takeEvery Int
n (forall a. Int -> [a] -> [a]
drop (Int
nforall a. Num a => a -> a -> a
-Int
1) [a]
xs)

readMaybe :: Read a => String -> Maybe a
readMaybe :: forall a. Read a => [Char] -> Maybe a
readMaybe [Char]
s = case forall a. Read a => ReadS a
reads [Char]
s of
                  [(a
val, [Char]
"")] -> forall a. a -> Maybe a
Just a
val
                  [(a, [Char])]
_           -> forall a. Maybe a
Nothing