Add minimalistic implementation for Push notifications parsing with hiredis-py

Author: uglide

redis/_parsers/__init__.py

@@ -1,4 +1,4 @@
-from .base import BaseParser, _AsyncRESPBase
+from .base import AsyncPushNotificationsParser, BaseParser, PushNotificationsParser, _AsyncRESPBase
 from .commands import AsyncCommandsParser, CommandsParser
 from .encoders import Encoder
 from .hiredis import _AsyncHiredisParser, _HiredisParser
@@ -11,10 +11,12 @@
     "_AsyncRESPBase",
     "_AsyncRESP2Parser",
     "_AsyncRESP3Parser",
+    "AsyncPushNotificationsParser",
     "CommandsParser",
     "Encoder",
     "BaseParser",
     "_HiredisParser",
     "_RESP2Parser",
     "_RESP3Parser",
+    "PushNotificationsParser",
 ]

redis/_parsers/base.py

@@ -1,7 +1,7 @@
 import sys
 from abc import ABC
 from asyncio import IncompleteReadError, StreamReader, TimeoutError
-from typing import List, Optional, Union
+from typing import Callable, List, Optional, Protocol, Union
 
 if sys.version_info.major >= 3 and sys.version_info.minor >= 11:
     from asyncio import timeout as async_timeout
