{-# 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 -> String
(Int -> Flag -> ShowS)
-> (Flag -> String) -> ([Flag] -> ShowS) -> Show Flag
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Flag] -> ShowS
$cshowList :: [Flag] -> ShowS
show :: Flag -> String
$cshow :: Flag -> String
showsPrec :: Int -> Flag -> ShowS
$cshowsPrec :: Int -> Flag -> ShowS
Show, Flag -> Flag -> Bool
(Flag -> Flag -> Bool) -> (Flag -> Flag -> Bool) -> Eq Flag
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 -> String
(Int -> AritySpec -> ShowS)
-> (AritySpec -> String)
-> ([AritySpec] -> ShowS)
-> Show AritySpec
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [AritySpec] -> ShowS
$cshowList :: [AritySpec] -> ShowS
show :: AritySpec -> String
$cshow :: AritySpec -> String
showsPrec :: Int -> AritySpec -> ShowS
$cshowsPrec :: Int -> AritySpec -> ShowS
Show)

data LastKeyPositionSpec = LastKeyPosition Integer | UnlimitedKeys Integer deriving (Int -> LastKeyPositionSpec -> ShowS
[LastKeyPositionSpec] -> ShowS
LastKeyPositionSpec -> String
(Int -> LastKeyPositionSpec -> ShowS)
-> (LastKeyPositionSpec -> String)
-> ([LastKeyPositionSpec] -> ShowS)
-> Show LastKeyPositionSpec
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [LastKeyPositionSpec] -> ShowS
$cshowList :: [LastKeyPositionSpec] -> ShowS
show :: LastKeyPositionSpec -> String
$cshow :: LastKeyPositionSpec -> String
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 -> String
(Int -> CommandInfo -> ShowS)
-> (CommandInfo -> String)
-> ([CommandInfo] -> ShowS)
-> Show CommandInfo
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [CommandInfo] -> ShowS
$cshowList :: [CommandInfo] -> ShowS
show :: CommandInfo -> String
$cshow :: CommandInfo -> String
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 <- (Reply -> Either Reply Flag) -> [Reply] -> Either Reply [Flag]
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
            CommandInfo -> Either Reply CommandInfo
forall (m :: * -> *) a. Monad m => a -> m a
return (CommandInfo -> Either Reply CommandInfo)
-> CommandInfo -> Either Reply CommandInfo
forall a b. (a -> b) -> a -> b
$ CommandInfo :: ByteString
-> AritySpec
-> [Flag]
-> Integer
-> LastKeyPositionSpec
-> Integer
-> CommandInfo
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 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
0 -> Integer -> AritySpec
Required Integer
i
            Integer
i -> Integer -> AritySpec
MinimumRequired (Integer -> AritySpec) -> Integer -> AritySpec
forall a b. (a -> b) -> a -> b
$ Integer -> Integer
forall a. Num a => a -> a
abs Integer
i
        parseFlag :: Reply -> Either Reply Flag
        parseFlag :: Reply -> Either Reply Flag
parseFlag (SingleLine ByteString
flag) = Flag -> Either Reply Flag
forall (m :: * -> *) a. Monad m => a -> m a
return (Flag -> Either Reply Flag) -> Flag -> Either Reply Flag
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 = Reply -> Either Reply Flag
forall a b. a -> Either a b
Left Reply
bad
        parseLastKeyPos :: Either Reply LastKeyPositionSpec
        parseLastKeyPos :: Either Reply LastKeyPositionSpec
parseLastKeyPos = LastKeyPositionSpec -> Either Reply LastKeyPositionSpec
forall (m :: * -> *) a. Monad m => a -> m a
return (LastKeyPositionSpec -> Either Reply LastKeyPositionSpec)
-> LastKeyPositionSpec -> Either Reply LastKeyPositionSpec
forall a b. (a -> b) -> a -> b
$ case Integer
lastKeyPos of
            Integer
i | Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 -> Integer -> LastKeyPositionSpec
UnlimitedKeys (-Integer
i Integer -> Integer -> Integer
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
        ])) =
        Reply -> Either Reply CommandInfo
forall a. RedisResult a => Reply -> Either Reply a
decode (Maybe [Reply] -> Reply
MultiBulk ([Reply] -> Maybe [Reply]
forall a. a -> Maybe a
Just [Reply
name, Reply
arity, Reply
flags, Reply
firstPos, Reply
lastPos, Reply
step]))

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

