StreamingCredentialProvider support

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/client.py

@@ -53,6 +53,8 @@
     list_or_args,
 )
 from redis.credentials import CredentialProvider
+from redis.event import EventDispatcher, AfterPooledConnectionsInstantiationEvent, ClientType, \
+    AfterSingleConnectionInstantiationEvent
 from redis.exceptions import (
     ConnectionError,
     ExecAbortError,
@@ -233,6 +235,7 @@ def __init__(
         redis_connect_func=None,
         credential_provider: Optional[CredentialProvider] = None,
         protocol: Optional[int] = 2,
+        event_dispatcher: Optional[EventDispatcher] = EventDispatcher(),
     ):
         """
         Initialize a new Redis client.
@@ -320,11 +323,22 @@ def __init__(
             # This arg only used if no pool is passed in
             self.auto_close_connection_pool = auto_close_connection_pool
             connection_pool = ConnectionPool(**kwargs)
+            event_dispatcher.dispatch(AfterPooledConnectionsInstantiationEvent(
+                [connection_pool],
+                ClientType.ASYNC,
+                credential_provider
+            ))
         else:
             # If a pool is passed in, do not close it
             self.auto_close_connection_pool = False
+            event_dispatcher.dispatch(AfterPooledConnectionsInstantiationEvent(
+                [connection_pool],
+                ClientType.ASYNC,
+                credential_provider
+            ))
 
         self.connection_pool = connection_pool
+        self._event_dispatcher = event_dispatcher
         self.single_connection_client = single_connection_client
         self.connection: Optional[Connection] = None
 
@@ -354,6 +368,10 @@ async def initialize(self: _RedisT) -> _RedisT:
             async with self._single_conn_lock:
                 if self.connection is None:
                     self.connection = await self.connection_pool.get_connection("_")
+
+            self._event_dispatcher.dispatch(
+                AfterSingleConnectionInstantiationEvent(self.connection, ClientType.ASYNC, self._single_conn_lock)
+            )
         return self
 
     def set_response_callback(self, command: str, callback: ResponseCallbackT):

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:
@@ -339,7 +345,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"]:
@@ -1039,6 +1047,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 = asyncio.Lock()
 
     def __repr__(self):
         return (
@@ -1058,13 +1067,14 @@ def can_get_connection(self) -> bool:
         )
 
     async def get_connection(self, command_name, *keys, **options):
-        """Get a connected connection from the pool"""
-        connection = self.get_available_connection()
-        try:
-            await self.ensure_connection(connection)
-        except BaseException:
-            await self.release(connection)
-            raise
+        async with self._lock:
+            """Get a connected connection from the pool"""
+            connection = self.get_available_connection()
+            try:
+                await self.ensure_connection(connection)
+            except BaseException:
+                await self.release(connection)
+                raise
 
         return connection
 
@@ -1147,6 +1157,14 @@ def set_retry(self, retry: "Retry") -> None:
         for conn in self._in_use_connections:
             conn.retry = retry
 
+    async def re_auth_callback(self, token):
+        async with self._lock:
+            for conn in self._available_connections:
+                await conn.send_command(
+                    'AUTH', token.try_get('oid'), token.get_value()
+                )
+                await conn.read_response()
+
 
 class BlockingConnectionPool(ConnectionPool):
     """

redis/client.py

@@ -1,3 +1,4 @@
+import asyncio
 import copy
 import re
 import threading
@@ -27,6 +28,8 @@
     UnixDomainSocketConnection,
 )
 from redis.credentials import CredentialProvider