@@ -158,6 +158,58 @@ async def read_response(
         raise NotImplementedError()
 
 
+_INVALIDATION_MESSAGE = [b"invalidate", "invalidate"]
+
+
+class PushNotificationsParser(Protocol):
+    """Protocol defining RESP3-specific parsing functionality"""
+
+    pubsub_push_handler_func: Callable
+    invalidation_push_handler_func: Optional[Callable] = None
+
+    def handle_pubsub_push_response(self, response):
+        """Handle pubsub push responses"""
+        raise NotImplementedError()
+
+    def handle_push_response(self, response, **kwargs):
+        if response[0] not in _INVALIDATION_MESSAGE:
+            return self.pubsub_push_handler_func(response)
+        if self.invalidation_push_handler_func:
+            return self.invalidation_push_handler_func(response)
+
+    def set_pubsub_push_handler(self, pubsub_push_handler_func):
+        self.pubsub_push_handler_func = pubsub_push_handler_func
+
+    def set_invalidation_push_handler(self, invalidation_push_handler_func):
+        self.invalidation_push_handler_func = invalidation_push_handler_func
+
+
+class AsyncPushNotificationsParser(Protocol):
+    """Protocol defining async RESP3-specific parsing functionality"""
+
+    pubsub_push_handler_func: Callable
+    invalidation_push_handler_func: Optional[Callable] = None
+
+    async def handle_pubsub_push_response(self, response):
+        """Handle pubsub push responses asynchronously"""
+        ...
+
+    async def handle_push_response(self, response, **kwargs):
+        """Handle push responses asynchronously"""
+        if response[0] not in _INVALIDATION_MESSAGE:
+            return await self.pubsub_push_handler_func(response)
+        if self.invalidation_push_handler_func:
+            return await self.invalidation_push_handler_func(response)
+
+    def set_pubsub_push_handler(self, pubsub_push_handler_func):
+        """Set the pubsub push handler function"""
+        self.pubsub_push_handler_func = pubsub_push_handler_func
+
+    def set_invalidation_push_handler(self, invalidation_push_handler_func):
+        """Set the invalidation push handler function"""
+        self.invalidation_push_handler_func = invalidation_push_handler_func
+
+
 class _AsyncRESPBase(AsyncBaseParser):
     """Base class for async resp parsing"""
 

redis/_parsers/hiredis.py

@@ -1,6 +1,7 @@
 import asyncio
 import socket
 import sys
+from logging import getLogger
 from typing import Callable, List, Optional, TypedDict, Union
 
 if sys.version_info.major >= 3 and sys.version_info.minor >= 11:
@@ -11,7 +12,12 @@
 from ..exceptions import ConnectionError, InvalidResponse, RedisError
 from ..typing import EncodableT
 from ..utils import HIREDIS_AVAILABLE
-from .base import AsyncBaseParser, BaseParser
+from .base import (
+    AsyncBaseParser,
+    BaseParser,
+    PushNotificationsParser,
+    AsyncPushNotificationsParser,
+)
 from .socket import (
     NONBLOCKING_EXCEPTION_ERROR_NUMBERS,
     NONBLOCKING_EXCEPTIONS,
@@ -32,21 +38,29 @@ class _HiredisReaderArgs(TypedDict, total=False):
     errors: Optional[str]
 
 
-class _HiredisParser(BaseParser):
+class _HiredisParser(BaseParser, PushNotificationsParser):
     "Parser class for connections using Hiredis"
 
     def __init__(self, socket_read_size):
         if not HIREDIS_AVAILABLE:
             raise RedisError("Hiredis is not installed")
         self.socket_read_size = socket_read_size
         self._buffer = bytearray(socket_read_size)
+        self.pubsub_push_handler_func = self.handle_pubsub_push_response
+        self.invalidation_push_handler_func = None
+        self._hiredis_PushNotificationType = None
 
     def __del__(self):
         try:
             self.on_disconnect()
         except Exception:
             pass
 
+    def handle_pubsub_push_response(self, response):
+        logger = getLogger("push_response")
+        logger.debug("Push response: " + str(response))
+        return response
+
     def on_connect(self, connection, **kwargs):
         import hiredis
 
@@ -64,6 +78,12 @@ def on_connect(self, connection, **kwargs):
         self._reader = hiredis.Reader(**kwargs)
         self._next_response = NOT_ENOUGH_DATA
 
+        try:
+            self._hiredis_PushNotificationType = hiredis.PushNotification
+        except AttributeError:
+            # hiredis < 3.2
+            self._hiredis_PushNotificationType = None
+
     def on_disconnect(self):
         self._sock = None
         self._reader = None
@@ -109,7 +129,7 @@ def read_from_socket(self, timeout=SENTINEL, raise_on_timeout=True):
             if custom_timeout:
                 sock.settimeout(self._socket_timeout)
 
-    def read_response(self, disable_decoding=False):
+    def read_response(self, disable_decoding=False, push_request=False):
         if not self._reader:
             raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
 
@@ -135,6 +155,16 @@ def read_response(self, disable_decoding=False):
         # happened
         if isinstance(response, ConnectionError):
             raise response
+        elif self._hiredis_PushNotificationType is not None and isinstance(
+            response, self._hiredis_PushNotificationType
+        ):
+            response = self.handle_push_response(response)
+            if not push_request:
+                return self.read_response(
+                    disable_decoding=disable_decoding, push_request=push_request
+                )
+            else:
+                return response
         elif (
             isinstance(response, list)
             and response
@@ -154,6 +184,8 @@ def __init__(self, socket_read_size: int):
             raise RedisError("Hiredis is not available.")
         super().__init__(socket_read_size=socket_read_size)
         self._reader = None
+        self.pubsub_push_handler_func = self.handle_pubsub_push_response
+        self.invalidation_push_handler_func = None
 
     def on_connect(self, connection):
         import hiredis
@@ -171,6 +203,14 @@ def on_connect(self, connection):
         self._reader = hiredis.Reader(**kwargs)
         self._connected = True
 
+        try:
+            self._hiredis_PushNotificationType = getattr(
+                hiredis, "PushNotificationType", None
+            )
+        except AttributeError:
+            # hiredis < 3.2
+            self._hiredis_PushNotificationType = None
+
     def on_disconnect(self):
         self._connected = False
 
@@ -195,7 +235,7 @@ async def read_from_socket(self):
         return True
 
     async def read_response(
-        self, disable_decoding: bool = False
+        self, disable_decoding: bool = False, push_request: bool = False
     ) -> Union[EncodableT, List[EncodableT]]:
         # If `on_disconnect()` has been called, prohibit any more reads
         # even if they could happen because data might be present.
@@ -219,6 +259,16 @@ async def read_response(
         # happened
         if isinstance(response, ConnectionError):
             raise response
+        elif self._hiredis_PushNotificationType is not None and isinstance(
+            response, self._hiredis_PushNotificationType
+        ):
+            response = await self.handle_push_response(response)
+            if not push_request:
+                return await self.read_response(
+                    disable_decoding=disable_decoding, push_request=push_request
+                )
+            else:
+                return response
         elif (
             isinstance(response, list)
             and response

redis/_parsers/resp3.py

@@ -3,13 +3,16 @@
 
 from ..exceptions import ConnectionError, InvalidResponse, ResponseError
 from ..typing import EncodableT
-from .base import _AsyncRESPBase, _RESPBase
+from .base import (
+    AsyncPushNotificationsParser,
+    PushNotificationsParser,
+    _AsyncRESPBase,
+    _RESPBase,
+)
 from .socket import SERVER_CLOSED_CONNECTION_ERROR
 
-_INVALIDATION_MESSAGE = [b"invalidate", "invalidate"]
 
-
-class _RESP3Parser(_RESPBase):
+class _RESP3Parser(_RESPBase, PushNotificationsParser):
     """RESP3 protocol implementation"""
 
     def __init__(self, socket_read_size):
@@ -113,9 +116,7 @@ def _read_response(self, disable_decoding=False, push_request=False):
                 )
                 for _ in range(int(response))
             ]
-            response = self.handle_push_response(
-                response, disable_decoding, push_request
-            )
+            response = self.handle_push_response(response)
             if not push_request:
                 return self._read_response(
                     disable_decoding=disable_decoding, push_request=push_request
@@ -129,20 +130,8 @@ def _read_response(self, disable_decoding=False, push_request=False):
             response = self.encoder.decode(response)
         return response
 
-    def handle_push_response(self, response, disable_decoding, push_request):
-        if response[0] not in _INVALIDATION_MESSAGE:
-            return self.pubsub_push_handler_func(response)
-        if self.invalidation_push_handler_func:
-            return self.invalidation_push_handler_func(response)
-
-    def set_pubsub_push_handler(self, pubsub_push_handler_func):
-        self.pubsub_push_handler_func = pubsub_push_handler_func
-
-    def set_invalidation_push_handler(self, invalidation_push_handler_func):
-        self.invalidation_push_handler_func = invalidation_push_handler_func
-
 
-class _AsyncRESP3Parser(_AsyncRESPBase):
+class _AsyncRESP3Parser(_AsyncRESPBase, AsyncPushNotificationsParser):
     def __init__(self, socket_read_size):
         super().__init__(socket_read_size)
         self.pubsub_push_handler_func = self.handle_pubsub_push_response
@@ -253,9 +242,7 @@ async def _read_response(
                 )
                 for _ in range(int(response))
             ]
-            response = await self.handle_push_response(
-                response, disable_decoding, push_request
-            )
+            response = await self.handle_push_response(response)
             if not push_request:
                 return await self._read_response(
                     disable_decoding=disable_decoding, push_request=push_request
@@ -268,15 +255,3 @@ async def _read_response(
         if isinstance(response, bytes) and disable_decoding is False:
             response = self.encoder.decode(response)
         return response
-
-    async def handle_push_response(self, response, disable_decoding, push_request):
-        if response[0] not in _INVALIDATION_MESSAGE:
-            return await self.pubsub_push_handler_func(response)
-        if self.invalidation_push_handler_func:
-            return await self.invalidation_push_handler_func(response)
-
-    def set_pubsub_push_handler(self, pubsub_push_handler_func):
-        self.pubsub_push_handler_func = pubsub_push_handler_func
-
-    def set_invalidation_push_handler(self, invalidation_push_handler_func):
-        self.invalidation_push_handler_func = invalidation_push_handler_func

redis/asyncio/client.py

@@ -70,7 +70,6 @@
 )
 from redis.typing import ChannelT, EncodableT, KeyT
 from redis.utils import (
-    HIREDIS_AVAILABLE,
     SSL_AVAILABLE,
     _set_info_logger,
     deprecated_args,
@@ -938,7 +937,7 @@ async def connect(self):
             self.connection.register_connect_callback(self.on_connect)
         else:
             await self.connection.connect()
-        if self.push_handler_func is not None and not HIREDIS_AVAILABLE:
+        if self.push_handler_func is not None:
             self.connection._parser.set_pubsub_push_handler(self.push_handler_func)
 
         self._event_dispatcher.dispatch(

redis/asyncio/connection.py

@@ -590,7 +590,7 @@ async def read_response(
                     response = await self._parser.read_response(
                         disable_decoding=disable_decoding
                     )
-            elif self.protocol in ["3", 3] and not HIREDIS_AVAILABLE:
+            elif self.protocol in ["3", 3]:
                 response = await self._parser.read_response(
                     disable_decoding=disable_decoding, push_request=push_request
                 )

redis/client.py

@@ -58,7 +58,6 @@
 from redis.lock import Lock
 from redis.retry import Retry
 from redis.utils import (
-    HIREDIS_AVAILABLE,
     _set_info_logger,
     deprecated_args,
     get_lib_version,
@@ -861,7 +860,7 @@ def execute_command(self, *args):
             # register a callback that re-subscribes to any channels we
             # were listening to when we were disconnected
             self.connection.register_connect_callback(self.on_connect)
-            if self.push_handler_func is not None and not HIREDIS_AVAILABLE:
+            if self.push_handler_func is not None:
                 self.connection._parser.set_pubsub_push_handler(self.push_handler_func)
             self._event_dispatcher.dispatch(
                 AfterPubSubConnectionInstantiationEvent(

redis/cluster.py

@@ -51,7 +51,6 @@
 from redis.lock import Lock
 from redis.retry import Retry
 from redis.utils import (
-    HIREDIS_AVAILABLE,
     deprecated_args,
     dict_merge,
     list_keys_to_dict,
@@ -1999,7 +1998,7 @@ def execute_command(self, *args):
             # register a callback that re-subscribes to any channels we
             # were listening to when we were disconnected
             self.connection.register_connect_callback(self.on_connect)
-            if self.push_handler_func is not None and not HIREDIS_AVAILABLE:
+            if self.push_handler_func is not None:
                 self.connection._parser.set_pubsub_push_handler(self.push_handler_func)
             self._event_dispatcher.dispatch(
                 AfterPubSubConnectionInstantiationEvent(

redis/connection.py

@@ -636,7 +636,7 @@ def read_response(
         host_error = self._host_error()
 
         try:
-            if self.protocol in ["3", 3] and not HIREDIS_AVAILABLE:
+            if self.protocol in ["3", 3]:
                 response = self._parser.read_response(
                     disable_decoding=disable_decoding, push_request=push_request
                 )

redis/utils.py

@@ -12,9 +12,10 @@
     import hiredis  # noqa
 
     # Only support Hiredis >= 3.0:
-    HIREDIS_AVAILABLE = int(hiredis.__version__.split(".")[0]) >= 3
+    hiredis_version = hiredis.__version__.split(".")
+    HIREDIS_AVAILABLE = int(hiredis_version[0]) >= 3 and int(hiredis_version[1]) >= 2
     if not HIREDIS_AVAILABLE:
-        raise ImportError("hiredis package should be >= 3.0.0")
+        raise ImportError("hiredis package should be >= 3.2.0")
 except ImportError:
     HIREDIS_AVAILABLE = False