diff --git a/benchmarks/command_packer_benchmark.py b/benchmarks/command_packer_benchmark.py index 4fb7196422..4a2cf9d4e5 100644 --- a/benchmarks/command_packer_benchmark.py +++ b/benchmarks/command_packer_benchmark.py @@ -9,6 +9,8 @@ def send_packed_command(self, command, check_health=True): if not self._sock: self.connect() try: + if isinstance(command, list): + command = SYM_EMPTY.join(command) self._sock.sendall(command) except OSError as e: self.disconnect() diff --git a/redis/_parsers/commands.py b/redis/_parsers/commands.py index 0d9dc40bd1..f7ad4ef9b6 100644 --- a/redis/_parsers/commands.py +++ b/redis/_parsers/commands.py @@ -114,6 +114,37 @@ def initialize(self, r): for cmd in uppercase_commands: commands[cmd.lower()] = commands.pop(cmd) self.commands = commands + for command in self.commands.values(): + first_key_pos = command["first_key_pos"] + last_key_pos = command["last_key_pos"] + flags = command["flags"] + if ( + first_key_pos > 0 + and first_key_pos == last_key_pos + and "movablekeys" not in flags + and "pubsub" not in flags + and command["name"] != "pubsub" + ): + command["_single_key_pos"] = first_key_pos + + def _is_keyed_command(self, *args): + """ + Determines whether the command is always keyed, never keyless. + + Must have been initialized, won't automatically do so. + """ + if len(args) < 2: + # The command has no keys in it + return False + + cmd_name = args[0].lower() + commands = self.commands + command = commands.get(cmd_name) + if command is None: + return False + + single_pos = command.get("_single_key_pos") + return single_pos is not None # As soon as this PR is merged into Redis, we should reimplement # our logic to use COMMAND INFO changes to determine the key positions @@ -134,34 +165,41 @@ def get_keys(self, redis_conn, *args): return None cmd_name = args[0].lower() - if cmd_name not in self.commands: + commands = self.commands + command = commands.get(cmd_name) + if command is None: # 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: + if cmd_name in commands: # save the split 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: + commands = self.commands + if cmd_name not in commands: raise RedisError( f"{cmd_name.upper()} command doesn't exist in Redis commands" ) + command = commands.get(cmd_name) - command = self.commands.get(cmd_name) - if "movablekeys" in command["flags"]: + single_pos = command.get("_single_key_pos") + if single_pos is not None: + return [args[single_pos]] + + flags = command["flags"] + if "movablekeys" in flags: keys = self._get_moveable_keys(redis_conn, *args) - elif "pubsub" in command["flags"] or command["name"] == "pubsub": + elif "pubsub" in 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 - ): + step_count = command["step_count"] + first_key_pos = command["first_key_pos"] + last_key_pos = command["last_key_pos"] + if step_count == 0 and first_key_pos == 0 and last_key_pos == 0: is_subcmd = False if "subcommands" in command: subcmd_name = f"{cmd_name}|{args[1].lower()}" @@ -169,19 +207,23 @@ def get_keys(self, redis_conn, *args): if str_if_bytes(subcmd[0]) == subcmd_name: command = self.parse_subcommand(subcmd) - if command["first_key_pos"] > 0: + step_count = command["step_count"] + first_key_pos = command["first_key_pos"] + last_key_pos = command["last_key_pos"] + + if first_key_pos > 0: 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"]) + last_key_pos += len(args) + keys = list( + map( + args.__getitem__, range(first_key_pos, last_key_pos + 1, step_count) + ) ) - keys = [args[pos] for pos in keys_pos] return keys @@ -421,7 +463,22 @@ async def initialize(self, node: Optional["ClusterNode"] = None) -> None: self.node = node commands = await self.node.execute_command("COMMAND") - self.commands = {cmd.lower(): command for cmd, command in commands.items()} + commands = {cmd.lower(): command for cmd, command in commands.items()} + + for command in commands.values(): + first_key_pos = command["first_key_pos"] + last_key_pos = command["last_key_pos"] + flags = command["flags"] + if ( + first_key_pos > 0 + and first_key_pos == last_key_pos + and "movablekeys" not in flags + and "pubsub" not in flags + and command["name"] != "pubsub" + ): + command["_single_key_pos"] = first_key_pos + + self.commands = commands # As soon as this PR is merged into Redis, we should reimplement # our logic to use COMMAND INFO changes to determine the key positions @@ -442,34 +499,41 @@ async def get_keys(self, *args: Any) -> Optional[Tuple[str, ...]]: return None cmd_name = args[0].lower() - if cmd_name not in self.commands: + commands = self.commands + command = commands.get(cmd_name) + if command is None: # 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: + if cmd_name in commands: # save the split 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 await self.initialize() - if cmd_name not in self.commands: + commands = self.commands + if cmd_name not in commands: raise RedisError( f"{cmd_name.upper()} command doesn't exist in Redis commands" ) + command = commands.get(cmd_name) + + single_pos = command.get("_single_key_pos") + if single_pos is not None: + return [args[single_pos]] - command = self.commands.get(cmd_name) - if "movablekeys" in command["flags"]: + flags = command["flags"] + if "movablekeys" in flags: keys = await self._get_moveable_keys(*args) - elif "pubsub" in command["flags"] or command["name"] == "pubsub": + elif "pubsub" in 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 - ): + step_count = command["step_count"] + first_key_pos = command["first_key_pos"] + last_key_pos = command["last_key_pos"] + if step_count == 0 and first_key_pos == 0 and last_key_pos == 0: is_subcmd = False if "subcommands" in command: subcmd_name = f"{cmd_name}|{args[1].lower()}" @@ -477,19 +541,20 @@ async def get_keys(self, *args: Any) -> Optional[Tuple[str, ...]]: if str_if_bytes(subcmd[0]) == subcmd_name: command = self.parse_subcommand(subcmd) - if command["first_key_pos"] > 0: + first_key_pos = command["first_key_pos"] + last_key_pos = command["last_key_pos"] + step_count = command["step_count"] + if first_key_pos > 0: 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] + last_key_pos += len(args) + keys = [ + args[pos] for pos in range(first_key_pos, last_key_pos + 1, step_count) + ] return keys diff --git a/redis/_parsers/encoders.py b/redis/_parsers/encoders.py index 0275ce0756..82551f31ba 100644 --- a/redis/_parsers/encoders.py +++ b/redis/_parsers/encoders.py @@ -15,6 +15,8 @@ def encode(self, value): "Return a bytestring or bytes-like representation of the value" if isinstance(value, (bytes, bytearray, memoryview)): return value + elif isinstance(value, str): + return value.encode(self.encoding, self.encoding_errors) elif isinstance(value, bool): # special case bool since it is a subclass of int raise DataError( @@ -22,17 +24,14 @@ def encode(self, value): "bytes, string, int or float first." ) elif isinstance(value, (int, float)): - value = repr(value).encode() - elif not isinstance(value, str): + return repr(value).encode() + else: # a value we don't know how to deal with. throw an error typename = type(value).__name__ raise DataError( f"Invalid input of type: '{typename}'. " f"Convert to a bytes, string, int or float first." ) - if isinstance(value, str): - value = value.encode(self.encoding, self.encoding_errors) - return value def decode(self, value, force=False): "Return a unicode string from the bytes-like representation" diff --git a/redis/_parsers/helpers.py b/redis/_parsers/helpers.py index d9df7a4d30..4e00d6fdf8 100644 --- a/redis/_parsers/helpers.py +++ b/redis/_parsers/helpers.py @@ -1,4 +1,4 @@ -import datetime +from datetime import datetime from redis.utils import str_if_bytes @@ -11,7 +11,7 @@ def timestamp_to_datetime(response): response = int(response) except ValueError: return None - return datetime.datetime.fromtimestamp(response) + return datetime.fromtimestamp(response) def parse_debug_object(response): @@ -875,7 +875,7 @@ def float_or_none(response): def bool_ok(response, **options): - return str_if_bytes(response) == "OK" + return response in (b"OK", "OK") def parse_zadd(response, **options): @@ -1259,7 +1259,7 @@ def parse_pubsub_numsub(response, **options): def parse_client_kill(response, **options): if isinstance(response, int): return response - return str_if_bytes(response) == "OK" + return bool_ok(response) def parse_acl_getuser(response, **options): @@ -1451,7 +1451,7 @@ def parse_set_result(response, **options): # Redis will return a getCommand result. # See `setGenericCommand` in t_string.c return response - return response and str_if_bytes(response) == "OK" + return response and bool_ok(response) def parse_function_list_unified(response, **options): diff --git a/redis/_parsers/hiredis.py b/redis/_parsers/hiredis.py index f88f667a1c..47faa1f9bf 100644 --- a/redis/_parsers/hiredis.py +++ b/redis/_parsers/hiredis.py @@ -349,17 +349,18 @@ async def read_response( if not self._connected: raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) from None + reader = self._reader if disable_decoding: - response = self._reader.gets(False) + response = reader.gets(False) else: - response = self._reader.gets() + response = reader.gets() while response is NOT_ENOUGH_DATA: await self.read_from_socket() if disable_decoding: - response = self._reader.gets(False) + response = reader.gets(False) else: - response = self._reader.gets() + response = reader.gets() # if the response is a ConnectionError or the response is a list and # the first item is a ConnectionError, raise it as something bad diff --git a/redis/client.py b/redis/client.py index 139306d7f6..9ae105eb92 100755 --- a/redis/client.py +++ b/redis/client.py @@ -927,22 +927,26 @@ def failure_callback(error, failure_count): def parse_response(self, connection, command_name, **options): """Parses a response from the Redis server""" - try: - if NEVER_DECODE in options: - response = connection.read_response(disable_decoding=True) - options.pop(NEVER_DECODE) - else: - response = connection.read_response() - except ResponseError: - if EMPTY_RESPONSE in options: - return options[EMPTY_RESPONSE] - raise + if not options: + # Fast-path for the common case of no options + response = connection.read_response() + else: + try: + if NEVER_DECODE in options: + response = connection.read_response(disable_decoding=True) + options.pop(NEVER_DECODE) + else: + response = connection.read_response() + except ResponseError: + if EMPTY_RESPONSE in options: + return options[EMPTY_RESPONSE] + raise - if EMPTY_RESPONSE in options: - options.pop(EMPTY_RESPONSE) + if EMPTY_RESPONSE in options: + options.pop(EMPTY_RESPONSE) - # Remove keys entry, it needs only for cache. - options.pop("keys", None) + # Remove keys entry, it needs only for cache. + options.pop("keys", None) if command_name in self.response_callbacks: return self.response_callbacks[command_name](response, **options) diff --git a/redis/cluster.py b/redis/cluster.py index 89c61f95bb..f125cc9cde 100644 --- a/redis/cluster.py +++ b/redis/cluster.py @@ -959,23 +959,36 @@ def __init__( self._policies_callback_mapping: dict[ Union[RequestPolicy, ResponsePolicy], Callable ] = { - RequestPolicy.DEFAULT_KEYLESS: lambda command_name: [ - self.get_random_primary_or_all_nodes(command_name) + RequestPolicy.DEFAULT_KEYLESS: lambda self, command, *args, **kwargs: [ + self.get_random_primary_or_all_nodes(command) ], - RequestPolicy.DEFAULT_KEYED: lambda command, - *args: self.get_nodes_from_slot(command, *args), - RequestPolicy.DEFAULT_NODE: lambda: [self.get_default_node()], - RequestPolicy.ALL_SHARDS: self.get_primaries, - RequestPolicy.ALL_NODES: self.get_nodes, - RequestPolicy.ALL_REPLICAS: self.get_replicas, - RequestPolicy.MULTI_SHARD: lambda *args, - **kwargs: self._split_multi_shard_command(*args, **kwargs), - RequestPolicy.SPECIAL: self.get_special_nodes, + RequestPolicy.DEFAULT_KEYED: lambda self, command, *args, **kwargs: ( + self.get_nodes_from_slot(command, *args) + ), + RequestPolicy.DEFAULT_NODE: lambda self, command, *args, **kwargs: [ + self.get_default_node() + ], + RequestPolicy.ALL_SHARDS: lambda self, command, *args, **kwargs: ( + self.get_primaries() + ), + RequestPolicy.ALL_NODES: lambda self, command, *args, **kwargs: ( + self.get_nodes() + ), + RequestPolicy.ALL_REPLICAS: lambda self, command, *args, **kwargs: ( + self.get_replicas() + ), + RequestPolicy.MULTI_SHARD: lambda self, command, *args, **kwargs: ( + self._split_multi_shard_command(*args, **kwargs) + ), + RequestPolicy.SPECIAL: lambda self, command, *args, **kwargs: ( + self.get_special_nodes() + ), ResponsePolicy.DEFAULT_KEYLESS: lambda res: res, ResponsePolicy.DEFAULT_KEYED: lambda res: res, } self._policy_resolver = policy_resolver + self._policy_cb_cache = {} self.commands_parser = CommandsParser(self) # Node where FT.AGGREGATE command is executed. @@ -1089,11 +1102,12 @@ def get_nodes_from_slot(self, command: str, *args): Returns a list of nodes that hold the specified keys' slots. """ # get the node that holds the key's slot + is_read = command in READ_COMMANDS slot = self.determine_slot(*args) node = self.nodes_manager.get_node_from_slot( slot, - self.read_from_replicas and command in READ_COMMANDS, - self.load_balancing_strategy if command in READ_COMMANDS else None, + self.read_from_replicas and is_read, + self.load_balancing_strategy if is_read else None, ) return [node] @@ -1339,33 +1353,40 @@ def _determine_nodes( """ Determines a nodes the command should be executed on. """ - command = args[0].upper() - if len(args) >= 2 and f"{args[0]} {args[1]}".upper() in self.command_flags: - command = f"{args[0]} {args[1]}".upper() - + arg0 = args[0] nodes_flag = kwargs.pop("nodes_flag", None) - if nodes_flag is not None: - # nodes flag passed by the user - command_flag = nodes_flag + if nodes_flag is None: + policy_cb = self._policy_cb_cache.get(arg0) else: - # get the nodes group for this command if it was predefined - command_flag = self.command_flags.get(command) - - if command_flag in self._command_flags_mapping: - request_policy = self._command_flags_mapping[command_flag] + policy_cb = None + if policy_cb is None: + command = arg0.upper() + if len(args) >= 2 and f"{arg0} {args[1]}".upper() in self.command_flags: + command = f"{arg0} {args[1]}".upper() + + if nodes_flag is not None: + # nodes flag passed by the user + command_flag = nodes_flag + else: + # get the nodes group for this command if it was predefined + command_flag = self.command_flags.get(command) - policy_callback = self._policies_callback_mapping[request_policy] + request_policy = self._command_flags_mapping.get( + command_flag, request_policy + ) - if request_policy == RequestPolicy.DEFAULT_KEYED: - nodes = policy_callback(command, *args) - elif request_policy == RequestPolicy.MULTI_SHARD: - nodes = policy_callback(*args, **kwargs) - elif request_policy == RequestPolicy.DEFAULT_KEYLESS: - nodes = policy_callback(args[0]) + policy_cb = self._policies_callback_mapping[request_policy] + if nodes_flag is None and command == arg0: + if len(self._policy_cb_cache) > 5000: + # Prevent unbounded memory leak on abnormal use + self._policy_cb_cache.clear() + self._policy_cb_cache[arg0] = policy_cb else: - nodes = policy_callback() + command = arg0 + + nodes = policy_cb(self, command, *args, **kwargs) - if args[0].lower() == "ft.aggregate": + if arg0.lower() == "ft.aggregate": self._aggregate_nodes = nodes return nodes @@ -1451,14 +1472,15 @@ def determine_slot(self, *args) -> Optional[int]: # CLIENT TRACKING is a special case. # It doesn't have any keys, it needs to be sent to the provided nodes # By default it will be sent to all nodes. - if command.upper() == "CLIENT TRACKING": + commandu = command.upper() + if commandu == "CLIENT TRACKING": return None # EVAL and EVALSHA are common enough that it's wasteful to go to the # redis server to parse the keys. Besides, there is a bug in redis<7.0 # where `self._get_command_keys()` fails anyway. So, we special case # EVAL/EVALSHA. - if command.upper() in ("EVAL", "EVALSHA"): + if commandu in ("EVAL", "EVALSHA"): # command syntax: EVAL "script body" num_keys ... if len(args) <= 2: raise RedisClusterException(f"Invalid args in command: {args}") @@ -1471,10 +1493,10 @@ def determine_slot(self, *args) -> Optional[int]: keys = eval_keys else: keys = self._get_command_keys(*args) - if keys is None or len(keys) == 0: + if not keys: # FCALL can call a function with 0 keys, that means the function # can be run on any node so we can just return a random slot - if command.upper() in ("FCALL", "FCALL_RO"): + if commandu in ("FCALL", "FCALL_RO"): return random.randrange(0, REDIS_CLUSTER_HASH_SLOTS) raise RedisClusterException( "No way to dispatch this command to Redis Cluster. " @@ -3610,23 +3632,36 @@ def __init__( self._policies_callback_mapping: dict[ Union[RequestPolicy, ResponsePolicy], Callable ] = { - RequestPolicy.DEFAULT_KEYLESS: lambda command_name: [ - self.get_random_primary_or_all_nodes(command_name) + RequestPolicy.DEFAULT_KEYLESS: lambda self, command, *args, **kwargs: [ + self.get_random_primary_or_all_nodes(command) ], - RequestPolicy.DEFAULT_KEYED: lambda command, - *args: self.get_nodes_from_slot(command, *args), - RequestPolicy.DEFAULT_NODE: lambda: [self.get_default_node()], - RequestPolicy.ALL_SHARDS: self.get_primaries, - RequestPolicy.ALL_NODES: self.get_nodes, - RequestPolicy.ALL_REPLICAS: self.get_replicas, - RequestPolicy.MULTI_SHARD: lambda *args, - **kwargs: self._split_multi_shard_command(*args, **kwargs), - RequestPolicy.SPECIAL: self.get_special_nodes, + RequestPolicy.DEFAULT_KEYED: lambda self, command, *args, **kwargs: ( + self.get_nodes_from_slot(command, *args) + ), + RequestPolicy.DEFAULT_NODE: lambda self, command, *args, **kwargs: [ + self.get_default_node() + ], + RequestPolicy.ALL_SHARDS: lambda self, command, *args, **kwargs: ( + self.get_primaries() + ), + RequestPolicy.ALL_NODES: lambda self, command, *args, **kwargs: ( + self.get_nodes() + ), + RequestPolicy.ALL_REPLICAS: lambda self, command, *args, **kwargs: ( + self.get_replicas() + ), + RequestPolicy.MULTI_SHARD: lambda self, command, *args, **kwargs: ( + self._split_multi_shard_command(*args, **kwargs) + ), + RequestPolicy.SPECIAL: lambda self, command, *args, **kwargs: ( + self.get_special_nodes() + ), ResponsePolicy.DEFAULT_KEYLESS: lambda res: res, ResponsePolicy.DEFAULT_KEYED: lambda res: res, } self._policy_resolver = policy_resolver + self._policy_cb_cache = {} if event_dispatcher is None: self._event_dispatcher = EventDispatcher() @@ -4226,15 +4261,35 @@ def _send_cluster_commands( node_objs: dict = {} nodes_written = 0 nodes_read = 0 + pipe = self._pipe + + # commonly used policies for reuse + default_keyless = CommandPolicies() + default_keyed = CommandPolicies( + request_policy=RequestPolicy.DEFAULT_KEYED, + response_policy=ResponsePolicy.DEFAULT_KEYED, + ) + policy_resolver = pipe._policy_resolver + pipe_command_flags = pipe.command_flags + command_flags = self.command_flags + no_default_node = not pipe.get_default_node() + + policy_cache = {} + sentinel = object() try: # as we move through each command that still needs to be processed, # we figure out the slot number that command maps to, then from # the slot determine the node. for c in attempt: - command_policies = self._pipe._policy_resolver.resolve( - c.args[0].lower() - ) + args = c.args + arg0 = args[0] + + command_policies = policy_cache.get(arg0, sentinel) + if command_policies is sentinel: + command_policies = policy_resolver.resolve(arg0.lower()) + policy_cache[arg0] = command_policies + # refer to our internal node -> slot table that # tells us where a given command should route to. # (it might be possible we have a cached node that no longer @@ -4244,60 +4299,60 @@ def _send_cluster_commands( target_nodes = self._parse_target_nodes(passed_targets) if not command_policies: - command_policies = CommandPolicies() + command_policies = default_keyless else: if not command_policies: - command = c.args[0].upper() - if ( - len(c.args) >= 2 - and f"{c.args[0]} {c.args[1]}".upper() - in self._pipe.command_flags - ): - command = f"{c.args[0]} {c.args[1]}".upper() + if len(args) >= 2: + command = f"{arg0} {args[1]}".upper() + if command not in pipe_command_flags: + command = arg0.upper() + else: + command = arg0.upper() # We only could resolve key properties if command is not # in a list of pre-defined request policies - command_flag = self.command_flags.get(command) + command_flag = command_flags.get(command) if not command_flag: # Fallback to default policy - if not self._pipe.get_default_node(): + if no_default_node: keys = None else: - keys = self._pipe._get_command_keys(*c.args) + keys = pipe._get_command_keys(*args) if not keys or len(keys) == 0: - command_policies = CommandPolicies() + command_policies = default_keyless else: - command_policies = CommandPolicies( - request_policy=RequestPolicy.DEFAULT_KEYED, - response_policy=ResponsePolicy.DEFAULT_KEYED, - ) + command_policies = default_keyed + if ( + command == arg0 + and pipe.commands_parser._is_keyed_command(*args) + ): + # safe to cache + policy_cache[arg0] = command_policies else: - if command_flag in self._pipe._command_flags_mapping: + if command_flag in pipe._command_flags_mapping: command_policies = CommandPolicies( - request_policy=self._pipe._command_flags_mapping[ + request_policy=pipe._command_flags_mapping[ command_flag ] ) else: - command_policies = CommandPolicies() + command_policies = default_keyless target_nodes = self._determine_nodes( - *c.args, + *args, request_policy=command_policies.request_policy, node_flag=passed_targets, ) if not target_nodes: raise RedisClusterException( - f"No targets were found to execute {c.args} command on" + f"No targets were found to execute {args} command on" ) c.command_policies = command_policies if len(target_nodes) > 1: - raise RedisClusterException( - f"Too many targets for command {c.args}" - ) + raise RedisClusterException(f"Too many targets for command {args}") node = target_nodes[0] - if node == self._pipe.get_default_node(): + if node == pipe.get_default_node(): is_default_node = True # now that we know the name of the node @@ -4305,7 +4360,7 @@ def _send_cluster_commands( # we can build a list of commands for each node. node_name = node.name if node_name not in nodes: - redis_node = self._pipe.get_redis_connection(node) + redis_node = pipe.get_redis_connection(node) try: connection = get_connection(redis_node) except (ConnectionError, TimeoutError): @@ -4316,7 +4371,7 @@ def _send_cluster_commands( # Retry object. Reinitialize the node -> slot table. self._nodes_manager.initialize() if is_default_node: - self._pipe.replace_default_node() + pipe.replace_default_node() nodes = {} raise nodes[node_name] = NodeCommands( @@ -4424,16 +4479,16 @@ def _send_cluster_commands( # If a lot of commands have failed, we'll be setting the # flag to rebuild the slots table from scratch. # So MOVED errors should correct themselves fairly quickly. - self._pipe.reinitialize_counter += 1 - if self._pipe._should_reinitialized(): + pipe.reinitialize_counter += 1 + if pipe._should_reinitialized(): self._nodes_manager.initialize() if is_default_node: - self._pipe.replace_default_node() + pipe.replace_default_node() for c in attempt: try: # send each command individually like we # do in the main client. - c.result = self._pipe.parent_execute_command(*c.args, **c.options) + c.result = pipe.parent_execute_command(*c.args, **c.options) except RedisError as e: c.result = e @@ -4441,16 +4496,12 @@ def _send_cluster_commands( # to the sequence of commands issued in the stack in pipeline.execute() response = [] for c in sorted(stack, key=lambda x: x.position): - if c.args[0] in self._pipe.cluster_response_callbacks: + if c.args[0] in pipe.cluster_response_callbacks: # Remove keys entry, it needs only for cache. c.options.pop("keys", None) - c.result = self._pipe._policies_callback_mapping[ + c.result = pipe._policies_callback_mapping[ c.command_policies.response_policy - ]( - self._pipe.cluster_response_callbacks[c.args[0]]( - c.result, **c.options - ) - ) + ](pipe.cluster_response_callbacks[c.args[0]](c.result, **c.options)) response.append(c.result) if raise_on_error: @@ -4486,36 +4537,40 @@ def _determine_nodes( ) -> List["ClusterNode"]: # Determine which nodes should be executed the command on. # Returns a list of target nodes. - command = args[0].upper() - if ( - len(args) >= 2 - and f"{args[0]} {args[1]}".upper() in self._pipe.command_flags - ): - command = f"{args[0]} {args[1]}".upper() - + pipe = self._pipe + arg0 = args[0] nodes_flag = kwargs.pop("nodes_flag", None) - if nodes_flag is not None: - # nodes flag passed by the user - command_flag = nodes_flag + if nodes_flag is None: + policy_cb = pipe._policy_cb_cache.get(arg0) else: - # get the nodes group for this command if it was predefined - command_flag = self._pipe.command_flags.get(command) - - if command_flag in self._pipe._command_flags_mapping: - request_policy = self._pipe._command_flags_mapping[command_flag] + policy_cb = None + if policy_cb is None: + command = arg0.upper() + if len(args) >= 2 and f"{arg0} {args[1]}".upper() in pipe.command_flags: + command = f"{arg0} {args[1]}".upper() + + if nodes_flag is not None: + # nodes flag passed by the user + command_flag = nodes_flag + else: + # get the nodes group for this command if it was predefined + command_flag = pipe.command_flags.get(command) - policy_callback = self._pipe._policies_callback_mapping[request_policy] + request_policy = pipe._command_flags_mapping.get( + command_flag, request_policy + ) + policy_cb = pipe._policies_callback_mapping[request_policy] - if request_policy == RequestPolicy.DEFAULT_KEYED: - nodes = policy_callback(command, *args) - elif request_policy == RequestPolicy.MULTI_SHARD: - nodes = policy_callback(*args, **kwargs) - elif request_policy == RequestPolicy.DEFAULT_KEYLESS: - nodes = policy_callback(args[0]) + if nodes_flag is None and command == arg0: + if len(pipe._policy_cb_cache) > 5000: + # Prevent unbounded memory leak on abnormal use + pipe._policy_cb_cache.clear() + pipe._policy_cb_cache[arg0] = policy_cb else: - nodes = policy_callback() + command = arg0 + nodes = policy_cb(pipe, command, *args, **kwargs) - if args[0].lower() == "ft.aggregate": + if arg0.lower() == "ft.aggregate": self._aggregate_nodes = nodes return nodes diff --git a/redis/commands/policies.py b/redis/commands/policies.py index af704e91e4..de97848064 100644 --- a/redis/commands/policies.py +++ b/redis/commands/policies.py @@ -193,26 +193,29 @@ def __init__( self._fallback = fallback def resolve(self, command_name: str) -> Optional[CommandPolicies]: - parts = command_name.split(".") + parts = command_name.split(".", 2) if len(parts) > 2: raise ValueError(f"Wrong command or module name: {command_name}") module, command = parts if len(parts) == 2 else ("core", parts[0]) - if self._policies.get(module, None) is None: + module_policies = self._policies.get(module, None) + + if module_policies is None: if self._fallback is not None: return self._fallback.resolve(command_name) else: return None - if self._policies.get(module).get(command, None) is None: + command_policy = module_policies.get(command, None) + if command_policy is None: if self._fallback is not None: return self._fallback.resolve(command_name) else: return None - return self._policies.get(module).get(command) + return command_policy @abstractmethod def with_fallback(self, fallback: "PolicyResolver") -> "PolicyResolver": @@ -231,26 +234,28 @@ def __init__( self._fallback = fallback async def resolve(self, command_name: str) -> Optional[CommandPolicies]: - parts = command_name.split(".") + parts = command_name.split(".", 2) if len(parts) > 2: raise ValueError(f"Wrong command or module name: {command_name}") module, command = parts if len(parts) == 2 else ("core", parts[0]) - if self._policies.get(module, None) is None: + module_policies = self._policies.get(module, None) + if module_policies is None: if self._fallback is not None: return await self._fallback.resolve(command_name) else: return None - if self._policies.get(module).get(command, None) is None: + command_policy = module_policies.get(command, None) + if command_policy is None: if self._fallback is not None: return await self._fallback.resolve(command_name) else: return None - return self._policies.get(module).get(command) + return command_policy @abstractmethod def with_fallback(self, fallback: "AsyncPolicyResolver") -> "AsyncPolicyResolver":