+from redis.event import EventDispatcher, AfterPooledConnectionsInstantiationEvent, ClientType, \
+    AfterSingleConnectionInstantiationEvent
 from redis.exceptions import (
     ConnectionError,
     ExecAbortError,
@@ -213,6 +216,7 @@ def __init__(
         protocol: Optional[int] = 2,
         cache: Optional[CacheInterface] = None,
         cache_config: Optional[CacheConfig] = None,
+        event_dispatcher: Optional[EventDispatcher] = EventDispatcher(),
     ) -> None:
         """
         Initialize a new Redis client.
@@ -313,9 +317,19 @@ def __init__(
                         }
                     )
             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
 
@@ -325,9 +339,13 @@ def __init__(
         ]:
             raise RedisError("Client caching is only supported with RESP version 3")
 
+        self._connection_lock = threading.Lock()
         self.connection = None
         if single_connection_client:
             self.connection = self.connection_pool.get_connection("_")
+            event_dispatcher.dispatch(
+                AfterSingleConnectionInstantiationEvent(self.connection, ClientType.SYNC, self._connection_lock)
+            )
 
         self.response_callbacks = CaseInsensitiveDict(_RedisCallbacks)
 

redis/cluster.py

@@ -15,6 +15,7 @@
 from redis.commands.helpers import list_or_args
 from redis.connection import ConnectionPool, DefaultParser, parse_url
 from redis.crc import REDIS_CLUSTER_HASH_SLOTS, key_slot
+from redis.event import EventDispatcher, EventDispatcherInterface, AfterPooledConnectionsInstantiationEvent, ClientType
 from redis.exceptions import (
     AskError,
     AuthenticationError,
@@ -505,6 +506,7 @@ def __init__(
         address_remap: Optional[Callable[[Tuple[str, int]], Tuple[str, int]]] = None,
         cache: Optional[CacheInterface] = None,
         cache_config: Optional[CacheConfig] = None,
+        event_dispatcher: Optional[EventDispatcher] = EventDispatcher(),
         **kwargs,
     ):
         """
@@ -646,6 +648,7 @@ def __init__(
             address_remap=address_remap,
             cache=cache,
             cache_config=cache_config,
+            event_dispatcher=event_dispatcher,
             **kwargs,
         )
 
@@ -1332,6 +1335,7 @@ def __init__(
         cache: Optional[CacheInterface] = None,
         cache_config: Optional[CacheConfig] = None,
         cache_factory: Optional[CacheFactoryInterface] = None,
+        event_dispatcher: Optional[EventDispatcher] = EventDispatcher(),
         **kwargs,
     ):
         self.nodes_cache = {}
@@ -1353,6 +1357,8 @@ def __init__(
         if lock is None:
             lock = threading.Lock()
         self._lock = lock
+        self._event_dispatcher = event_dispatcher
+        self._credential_provider = self.connection_kwargs.get("credential_provider", None)
         self.initialize()
 
     def get_node(self, host=None, port=None, node_name=None):
@@ -1479,11 +1485,21 @@ def create_redis_connections(self, nodes):
         """
         This function will create a redis connection to all nodes in :nodes:
         """
+        connection_pools = []
         for node in nodes:
             if node.redis_connection is None:
                 node.redis_connection = self.create_redis_node(
                     host=node.host, port=node.port, **self.connection_kwargs
                 )
+                connection_pools.append(node.redis_connection.connection_pool)
+
+        self._event_dispatcher.dispatch(
+            AfterPooledConnectionsInstantiationEvent(
+                connection_pools,
+                ClientType.SYNC,
+                self._credential_provider
+            )
+        )
 
     def create_redis_node(self, host, port, **kwargs):
         if self.from_url:

redis/connection.py

@@ -23,7 +23,8 @@
 
 from ._parsers import Encoder, _HiredisParser, _RESP2Parser, _RESP3Parser
 from .backoff import NoBackoff
-from .credentials import CredentialProvider, UsernamePasswordCredentialProvider
+from .credentials import CredentialProvider, UsernamePasswordCredentialProvider, StreamingCredentialProvider
+from .event import EventDispatcherInterface, EventDispatcher, BeforeCommandExecutionEvent
 from .exceptions import (
     AuthenticationError,
     AuthenticationWrongNumberOfArgsError,
@@ -229,6 +230,7 @@
         credential_provider: Optional[CredentialProvider] = None,
         protocol: Optional[int] = 2,
         command_packer: Optional[Callable[[], None]] = None,
+        event_dispatcher: Optional[EventDispatcher] = EventDispatcher()
     ):
         """
         Initialize a new Connection.
@@ -283,6 +285,8 @@
         self.set_parser(parser_class)
         self._connect_callbacks = []
         self._buffer_cutoff = 6000
+        self._event_dispatcher = event_dispatcher
+        self._init_auth_args = None
         try:
             p = int(protocol)
         except TypeError:
@@ -408,6 +412,8 @@
                 or UsernamePasswordCredentialProvider(self.username, self.password)
             )
             auth_args = cred_provider.get_credentials()
+            self._init_auth_args = hash(auth_args)
+            print(auth_args)
 
         # if resp version is specified and we have auth args,
         # we need to send them via HELLO
@@ -1318,6 +1324,7 @@
         connection_kwargs.pop("cache", None)
         connection_kwargs.pop("cache_config", None)
 
+
         # a lock to protect the critical section in _checkpid().
         # this lock is acquired when the process id changes, such as
         # after a fork. during this time, multiple threads in the child
@@ -1517,6 +1524,14 @@
         for conn in self._in_use_connections:
             conn.retry = retry
 
+    def re_auth_callback(self, token):
+        with self._lock:
+            for conn in self._available_connections:
+                conn.send_command(
+                    'AUTH', token.try_get('oid'), token.get_value()
+                )
+                conn.read_response()
+
 
 class BlockingConnectionPool(ConnectionPool):
     """

redis/credentials.py

@@ -1,4 +1,8 @@
-from typing import Optional, Tuple, Union
+import logging
+from abc import ABC, abstractmethod
+from typing import Optional, Tuple, Union, Callable, Any
+
+logger = logging.getLogger(__name__)
 
 
 class CredentialProvider:
@@ -9,6 +13,33 @@ class CredentialProvider:
     def get_credentials(self) -> Union[Tuple[str], Tuple[str, str]]:
         raise NotImplementedError("get_credentials must be implemented")
 
+    async def get_credentials_async(self) -> Union[Tuple[str], Tuple[str, str]]:
+        logger.warning("This method is added for backward compatability. Please override it in your implementation.")
+        return self.get_credentials()
+
+
+class StreamingCredentialProvider(CredentialProvider, ABC):
+    """
+    Credential provider that streams credentials in the background.
+    """
+    @abstractmethod
+    def on_next(self, callback: Callable[[Any], None]):
+        """
+        Specifies the callback that should be invoked when the next credentials will be retrieved.
+
+        :param callback: Callback with
+        :return:
+        """
+        pass
+
+    @abstractmethod
+    def on_error(self, callback: Callable[[Exception], None]):
+        pass
+
+    @abstractmethod
+    def is_streaming(self) -> bool:
+        pass
+
 
 class UsernamePasswordCredentialProvider(CredentialProvider):
     """
@@ -24,3 +55,6 @@ def get_credentials(self):
         if self.username:
             return self.username, self.password
         return (self.password,)
+
+    async def get_credentials_async(self) -> Union[Tuple[str], Tuple[str, str]]:
+        return self.get_credentials()