1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
|
from redis.exceptions import RedisError, ResponseError
from redis.utils import str_if_bytes
class CommandsParser:
"""
Parses Redis commands to get command keys.
COMMAND output is used to determine key locations.
Commands that do not have a predefined key location are flagged with
'movablekeys', and these commands' keys are determined by the command
'COMMAND GETKEYS'.
"""
def __init__(self, redis_connection):
self.commands = {}
self.initialize(redis_connection)
def initialize(self, r):
commands = r.execute_command("COMMAND")
uppercase_commands = []
for cmd in commands:
if any(x.isupper() for x in cmd):
uppercase_commands.append(cmd)
for cmd in uppercase_commands:
commands[cmd.lower()] = commands.pop(cmd)
self.commands = commands
def parse_subcommand(self, command, **options):
cmd_dict = {}
cmd_name = str_if_bytes(command[0])
cmd_dict["name"] = cmd_name
cmd_dict["arity"] = int(command[1])
cmd_dict["flags"] = [str_if_bytes(flag) for flag in command[2]]
cmd_dict["first_key_pos"] = command[3]
cmd_dict["last_key_pos"] = command[4]
cmd_dict["step_count"] = command[5]
if len(command) > 7:
cmd_dict["tips"] = command[7]
cmd_dict["key_specifications"] = command[8]
cmd_dict["subcommands"] = command[9]
return cmd_dict
# As soon as this PR is merged into Redis, we should reimplement
# our logic to use COMMAND INFO changes to determine the key positions
# https://github.com/redis/redis/pull/8324
def get_keys(self, redis_conn, *args):
"""
Get the keys from the passed command.
NOTE: Due to a bug in redis<7.0, this function does not work properly
for EVAL or EVALSHA when the `numkeys` arg is 0.
- issue: https://github.com/redis/redis/issues/9493
- fix: https://github.com/redis/redis/pull/9733
So, don't use this function with EVAL or EVALSHA.
"""
if len(args) < 2:
# The command has no keys in it
return None
cmd_name = args[0].lower()
if cmd_name not in self.commands:
# try to split the command name and to take only the main command,
# e.g. 'memory' for 'memory usage'
cmd_name_split = cmd_name.split()
cmd_name = cmd_name_split[0]
if cmd_name in self.commands:
# save the splitted command to args
args = cmd_name_split + list(args[1:])
else:
# We'll try to reinitialize the commands cache, if the engine
# version has changed, the commands may not be current
self.initialize(redis_conn)
if cmd_name not in self.commands:
raise RedisError(
f"{cmd_name.upper()} command doesn't exist in Redis commands"
)
command = self.commands.get(cmd_name)
if "movablekeys" in command["flags"]:
keys = self._get_moveable_keys(redis_conn, *args)
elif "pubsub" in command["flags"] or command["name"] == "pubsub":
keys = self._get_pubsub_keys(*args)
else:
if (
command["step_count"] == 0
and command["first_key_pos"] == 0
and command["last_key_pos"] == 0
):
is_subcmd = False
if "subcommands" in command:
subcmd_name = f"{cmd_name}|{args[1].lower()}"
for subcmd in command["subcommands"]:
if str_if_bytes(subcmd[0]) == subcmd_name:
command = self.parse_subcommand(subcmd)
is_subcmd = True
# The command doesn't have keys in it
if not is_subcmd:
return None
last_key_pos = command["last_key_pos"]
if last_key_pos < 0:
last_key_pos = len(args) - abs(last_key_pos)
keys_pos = list(
range(command["first_key_pos"], last_key_pos + 1, command["step_count"])
)
keys = [args[pos] for pos in keys_pos]
return keys
def _get_moveable_keys(self, redis_conn, *args):
"""
NOTE: Due to a bug in redis<7.0, this function does not work properly
for EVAL or EVALSHA when the `numkeys` arg is 0.
- issue: https://github.com/redis/redis/issues/9493
- fix: https://github.com/redis/redis/pull/9733
So, don't use this function with EVAL or EVALSHA.
"""
pieces = []
cmd_name = args[0]
# The command name should be splitted into separate arguments,
# e.g. 'MEMORY USAGE' will be splitted into ['MEMORY', 'USAGE']
pieces = pieces + cmd_name.split()
pieces = pieces + list(args[1:])
try:
keys = redis_conn.execute_command("COMMAND GETKEYS", *pieces)
except ResponseError as e:
message = e.__str__()
if (
"Invalid arguments" in message
or "The command has no key arguments" in message
):
return None
else:
raise e
return keys
def _get_pubsub_keys(self, *args):
"""
Get the keys from pubsub command.
Although PubSub commands have predetermined key locations, they are not
supported in the 'COMMAND's output, so the key positions are hardcoded
in this method
"""
if len(args) < 2:
# The command has no keys in it
return None
args = [str_if_bytes(arg) for arg in args]
command = args[0].upper()
keys = None
if command == "PUBSUB":
# the second argument is a part of the command name, e.g.
# ['PUBSUB', 'NUMSUB', 'foo'].
pubsub_type = args[1].upper()
if pubsub_type in ["CHANNELS", "NUMSUB"]:
keys = args[2:]
elif command in ["SUBSCRIBE", "PSUBSCRIBE", "UNSUBSCRIBE", "PUNSUBSCRIBE"]:
# format example:
# SUBSCRIBE channel [channel ...]
keys = list(args[1:])
elif command == "PUBLISH":
# format example:
# PUBLISH channel message
keys = [args[1]]
return keys
|
