77 from redis .asyncio .cluster import ClusterNode
88
99
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 ):
1157 """
1258 Parses Redis commands to get command keys.
1359 COMMAND output is used to determine key locations.
@@ -30,21 +76,6 @@ def initialize(self, r):
3076 commands [cmd .lower ()] = commands .pop (cmd )
3177 self .commands = commands
3278
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-
4879 # As soon as this PR is merged into Redis, we should reimplement
4980 # our logic to use COMMAND INFO changes to determine the key positions
5081 # https://github.com/redis/redis/pull/8324
@@ -138,37 +169,8 @@ def _get_moveable_keys(self, redis_conn, *args):
138169 raise e
139170 return keys
140171
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
169172
170-
171- class AsyncCommandsParser :
173+ class AsyncCommandsParser (AbstractCommandsParser ):
172174 """
173175 Parses Redis commands to get command keys.
174176
@@ -194,52 +196,75 @@ async def initialize(self, node: Optional["ClusterNode"] = None) -> None:
194196 self .node = node
195197
196198 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 ()}
205200
206201 # As soon as this PR is merged into Redis, we should reimplement
207202 # our logic to use COMMAND INFO changes to determine the key positions
208203 # https://github.com/redis/redis/pull/8324
209204 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+ """
210215 if len (args ) < 2 :
211216 # The command has no keys in it
212217 return None
213218
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,
218222 # 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 :
222229 # We'll try to reinitialize the commands cache, if the engine
223230 # version has changed, the commands may not be current
224231 await self .initialize ()
225232 if cmd_name not in self .commands :
226233 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"
228235 )
229236
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
231255
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 ]
238266
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
243268
244269 async def _get_moveable_keys (self , * args : Any ) -> Optional [Tuple [str , ...]]:
245270 try :
0 commit comments