diff --git a/redis/asyncio/connection.py b/redis/asyncio/connection.py index 9cd0d082a4..939a95f9b7 100644 --- a/redis/asyncio/connection.py +++ b/redis/asyncio/connection.py @@ -74,6 +74,7 @@ ) from redis.asyncio.retry import Retry from redis.backoff import NoBackoff +from redis.connection import HiredisRespSerializer from redis.credentials import CredentialProvider, UsernamePasswordCredentialProvider from redis.exceptions import ( AuthenticationError, @@ -689,6 +690,11 @@ def __init__( self._connect_callbacks: List[weakref.WeakMethod[ConnectCallbackT]] = [] self._buffer_cutoff = 6000 self._re_auth_token: Optional[TokenInterface] = None + self._command_packer = ( + HiredisRespSerializer(self._pack_command_python, self.encoder.encode) + if HIREDIS_AVAILABLE + else None + ) self._should_reconnect = False try: @@ -1316,6 +1322,11 @@ async def _read_response_from_parser( return await self._parser.read_response(disable_decoding=disable_decoding) def pack_command(self, *args: EncodableT) -> List[bytes]: + if self._command_packer is not None: + return self._command_packer.pack(*args) + return self._pack_command_python(*args) + + def _pack_command_python(self, *args: EncodableT) -> List[bytes]: """Pack a series of arguments into the Redis protocol""" output = [] # the client might have included 1 or more literal arguments in diff --git a/redis/connection.py b/redis/connection.py index 189d16d91f..23df3cc3ba 100644 --- a/redis/connection.py +++ b/redis/connection.py @@ -122,8 +122,14 @@ class HiredisRespSerializer: + def __init__(self, fallback, encode): + self._fallback = fallback + self._encode = encode + def pack(self, *args: List): """Pack a series of arguments into the Redis protocol""" + if any(isinstance(arg, memoryview) for arg in args): + return self._fallback(*args) output = [] if isinstance(args[0], str): @@ -131,7 +137,7 @@ def pack(self, *args: List): elif b" " in args[0]: args = tuple(args[0].split()) + args[1:] args = tuple( - bytes(arg) if isinstance(arg, (bytearray, memoryview)) else arg + bytes(arg) if isinstance(arg, bytearray) else self._encode(arg) for arg in args ) try: @@ -976,7 +982,8 @@ def _construct_command_packer(self, packer): if packer is not None: return packer elif HIREDIS_AVAILABLE: - return HiredisRespSerializer() + fallback = PythonRespSerializer(self._buffer_cutoff, self.encoder.encode) + return HiredisRespSerializer(fallback.pack, self.encoder.encode) else: return PythonRespSerializer(self._buffer_cutoff, self.encoder.encode) diff --git a/tests/test_asyncio/test_connection.py b/tests/test_asyncio/test_connection.py index c56b5d9dd4..a960c4adbb 100644 --- a/tests/test_asyncio/test_connection.py +++ b/tests/test_asyncio/test_connection.py @@ -18,6 +18,7 @@ from redis.asyncio import ConnectionPool, Redis from redis.asyncio.connection import ( Connection, + HiredisRespSerializer, SSLConnection, UnixDomainSocketConnection, parse_url, @@ -83,6 +84,39 @@ def test_connection_default_parser_matches_default_protocol(): assert conn.protocol == 3 +@pytest.mark.skipif(not HIREDIS_AVAILABLE, reason="hiredis is not installed") +def test_connection_uses_hiredis_command_packer(monkeypatch): + calls = [] + + def pack_command(args): + calls.append(args) + return b"packed" + + monkeypatch.setattr("redis.connection.hiredis.pack_command", pack_command) + connection = Connection() + + assert isinstance(connection._command_packer, HiredisRespSerializer) + assert connection.pack_command("SET", "key", "value") == [b"packed"] + assert calls == [(b"SET", b"key", b"value")] + + +def test_connection_uses_python_command_packer_without_hiredis(monkeypatch): + monkeypatch.setattr("redis.asyncio.connection.HIREDIS_AVAILABLE", False) + connection = Connection() + + assert connection._command_packer is None + assert connection.pack_command("PING") == [b"*1\r\n$4\r\nPING\r\n"] + + +@pytest.mark.parametrize("protocol", [2, 3]) +def test_async_pack_command_respects_connection_encoding(protocol): + connection = Connection(encoding="latin-1", protocol=protocol) + + assert connection.pack_command("SET", "key", "café") == [ + b"*3\r\n$3\r\nSET\r\n$3\r\nkey\r\n$4\r\ncaf\xe9\r\n" + ] + + @pytest.mark.parametrize( ("buffer", "eof", "expected"), [ diff --git a/tests/test_connection.py b/tests/test_connection.py index 7963b19057..86311a878e 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -733,6 +733,15 @@ def test_pack_command(Class): assert actual == expected, f"actual = {actual}, expected = {expected}" +@pytest.mark.parametrize("protocol", [2, 3]) +def test_pack_command_respects_connection_encoding(protocol): + connection = Connection(encoding="latin-1", protocol=protocol) + + assert connection.pack_command("SET", "key", "café") == [ + b"*3\r\n$3\r\nSET\r\n$3\r\nkey\r\n$4\r\ncaf\xe9\r\n" + ] + + @pytest.mark.fixed_client def test_create_single_connection_client_from_url(): client = redis.Redis.from_url(