Added support for Pub/Sub

vladvildanov · vladvildanov · commit 0de0f4d5f530 · 2024-12-17T16:48:13.000+02:00
diff --git a/redis/asyncio/client.py b/redis/asyncio/client.py
@@ -54,7 +54,7 @@
 )
 from redis.credentials import CredentialProvider
 from redis.event import EventDispatcher, AfterPooledConnectionsInstantiationEvent, ClientType, \
-    AfterSingleConnectionInstantiationEvent
+    AfterSingleConnectionInstantiationEvent, AfterPubSubConnectionInstantiationEvent
 from redis.exceptions import (
     ConnectionError,
     ExecAbortError,
@@ -539,7 +539,7 @@ def pubsub(self, **kwargs) -> "PubSub":
         subscribe to channels and listen for messages that get published to
         them.
         """
-        return PubSub(self.connection_pool, **kwargs)
+        return PubSub(self.connection_pool, event_dispatcher=self._event_dispatcher, **kwargs)
 
     def monitor(self) -> "Monitor":
         return Monitor(self.connection_pool)
@@ -777,6 +777,7 @@ def __init__(
         ignore_subscribe_messages: bool = False,
         encoder=None,
         push_handler_func: Optional[Callable] = None,
+        event_dispatcher: Optional["EventDispatcher"] = EventDispatcher(),
     ):
         self.connection_pool = connection_pool
         self.shard_hint = shard_hint
@@ -804,6 +805,7 @@ def __init__(
         self.pending_unsubscribe_channels = set()
         self.patterns = {}
         self.pending_unsubscribe_patterns = set()
+        self._event_dispatcher = event_dispatcher
         self._lock = asyncio.Lock()
 
     async def __aenter__(self):
@@ -894,6 +896,15 @@ async def connect(self):
         if self.push_handler_func is not None and not HIREDIS_AVAILABLE:
             self.connection._parser.set_pubsub_push_handler(self.push_handler_func)
 
+        self._event_dispatcher.dispatch(
+            AfterPubSubConnectionInstantiationEvent(
+                self.connection,
+                self.connection_pool,
+                ClientType.ASYNC,
+                self._lock
+            )
+        )
+
     async def _disconnect_raise_connect(self, conn, error):
         """
         Close the connection and raise an exception
diff --git a/redis/asyncio/connection.py b/redis/asyncio/connection.py
@@ -334,6 +334,9 @@ def _host_error(self) -> str:
     def _error_message(self, exception: BaseException) -> str:
         return format_error_message(self._host_error(), exception)
 
+    def get_protocol(self):
+        return self.protocol
+
     async def on_connect(self) -> None:
         """Initialize the connection, authenticate and select a database"""
         self._parser.on_connect(self)
diff --git a/redis/client.py b/redis/client.py
@@ -29,7 +29,7 @@
 )
 from redis.credentials import CredentialProvider
 from redis.event import EventDispatcher, AfterPooledConnectionsInstantiationEvent, ClientType, \
-    AfterSingleConnectionInstantiationEvent
+    AfterSingleConnectionInstantiationEvent, AfterPubSubConnectionInstantiationEvent
 from redis.exceptions import (
     ConnectionError,
     ExecAbortError,
@@ -332,6 +332,7 @@ def __init__(
             ))
 
         self.connection_pool = connection_pool
