StreamingCredentialProvider support

Author: vladvildanov

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 @@ def __init__(
         credential_provider: Optional[CredentialProvider] = None,
         protocol: Optional[int] = 2,
         command_packer: Optional[Callable[[], None]] = None,
+        event_dispatcher: Optional[EventDispatcherInterface] = EventDispatcher()
     ):
         """
         Initialize a new Connection.
@@ -283,6 +285,8 @@ def __init__(
         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,7 @@ def on_connect(self):
                 or UsernamePasswordCredentialProvider(self.username, self.password)
             )
             auth_args = cred_provider.get_credentials()
+            self._init_auth_args = hash(auth_args)
 
         # if resp version is specified and we have auth args,
         # we need to send them via HELLO
@@ -553,6 +558,10 @@ def send_packed_command(self, command, check_health=True):
 
     def send_command(self, *args, **kwargs):
         """Pack and send a command to the Redis server"""
+        if isinstance(self.credential_provider, StreamingCredentialProvider):
+            self._event_dispatcher.dispatch(
+                BeforeCommandExecutionEvent(args, self._init_auth_args, self, self.credential_provider)
+            )
         self.send_packed_command(
             self._command_packer.pack(*args),
             check_health=kwargs.get("check_health", True),
@@ -1318,6 +1327,12 @@ def __init__(
         connection_kwargs.pop("cache", None)
         connection_kwargs.pop("cache_config", None)
 
+        cred_provider = connection_kwargs.get("credential_provider")
+
+        if cred_provider is not None and isinstance(cred_provider, StreamingCredentialProvider):
+            cred_provider.on_next(self._re_auth)
+
+
         # 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 +1532,14 @@ def set_retry(self, retry: "Retry") -> None:
         for conn in self._in_use_connections:
             conn.retry = retry
 
+    def _re_auth(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

@@ -12,8 +12,21 @@ def get_credentials(self) -> Union[Tuple[str], Tuple[str, str]]:
 
 
 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

redis/event.py

@@ -0,0 +1,88 @@
+from abc import ABC, abstractmethod
+
+from redis.credentials import StreamingCredentialProvider
+
+
+class EventListenerInterface(ABC):
+    """
+    Represents a listener for given event object.
+    """
+    @abstractmethod
+    def listen(self, event: object):
+        pass
+
+
+class EventDispatcherInterface(ABC):
+    """
+    Represents a dispatcher that dispatches events to listeners associated with given event.
+    """
+    @abstractmethod
+    def dispatch(self, event: object):
+        pass
+
+
+class EventDispatcher(EventDispatcherInterface):
+
+    def __init__(self):
+        """
+        Mapping should be extended for any new events or listeners to be added.
+        """
+        self._event_listeners_mapping = {
+            BeforeCommandExecutionEvent: [
+                ReAuthBeforeCommandExecutionListener(),
+            ]
+        }
+
+    def dispatch(self, event: object):
+        listeners = self._event_listeners_mapping.get(type(event))
+
+        for listener in listeners:
+            listener.listen(event)
+
+
+class BeforeCommandExecutionEvent:
+    """
+    Event that will be fired before each command execution.
+    """
+    def __init__(self, command, initial_cred, connection, credential_provider: StreamingCredentialProvider):
+        self._command = command
+        self._initial_cred = initial_cred
+        self._credential_provider = credential_provider
+        self._connection = connection
+
+    @property
+    def command(self):
+        return self._command
+
+    @property
+    def initial_cred(self):
+        return self._initial_cred
+
+    @property
+    def connection(self):
+        return self._connection
+
+    @property
+    def credential_provider(self) -> StreamingCredentialProvider:
+        return self._credential_provider
+
+
+class ReAuthBeforeCommandExecutionListener(EventListenerInterface):
+    """
+    Listener that performs re-authentication (if needed) for StreamingCredentialProviders before command execution.
+    """
+    def __init__(self):
+        self._current_cred = None
+
+    def listen(self, event: BeforeCommandExecutionEvent):
+        if self._current_cred is None:
+            self._current_cred = event.initial_cred
+
+        credentials = event.credential_provider.get_credentials()
+
+        print(hash(credentials) != self._current_cred)
+
+        if hash(credentials) != self._current_cred:
+            self._current_cred = hash(credentials)
+            event.connection.send_command('AUTH', credentials[0], credentials[1])
+            event.connection.read_response()