Added support for async API

Author: vladvildanov

dev_requirements.txt

@@ -16,3 +16,4 @@ uvloop
 vulture>=2.3.0
 wheel>=0.30.0
 numpy>=1.24.0
+redispy-entraid-credentials @ git+https://github.com/redis-developer/redispy-entra-credentials.git/@main

redis/asyncio/connection.py

@@ -5,6 +5,7 @@
 import socket
 import ssl
 import sys
+import threading
 import warnings
 import weakref
 from abc import abstractmethod
@@ -27,6 +28,7 @@
 )
 from urllib.parse import ParseResult, parse_qs, unquote, urlparse
 
+from ..event import EventDispatcher, AsyncBeforeCommandExecutionEvent
 from ..utils import format_error_message
 
 # the functionality is available in 3.11.x but has a major issue before
@@ -39,7 +41,7 @@
 from redis.asyncio.retry import Retry
 from redis.backoff import NoBackoff
 from redis.connection import DEFAULT_RESP_VERSION
-from redis.credentials import CredentialProvider, UsernamePasswordCredentialProvider
+from redis.credentials import CredentialProvider, UsernamePasswordCredentialProvider, StreamingCredentialProvider
 from redis.exceptions import (
     AuthenticationError,
     AuthenticationWrongNumberOfArgsError,
@@ -148,6 +150,7 @@ def __init__(
         encoder_class: Type[Encoder] = Encoder,
         credential_provider: Optional[CredentialProvider] = None,
         protocol: Optional[int] = 2,
+        event_dispatcher: Optional[EventDispatcher] = EventDispatcher()
     ):
         if (username or password) and credential_provider is not None:
             raise DataError(
@@ -195,6 +198,9 @@ def __init__(
         self.set_parser(parser_class)
         self._connect_callbacks: List[weakref.WeakMethod[ConnectCallbackT]] = []
         self._buffer_cutoff = 6000
+        self._event_dispatcher = event_dispatcher
+        self._init_auth_args = None
+
         try:
             p = int(protocol)
         except TypeError:
@@ -333,7 +339,9 @@ async def on_connect(self) -> None:
                 self.credential_provider
                 or UsernamePasswordCredentialProvider(self.username, self.password)
             )
-            auth_args = cred_provider.get_credentials()
+            auth_args = await cred_provider.get_credentials_async()
+            self._init_auth_args = hash(auth_args)
+
             # if resp version is specified and we have auth args,
             # we need to send them via HELLO
         if auth_args and self.protocol not in [2, "2"]:
@@ -496,6 +504,10 @@ async def send_packed_command(
 
     async def send_command(self, *args: Any, **kwargs: Any) -> None:
         """Pack and send a command to the Redis server"""
+        if isinstance(self.credential_provider, StreamingCredentialProvider):
+            await self._event_dispatcher.dispatch_async(
+                AsyncBeforeCommandExecutionEvent(args, self._init_auth_args, self, self.credential_provider)
+            )
         await self.send_packed_command(
             self.pack_command(*args), check_health=kwargs.get("check_health", True)
         )
@@ -1033,6 +1045,7 @@ def __init__(
         self._available_connections: List[AbstractConnection] = []
         self._in_use_connections: Set[AbstractConnection] = set()
         self.encoder_class = self.connection_kwargs.get("encoder_class", Encoder)
+        self._lock = threading.Lock()
 
     def __repr__(self):
         return (

redis/client.py

@@ -314,20 +314,22 @@ def __init__(
                             "cache_config": cache_config,
                         }
                     )
-            connection_pool = ConnectionPool(**kwargs, event_dispatcher=event_dispatcher)
+            connection_pool = ConnectionPool(**kwargs)
             event_dispatcher.dispatch(AfterPooledConnectionsInstantiationEvent(
                 [connection_pool],
+                ClientType.SYNC,
                 credential_provider
             ))
             self.auto_close_connection_pool = True
         else:
             self.auto_close_connection_pool = False
+            event_dispatcher.dispatch(AfterPooledConnectionsInstantiationEvent(
+                [connection_pool],
+                ClientType.SYNC,
+                credential_provider
+            ))
 
         self.connection_pool = connection_pool
-        event_dispatcher.dispatch(AfterPooledConnectionsInstantiationEvent(
-            [connection_pool],
-            credential_provider
-        ))
 
         if (cache_config or cache) and self.connection_pool.get_protocol() not in [
             3,

redis/connection.py

@@ -230,7 +230,7 @@ def __init__(
         credential_provider: Optional[CredentialProvider] = None,
         protocol: Optional[int] = 2,
         command_packer: Optional[Callable[[], None]] = None,
-        event_dispatcher: Optional[EventDispatcherInterface] = EventDispatcher()
+        event_dispatcher: Optional[EventDispatcher] = EventDispatcher()
     ):
         """
         Initialize a new Connection.

redis/event.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from enum import Enum
 from typing import List, Union, Optional
 
 from redis.credentials import StreamingCredentialProvider, CredentialProvider
@@ -13,6 +14,15 @@ def listen(self, event: object):
         pass
 
 
+class AsyncEventListenerInterface(ABC):
+    """
+    Represents an async listener for given event object.
+    """
+    @abstractmethod
+    async def listen(self, event: object):
+        pass
+
+
 class EventDispatcherInterface(ABC):
     """
     Represents a dispatcher that dispatches events to listeners associated with given event.
@@ -21,6 +31,10 @@ class EventDispatcherInterface(ABC):
     def dispatch(self, event: object):
         pass
 
+    @abstractmethod
+    async def dispatch_async(self, event: object):
+        pass
+
 
 class EventDispatcher(EventDispatcherInterface):
     # TODO: Make dispatcher to accept external mappings.
@@ -34,7 +48,10 @@ def __init__(self):
             ],
             AfterPooledConnectionsInstantiationEvent: [
                 RegisterReAuthForPooledConnections()
-            ]
+            ],
+            AsyncBeforeCommandExecutionEvent: [
+                AsyncReAuthBeforeCommandExecutionListener(),
+            ],
         }
 
     def dispatch(self, event: object):
@@ -43,6 +60,12 @@ def dispatch(self, event: object):
         for listener in listeners:
             listener.listen(event)
 
+    async def dispatch_async(self, event: object):
+        listeners = self._event_listeners_mapping.get(type(event))
+
+        for listener in listeners:
+            await listener.listen(event)
+
 
 class BeforeCommandExecutionEvent:
     """
@@ -71,6 +94,10 @@ def credential_provider(self) -> StreamingCredentialProvider:
         return self._credential_provider
 
 
+class AsyncBeforeCommandExecutionEvent(BeforeCommandExecutionEvent):
+    pass
+
+
 class AfterPooledConnectionsInstantiationEvent:
     """
     Event that will be fired after pooled connection instances was created.
@@ -111,14 +138,33 @@ def listen(self, event: BeforeCommandExecutionEvent):
             event.connection.read_response()
 
 
-class RegisterReAuthForPooledConnections(EventListenerInterface):
+class AsyncReAuthBeforeCommandExecutionListener(AsyncEventListenerInterface):
+    """
+    Async listener that performs re-authentication (if needed) for StreamingCredentialProviders before command execution
+    """
     def __init__(self):
-        self._event = None
+        self._current_cred = None
+
+    async def listen(self, event: AsyncBeforeCommandExecutionEvent):
+        if self._current_cred is None:
+            self._current_cred = event.initial_cred
+
+        credentials = await event.credential_provider.get_credentials_async()
 
+        if hash(credentials) != self._current_cred:
+            self._current_cred = hash(credentials)
+            await event.connection.send_command('AUTH', credentials[0], credentials[1])
+            await event.connection.read_response()
+
+
+class RegisterReAuthForPooledConnections(EventListenerInterface):
     """
     Listener that registers a re-authentication callback for pooled connections.
     Required by :class:`StreamingCredentialProvider`.
     """
+    def __init__(self):
+        self._event = None
+
     def listen(self, event: AfterPooledConnectionsInstantiationEvent):
         if isinstance(event.credential_provider, StreamingCredentialProvider):
             self._event = event

requirements.txt

@@ -1 +1,2 @@
-async-timeout>=4.0.3
+async-timeout>=4.0.3
+redispy-entraid-credentials @ git+https://github.com/redis-developer/redispy-entra-credentials.git/@main