7
7
from redis .asyncio .cluster import ClusterNode
8
8
9
9
10
- class CommandsParser :
10
+ class AbstractCommandsParser :
11
+ def _get_pubsub_keys (self , * args ):
12
+ """
13
+ Get the keys from pubsub command.
14
+ Although PubSub commands have predetermined key locations, they are not
15
+ supported in the 'COMMAND's output, so the key positions are hardcoded
16
+ in this method
17
+ """
18
+ if len (args ) < 2 :
19
+ # The command has no keys in it
20
+ return None
21
+ args = [str_if_bytes (arg ) for arg in args ]
22
+ command = args [0 ].upper ()
23
+ keys = None
24
+ if command == "PUBSUB" :
25
+ # the second argument is a part of the command name, e.g.
26
+ # ['PUBSUB', 'NUMSUB', 'foo'].
27
+ pubsub_type = args [1 ].upper ()
28
+ if pubsub_type in ["CHANNELS" , "NUMSUB" , "SHARDCHANNELS" , "SHARDNUMSUB" ]:
29
+ keys = args [2 :]
30
+ elif command in ["SUBSCRIBE" , "PSUBSCRIBE" , "UNSUBSCRIBE" , "PUNSUBSCRIBE" ]:
31
+ # format example:
32
+ # SUBSCRIBE channel [channel ...]
33
+ keys = list (args [1 :])
34
+ elif command in ["PUBLISH" , "SPUBLISH" ]:
35
+ # format example:
36
+ # PUBLISH channel message
37
+ keys = [args [1 ]]
38
+ return keys
39
+
40
+ def parse_subcommand (self , command , ** options ):
41
+ cmd_dict = {}
42
+ cmd_name = str_if_bytes (command [0 ])
43
+ cmd_dict ["name" ] = cmd_name
44
+ cmd_dict ["arity" ] = int (command [1 ])
45
+ cmd_dict ["flags" ] = [str_if_bytes (flag ) for flag in command [2 ]]
46
+ cmd_dict ["first_key_pos" ] = command [3 ]
47
+ cmd_dict ["last_key_pos" ] = command [4 ]
48
+ cmd_dict ["step_count" ] = command [5 ]
49
+ if len (command ) > 7 :
50
+ cmd_dict ["tips" ] = command [7 ]
51
+ cmd_dict ["key_specifications" ] = command [8 ]
52
+ cmd_dict ["subcommands" ] = command [9 ]
53
+ return cmd_dict
54
+
55
+
56
+ class CommandsParser (AbstractCommandsParser ):
11
57
"""
12
58
Parses Redis commands to get command keys.
13
59
COMMAND output is used to determine key locations.
@@ -30,21 +76,6 @@ def initialize(self, r):
30
76
commands [cmd .lower ()] = commands .pop (cmd )
31
77
self .commands = commands
32
78
33
- def parse_subcommand (self , command , ** options ):
34
- cmd_dict = {}
35
- cmd_name = str_if_bytes (command [0 ])
36
- cmd_dict ["name" ] = cmd_name
37
- cmd_dict ["arity" ] = int (command [1 ])
38
- cmd_dict ["flags" ] = [str_if_bytes (flag ) for flag in command [2 ]]
39
- cmd_dict ["first_key_pos" ] = command [3 ]
40
- cmd_dict ["last_key_pos" ] = command [4 ]
41
- cmd_dict ["step_count" ] = command [5 ]
42
- if len (command ) > 7 :
43
- cmd_dict ["tips" ] = command [7 ]
44
- cmd_dict ["key_specifications" ] = command [8 ]
45
- cmd_dict ["subcommands" ] = command [9 ]
46
- return cmd_dict
47
-
48
79
# As soon as this PR is merged into Redis, we should reimplement
49
80
# our logic to use COMMAND INFO changes to determine the key positions
50
81
# https://github.com/redis/redis/pull/8324
@@ -138,37 +169,8 @@ def _get_moveable_keys(self, redis_conn, *args):
138
169
raise e
139
170
return keys
140
171
141
- def _get_pubsub_keys (self , * args ):
142
- """
143
- Get the keys from pubsub command.
144
- Although PubSub commands have predetermined key locations, they are not
145
- supported in the 'COMMAND's output, so the key positions are hardcoded
146
- in this method
147
- """
148
- if len (args ) < 2 :
149
- # The command has no keys in it
150
- return None
151
- args = [str_if_bytes (arg ) for arg in args ]
152
- command = args [0 ].upper ()
153
- keys = None
154
- if command == "PUBSUB" :
155
- # the second argument is a part of the command name, e.g.
156
- # ['PUBSUB', 'NUMSUB', 'foo'].
157
- pubsub_type = args [1 ].upper ()
158
- if pubsub_type in ["CHANNELS" , "NUMSUB" , "SHARDCHANNELS" , "SHARDNUMSUB" ]:
159
- keys = args [2 :]
160
- elif command in ["SUBSCRIBE" , "PSUBSCRIBE" , "UNSUBSCRIBE" , "PUNSUBSCRIBE" ]:
161
- # format example:
162
- # SUBSCRIBE channel [channel ...]
163
- keys = list (args [1 :])
164
- elif command in ["PUBLISH" , "SPUBLISH" ]:
165
- # format example:
166
- # PUBLISH channel message
167
- keys = [args [1 ]]
168
- return keys
169
172
170
-
171
- class AsyncCommandsParser :
173
+ class AsyncCommandsParser (AbstractCommandsParser ):
172
174
"""
173
175
Parses Redis commands to get command keys.
174
176
@@ -194,52 +196,75 @@ async def initialize(self, node: Optional["ClusterNode"] = None) -> None:
194
196
self .node = node
195
197
196
198
commands = await self .node .execute_command ("COMMAND" )
197
- for cmd , command in commands .items ():
198
- if "movablekeys" in command ["flags" ]:
199
- commands [cmd ] = - 1
200
- elif command ["first_key_pos" ] == 0 and command ["last_key_pos" ] == 0 :
201
- commands [cmd ] = 0
202
- elif command ["first_key_pos" ] == 1 and command ["last_key_pos" ] == 1 :
203
- commands [cmd ] = 1
204
- self .commands = {cmd .upper (): command for cmd , command in commands .items ()}
199
+ self .commands = {cmd .lower (): command for cmd , command in commands .items ()}
205
200
206
201
# As soon as this PR is merged into Redis, we should reimplement
207
202
# our logic to use COMMAND INFO changes to determine the key positions
208
203
# https://github.com/redis/redis/pull/8324
209
204
async def get_keys (self , * args : Any ) -> Optional [Tuple [str , ...]]:
205
+ """
206
+ Get the keys from the passed command.
207
+
208
+ NOTE: Due to a bug in redis<7.0, this function does not work properly
209
+ for EVAL or EVALSHA when the `numkeys` arg is 0.
210
+ - issue: https://github.com/redis/redis/issues/9493
211
+ - fix: https://github.com/redis/redis/pull/9733
212
+
213
+ So, don't use this function with EVAL or EVALSHA.
214
+ """
210
215
if len (args ) < 2 :
211
216
# The command has no keys in it
212
217
return None
213
218
214
- try :
215
- command = self .commands [args [0 ]]
216
- except KeyError :
217
- # try to split the command name and to take only the main command
219
+ cmd_name = args [0 ].lower ()
220
+ if cmd_name not in self .commands :
221
+ # try to split the command name and to take only the main command,
218
222
# e.g. 'memory' for 'memory usage'
219
- args = args [0 ].split () + list (args [1 :])
220
- cmd_name = args [0 ].upper ()
221
- if cmd_name not in self .commands :
223
+ cmd_name_split = cmd_name .split ()
224
+ cmd_name = cmd_name_split [0 ]
225
+ if cmd_name in self .commands :
226
+ # save the splitted command to args
227
+ args = cmd_name_split + list (args [1 :])
228
+ else :
222
229
# We'll try to reinitialize the commands cache, if the engine
223
230
# version has changed, the commands may not be current
224
231
await self .initialize ()
225
232
if cmd_name not in self .commands :
226
233
raise RedisError (
227
- f"{ cmd_name } command doesn't exist in Redis commands"
234
+ f"{ cmd_name . upper () } command doesn't exist in Redis commands"
228
235
)
229
236
230
- command = self .commands [cmd_name ]
237
+ command = self .commands .get (cmd_name )
238
+ if "movablekeys" in command ["flags" ]:
239
+ keys = await self ._get_moveable_keys (* args )
240
+ elif "pubsub" in command ["flags" ] or command ["name" ] == "pubsub" :
241
+ keys = self ._get_pubsub_keys (* args )
242
+ else :
243
+ if (
244
+ command ["step_count" ] == 0
245
+ and command ["first_key_pos" ] == 0
246
+ and command ["last_key_pos" ] == 0
247
+ ):
248
+ is_subcmd = False
249
+ if "subcommands" in command :
250
+ subcmd_name = f"{ cmd_name } |{ args [1 ].lower ()} "
251
+ for subcmd in command ["subcommands" ]:
252
+ if str_if_bytes (subcmd [0 ]) == subcmd_name :
253
+ command = self .parse_subcommand (subcmd )
254
+ is_subcmd = True
231
255
232
- if command == 1 :
233
- return (args [1 ],)
234
- if command == 0 :
235
- return None
236
- if command == - 1 :
237
- return await self ._get_moveable_keys (* args )
256
+ # The command doesn't have keys in it
257
+ if not is_subcmd :
258
+ return None
259
+ last_key_pos = command ["last_key_pos" ]
260
+ if last_key_pos < 0 :
261
+ last_key_pos = len (args ) - abs (last_key_pos )
262
+ keys_pos = list (
263
+ range (command ["first_key_pos" ], last_key_pos + 1 , command ["step_count" ])
264
+ )
265
+ keys = [args [pos ] for pos in keys_pos ]
238
266
239
- last_key_pos = command ["last_key_pos" ]
240
- if last_key_pos < 0 :
241
- last_key_pos = len (args ) + last_key_pos
242
- return args [command ["first_key_pos" ] : last_key_pos + 1 : command ["step_count" ]]
267
+ return keys
243
268
244
269
async def _get_moveable_keys (self , * args : Any ) -> Optional [Tuple [str , ...]]:
245
270
try :
0 commit comments