+        self._event_dispatcher = event_dispatcher
 
         if (cache_config or cache) and self.connection_pool.get_protocol() not in [
             3,
@@ -518,7 +519,7 @@ def pubsub(self, **kwargs):
         subscribe to channels and listen for messages that get published to
         them.
         """
-        return PubSub(self.connection_pool, **kwargs)
+        return PubSub(self.connection_pool, event_dispatcher=self._event_dispatcher, **kwargs)
 
     def monitor(self):
         return Monitor(self.connection_pool)
@@ -709,6 +710,7 @@ def __init__(
         ignore_subscribe_messages: bool = False,
         encoder: Optional["Encoder"] = None,
         push_handler_func: Union[None, Callable[[str], None]] = None,
+        event_dispatcher: Optional["EventDispatcher"] = EventDispatcher(),
     ):
         self.connection_pool = connection_pool
         self.shard_hint = shard_hint
@@ -719,6 +721,8 @@ def __init__(
         # to lookup channel and pattern names for callback handlers.
         self.encoder = encoder
         self.push_handler_func = push_handler_func
+        self._event_dispatcher = event_dispatcher
+        self._lock = threading.Lock()
         if self.encoder is None:
             self.encoder = self.connection_pool.get_encoder()
         self.health_check_response_b = self.encoder.encode(self.HEALTH_CHECK_MESSAGE)
@@ -809,11 +813,20 @@ def execute_command(self, *args):
             self.connection.register_connect_callback(self.on_connect)
             if self.push_handler_func is not None and not HIREDIS_AVAILABLE:
                 self.connection._parser.set_pubsub_push_handler(self.push_handler_func)
+            self._event_dispatcher.dispatch(
+                AfterPubSubConnectionInstantiationEvent(
+                    self.connection,
+                    self.connection_pool,
+                    ClientType.SYNC,
+                    self._lock
+                )
+            )
         connection = self.connection
         kwargs = {"check_health": not self.subscribed}
         if not self.subscribed:
             self.clean_health_check_responses()
-        self._execute(connection, connection.send_command, *args, **kwargs)
+        with self._lock:
+            self._execute(connection, connection.send_command, *args, **kwargs)
 
     def clean_health_check_responses(self) -> None:
         """
diff --git a/redis/connection.py b/redis/connection.py
@@ -153,6 +153,10 @@ def deregister_connect_callback(self, callback):
     def set_parser(self, parser_class):
         pass
 
+    @abstractmethod
+    def get_protocol(self):
+        pass
+
     @abstractmethod
     def connect(self):
         pass
diff --git a/redis/event.py b/redis/event.py
@@ -58,6 +58,9 @@ def __init__(self):
             AfterSingleConnectionInstantiationEvent: [
                 RegisterReAuthForSingleConnection()
             ],
+            AfterPubSubConnectionInstantiationEvent: [
+                RegisterReAuthForPubSub()
+            ],
             AfterAsyncClusterInstantiationEvent: [
                 RegisterReAuthForAsyncClusterNodes()
             ],
@@ -158,6 +161,36 @@ def connection_lock(self) -> Union[threading.Lock, asyncio.Lock]:
         return self._connection_lock
 
 
+class AfterPubSubConnectionInstantiationEvent:
+    def __init__(
+            self,
+            pubsub_connection,
+            connection_pool,
+            client_type: ClientType,
+            connection_lock: Union[threading.Lock, asyncio.Lock]
+    ):
+        self._pubsub_connection = pubsub_connection
+        self._connection_pool = connection_pool
+        self._client_type = client_type
+        self._connection_lock = connection_lock
+
+    @property
+    def pubsub_connection(self):
+        return self._pubsub_connection
+
+    @property
+    def connection_pool(self):
+        return self._connection_pool
+
+    @property
+    def client_type(self) -> ClientType:
+        return self._client_type
+
+    @property
+    def connection_lock(self) -> Union[threading.Lock, asyncio.Lock]:
+        return self._connection_lock
+
+
 class AfterAsyncClusterInstantiationEvent:
     """
     Event that will be fired after async cluster instance was created.
@@ -266,3 +299,40 @@ def listen(self, event: AfterAsyncClusterInstantiationEvent):
     async def _re_auth(self, token: TokenInterface):
         for key in self._event.nodes:
             await self._event.nodes[key].re_auth_callback(token)
+
+
+class RegisterReAuthForPubSub(EventListenerInterface):
+    def __init__(self):
+        self._connection = None
+        self._connection_pool = None
+        self._client_type = None
+        self._connection_lock = None
+
+    def listen(self, event: AfterPubSubConnectionInstantiationEvent):
+        if (
+                isinstance(event.pubsub_connection.credential_provider, StreamingCredentialProvider)
+                and event.pubsub_connection.get_protocol() in [3, "3"]
+        ):
+            self._connection = event.pubsub_connection
+            self._connection_pool = event.connection_pool
+            self._client_type = event.client_type
+            self._connection_lock = event.connection_lock
+
+            if self._client_type == ClientType.SYNC:
+                self._connection.credential_provider.on_next(self._re_auth)
+            else:
+                self._connection.credential_provider.on_next(self._re_auth_async)
+
+    def _re_auth(self, token: TokenInterface):
+        with self._connection_lock:
+            self._connection.send_command('AUTH', token.try_get('oid'), token.get_value())
+            self._connection.read_response()
+
+        self._connection_pool.re_auth_callback(token)
+
+    async def _re_auth_async(self, token: TokenInterface):
+        async with self._connection_lock:
+            await self._connection.send_command('AUTH', token.try_get('oid'), token.get_value())
+            await self._connection.read_response()
+
+        await self._connection_pool.re_auth_callback(token)
diff --git a/tests/test_asyncio/test_credentials.py b/tests/test_asyncio/test_credentials.py
@@ -414,6 +414,115 @@ async def re_auth_callback(token):
         mock_another_connection.read_response.assert_has_calls([call()])
         mock_failed_connection.read_response.assert_has_calls([call(), call(), call()])
 
+    @pytest.mark.parametrize(
+        "credential_provider",
+        [
+            {
+                "cred_provider_class": EntraIdCredentialsProvider,
+                "cred_provider_kwargs": {"expiration_refresh_ratio": 0.00005},
+                "mock_idp": True,
+            }
+        ],
+        indirect=True,
+    )
+    async def test_re_auth_pub_sub_in_resp3(self, credential_provider):
+        mock_pubsub_connection = Mock(spec=Connection)
+        mock_pubsub_connection.get_protocol.return_value = 3
+        mock_pubsub_connection.credential_provider = credential_provider
+        mock_pubsub_connection.retry = Retry(NoBackoff(), 3)
+        mock_another_connection = Mock(spec=Connection)
+        mock_another_connection.retry = Retry(NoBackoff(), 3)
+
+        mock_pool = Mock(spec=ConnectionPool)
+        mock_pool.connection_kwargs = {
+            "credential_provider": credential_provider,
+        }
+        mock_pool.get_connection.side_effect = [mock_pubsub_connection, mock_another_connection]
+        mock_pool._available_connections = [mock_another_connection]
+        mock_pool._lock = AsyncLock()
+        auth_token = None
+
+        async def re_auth_callback(token):
+            nonlocal auth_token
+            auth_token = token
+            async with mock_pool._lock:
+                for conn in mock_pool._available_connections:
+                    await conn.send_command('AUTH', token.try_get('oid'), token.get_value())
+                    await conn.read_response()
+
+        mock_pool.re_auth_callback = re_auth_callback
+
+        r = Redis(
+            connection_pool=mock_pool,
+            credential_provider=credential_provider,
+        )
+        p = r.pubsub()
+        await p.subscribe('test')
+        await credential_provider.get_credentials_async()
+        await async_sleep(0.5)
+
+        mock_pubsub_connection.send_command.assert_has_calls([
+            call('SUBSCRIBE', 'test', check_health=True),
+            call('AUTH', auth_token.try_get('oid'), auth_token.get_value()),
+        ])
+        mock_another_connection.send_command.assert_has_calls([
+            call('AUTH', auth_token.try_get('oid'), auth_token.get_value())
+        ])
+
+    @pytest.mark.parametrize(
+        "credential_provider",
+        [
+            {
+                "cred_provider_class": EntraIdCredentialsProvider,
+                "cred_provider_kwargs": {"expiration_refresh_ratio": 0.00005},
+                "mock_idp": True,
+            }
+        ],
+        indirect=True,
+    )
+    async def test_do_not_re_auth_pub_sub_in_resp2(self, credential_provider):
+        mock_pubsub_connection = Mock(spec=Connection)
+        mock_pubsub_connection.get_protocol.return_value = 2
+        mock_pubsub_connection.credential_provider = credential_provider
+        mock_pubsub_connection.retry = Retry(NoBackoff(), 3)
+        mock_another_connection = Mock(spec=Connection)
+        mock_another_connection.retry = Retry(NoBackoff(), 3)
+
+        mock_pool = Mock(spec=ConnectionPool)
+        mock_pool.connection_kwargs = {
+            "credential_provider": credential_provider,
+        }
+        mock_pool.get_connection.side_effect = [mock_pubsub_connection, mock_another_connection]
+        mock_pool._available_connections = [mock_another_connection]
+        mock_pool._lock = AsyncLock()
+        auth_token = None
+
+        async def re_auth_callback(token):
+            nonlocal auth_token
+            auth_token = token
+            async with mock_pool._lock:
+                for conn in mock_pool._available_connections:
+                    await conn.send_command('AUTH', token.try_get('oid'), token.get_value())
+                    await conn.read_response()
+
+        mock_pool.re_auth_callback = re_auth_callback
+
+        r = Redis(
+            connection_pool=mock_pool,
+            credential_provider=credential_provider,
+        )
+        p = r.pubsub()
+        await p.subscribe('test')
+        await credential_provider.get_credentials_async()
+        await async_sleep(0.5)
+
+        mock_pubsub_connection.send_command.assert_has_calls([
+            call('SUBSCRIBE', 'test', check_health=True),
+        ])
+        mock_another_connection.send_command.assert_has_calls([
+            call('AUTH', auth_token.try_get('oid'), auth_token.get_value())
+        ])
+
 
 @pytest.mark.asyncio
 @pytest.mark.onlynoncluster
diff --git a/tests/test_credentials.py b/tests/test_credentials.py