newInfoMap :: [CommandInfo] -> InfoMap
newInfoMap :: [CommandInfo] -> InfoMap
newInfoMap = HashMap String CommandInfo -> InfoMap
InfoMap (HashMap String CommandInfo -> InfoMap)
-> ([CommandInfo] -> HashMap String CommandInfo)
-> [CommandInfo]
-> InfoMap
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(String, CommandInfo)] -> HashMap String CommandInfo
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
HM.fromList ([(String, CommandInfo)] -> HashMap String CommandInfo)
-> ([CommandInfo] -> [(String, CommandInfo)])
-> [CommandInfo]
-> HashMap String CommandInfo
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (CommandInfo -> (String, CommandInfo))
-> [CommandInfo] -> [(String, CommandInfo)]
forall a b. (a -> b) -> [a] -> [b]
map (\CommandInfo
c -> (ByteString -> String
Char8.unpack (ByteString -> String) -> ByteString -> String
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.
    [ByteString] -> Maybe [ByteString]
forall a. a -> Maybe a
Just [ByteString
key]
keysForRequest InfoMap
_ [ByteString
"QUIT"] =
    -- The `QUIT` command is not listed in the `COMMAND` output.
    [ByteString] -> Maybe [ByteString]
forall a. a -> Maybe a
Just []
keysForRequest (InfoMap HashMap String CommandInfo
infoMap) request :: [ByteString]
request@(ByteString
command:[ByteString]
_) = do
    CommandInfo
info <- String -> HashMap String CommandInfo -> Maybe CommandInfo
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HM.lookup ((Char -> Char) -> ShowS
forall a b. (a -> b) -> [a] -> [b]
map Char -> Char
toLower ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$ ByteString -> String
Char8.unpack ByteString
command) HashMap String CommandInfo
infoMap
    CommandInfo -> [ByteString] -> Maybe [ByteString]
keysForRequest' CommandInfo
info [ByteString]
request
keysForRequest InfoMap
_ [] = Maybe [ByteString]
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 Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 =
        [ByteString] -> Maybe [ByteString]
forall a. a -> Maybe a
Just []
    | Bool
otherwise = do
        let possibleKeys :: [ByteString]
possibleKeys = case CommandInfo -> LastKeyPositionSpec
lastKeyPosition CommandInfo
info of
                LastKeyPosition Integer
end -> Int -> [ByteString] -> [ByteString]
forall a. Int -> [a] -> [a]
take (Integer -> Int
forall a. Enum a => a -> Int
fromEnum (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ Integer
1 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
end Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- CommandInfo -> Integer
firstKeyPosition CommandInfo
info) ([ByteString] -> [ByteString]) -> [ByteString] -> [ByteString]
forall a b. (a -> b) -> a -> b
$ Int -> [ByteString] -> [ByteString]
forall a. Int -> [a] -> [a]
drop (Integer -> Int
forall a. Enum a => a -> Int
fromEnum (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ CommandInfo -> Integer
firstKeyPosition CommandInfo
info) [ByteString]
request
                UnlimitedKeys Integer
end ->
                    Int -> [ByteString] -> [ByteString]
forall a. Int -> [a] -> [a]
drop (Integer -> Int
forall a. Enum a => a -> Int
fromEnum (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ CommandInfo -> Integer
firstKeyPosition CommandInfo
info) ([ByteString] -> [ByteString]) -> [ByteString] -> [ByteString]
forall a b. (a -> b) -> a -> b
$
                       Int -> [ByteString] -> [ByteString]
forall a. Int -> [a] -> [a]
take ([ByteString] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ByteString]
request Int -> Int -> Int
forall a. Num a => a -> a -> a
- Integer -> Int
forall a. Enum a => a -> Int
fromEnum Integer
end) [ByteString]
request
        [ByteString] -> Maybe [ByteString]
forall (m :: * -> *) a. Monad m => a -> m a
return ([ByteString] -> Maybe [ByteString])
-> [ByteString] -> Maybe [ByteString]
forall a b. (a -> b) -> a -> b
$ Int -> [ByteString] -> [ByteString]
forall a. Int -> [a] -> [a]
takeEvery (Integer -> Int
forall a. Enum a => a -> Int
fromEnum (Integer -> Int) -> Integer -> Int
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 Flag -> [Flag] -> Bool
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]
_) = [ByteString] -> Maybe [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]
_ = Maybe [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) = [ByteString] -> Maybe [ByteString]
forall a. a -> Maybe a
Just ([ByteString] -> Maybe [ByteString])
-> [ByteString] -> Maybe [ByteString]
forall a b. (a -> b) -> a -> b
$ Int -> [ByteString] -> [ByteString]
forall a. Int -> [a] -> [a]
take ([ByteString] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ByteString]
rest Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) [ByteString]
rest
readXreadKeys [ByteString]
_ = Maybe [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) = [ByteString] -> Maybe [ByteString]
forall a. a -> Maybe a
Just ([ByteString] -> Maybe [ByteString])
-> [ByteString] -> Maybe [ByteString]
forall a b. (a -> b) -> a -> b
$ Int -> [ByteString] -> [ByteString]
forall a. Int -> [a] -> [a]
take ([ByteString] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ByteString]
rest Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) [ByteString]
rest
readXreadgroupKeys [ByteString]
_ = Maybe [ByteString]
forall a. Maybe a
Nothing

readNumKeys :: [BS.ByteString] -> Maybe [BS.ByteString]
readNumKeys :: [ByteString] -> Maybe [ByteString]
readNumKeys (ByteString
rawNumKeys:[ByteString]
rest) = do
    Int
numKeys <- String -> Maybe Int
forall a. Read a => String -> Maybe a
readMaybe (ByteString -> String
Char8.unpack ByteString
rawNumKeys)
    [ByteString] -> Maybe [ByteString]
forall (m :: * -> *) a. Monad m => a -> m a
return ([ByteString] -> Maybe [ByteString])
-> [ByteString] -> Maybe [ByteString]
forall a b. (a -> b) -> a -> b
$ Int -> [ByteString] -> [ByteString]
forall a. Int -> [a] -> [a]
take Int
numKeys [ByteString]
rest
readNumKeys [ByteString]
_ = Maybe [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 :: Int -> [a] -> [a]
takeEvery Int
_ [] = []
takeEvery Int
n (a
x:[a]
xs) = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
takeEvery Int
n (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) [a]
xs)

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