redis 5.2.1__py3-none-any.whl → 5.3.0b3__py3-none-any.whl

redis/auth/token_manager.py

@@ -0,0 +1,370 @@
+import asyncio
+import logging
+import threading
+from datetime import datetime, timezone
+from time import sleep
+from typing import Any, Awaitable, Callable, Union
+
+from redis.auth.err import RequestTokenErr, TokenRenewalErr
+from redis.auth.idp import IdentityProviderInterface
+from redis.auth.token import TokenResponse
+
+logger = logging.getLogger(__name__)
+
+
+class CredentialsListener:
+    """
+    Listeners that will be notified on events related to credentials.
+    Accepts callbacks and awaitable callbacks.
+    """
+
+    def __init__(self):
+        self._on_next = None
+        self._on_error = None
+
+    @property
+    def on_next(self) -> Union[Callable[[Any], None], Awaitable]:
+        return self._on_next
+
+    @on_next.setter
+    def on_next(self, callback: Union[Callable[[Any], None], Awaitable]) -> None:
+        self._on_next = callback
+
+    @property
+    def on_error(self) -> Union[Callable[[Exception], None], Awaitable]:
+        return self._on_error
+
+    @on_error.setter
+    def on_error(self, callback: Union[Callable[[Exception], None], Awaitable]) -> None:
+        self._on_error = callback
+
+
+class RetryPolicy:
+    def __init__(self, max_attempts: int, delay_in_ms: float):
+        self.max_attempts = max_attempts
+        self.delay_in_ms = delay_in_ms
+
+    def get_max_attempts(self) -> int:
+        """
+        Retry attempts before exception will be thrown.
+
+        :return: int
+        """
+        return self.max_attempts
+
+    def get_delay_in_ms(self) -> float:
+        """
+        Delay between retries in seconds.
+
+        :return: int
+        """
+        return self.delay_in_ms
+
+
+class TokenManagerConfig:
+    def __init__(
+        self,
+        expiration_refresh_ratio: float,
+        lower_refresh_bound_millis: int,
+        token_request_execution_timeout_in_ms: int,
+        retry_policy: RetryPolicy,
+    ):
+        self._expiration_refresh_ratio = expiration_refresh_ratio
+        self._lower_refresh_bound_millis = lower_refresh_bound_millis
+        self._token_request_execution_timeout_in_ms = (
+            token_request_execution_timeout_in_ms
+        )
+        self._retry_policy = retry_policy
+
+    def get_expiration_refresh_ratio(self) -> float:
+        """
+        Represents the ratio of a token's lifetime at which a refresh should be triggered. # noqa: E501
+        For example, a value of 0.75 means the token should be refreshed
+        when 75% of its lifetime has elapsed (or when 25% of its lifetime remains).
+
+        :return: float
+        """
+
+        return self._expiration_refresh_ratio
+
+    def get_lower_refresh_bound_millis(self) -> int:
+        """
+        Represents the minimum time in milliseconds before token expiration
+        to trigger a refresh, in milliseconds.
+        This value sets a fixed lower bound for when a token refresh should occur,
+        regardless of the token's total lifetime.
+        If set to 0 there will be no lower bound and the refresh will be triggered
+        based on the expirationRefreshRatio only.
+
+        :return: int
+        """
+        return self._lower_refresh_bound_millis
+
+    def get_token_request_execution_timeout_in_ms(self) -> int:
+        """
+        Represents the maximum time in milliseconds to wait
+        for a token request to complete.
+
+        :return: int
+        """
+        return self._token_request_execution_timeout_in_ms
+
+    def get_retry_policy(self) -> RetryPolicy:
+        """
+        Represents the retry policy for token requests.
+
+        :return: RetryPolicy
+        """
+        return self._retry_policy
+
+
+class TokenManager:
+    def __init__(
+        self, identity_provider: IdentityProviderInterface, config: TokenManagerConfig
+    ):
+        self._idp = identity_provider
+        self._config = config
+        self._next_timer = None
+        self._listener = None
+        self._init_timer = None
+        self._retries = 0
+
+    def __del__(self):
+        logger.info("Token manager are disposed")
+        self.stop()
+
+    def start(
+        self,
+        listener: CredentialsListener,
+        skip_initial: bool = False,
+    ) -> Callable[[], None]:
+        self._listener = listener
+
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            # Run loop in a separate thread to unblock main thread.
+            loop = asyncio.new_event_loop()
+            thread = threading.Thread(
+                target=_start_event_loop_in_thread, args=(loop,), daemon=True
+            )
+            thread.start()
+
+        # Event to block for initial execution.
+        init_event = asyncio.Event()
+        self._init_timer = loop.call_later(
+            0, self._renew_token, skip_initial, init_event
+        )
+        logger.info("Token manager started")
+
+        # Blocks in thread-safe manner.
+        asyncio.run_coroutine_threadsafe(init_event.wait(), loop).result()
+        return self.stop
+
+    async def start_async(
+        self,
+        listener: CredentialsListener,
+        block_for_initial: bool = False,
+        initial_delay_in_ms: float = 0,
+        skip_initial: bool = False,
+    ) -> Callable[[], None]:
+        self._listener = listener
+
+        loop = asyncio.get_running_loop()
+        init_event = asyncio.Event()
+
+        # Wraps the async callback with async wrapper to schedule with loop.call_later()
+        wrapped = _async_to_sync_wrapper(
+            loop, self._renew_token_async, skip_initial, init_event
+        )
+        self._init_timer = loop.call_later(initial_delay_in_ms / 1000, wrapped)
+        logger.info("Token manager started")
+
+        if block_for_initial:
+            await init_event.wait()
+
+        return self.stop
+
+    def stop(self):
+        if self._init_timer is not None:
+            self._init_timer.cancel()
+        if self._next_timer is not None:
+            self._next_timer.cancel()
+
+    def acquire_token(self, force_refresh=False) -> TokenResponse:
+        try:
+            token = self._idp.request_token(force_refresh)
+        except RequestTokenErr as e:
+            if self._retries < self._config.get_retry_policy().get_max_attempts():
+                self._retries += 1
+                sleep(self._config.get_retry_policy().get_delay_in_ms() / 1000)
+                return self.acquire_token(force_refresh)
+            else:
+                raise e
+
+        self._retries = 0
+        return TokenResponse(token)
+
+    async def acquire_token_async(self, force_refresh=False) -> TokenResponse:
+        try:
+            token = self._idp.request_token(force_refresh)
+        except RequestTokenErr as e:
+            if self._retries < self._config.get_retry_policy().get_max_attempts():
+                self._retries += 1
+                await asyncio.sleep(
+                    self._config.get_retry_policy().get_delay_in_ms() / 1000
+                )
+                return await self.acquire_token_async(force_refresh)
+            else:
+                raise e
+
+        self._retries = 0
+        return TokenResponse(token)
+
+    def _calculate_renewal_delay(self, expire_date: float, issue_date: float) -> float:
+        delay_for_lower_refresh = self._delay_for_lower_refresh(expire_date)
+        delay_for_ratio_refresh = self._delay_for_ratio_refresh(expire_date, issue_date)
+        delay = min(delay_for_ratio_refresh, delay_for_lower_refresh)
+
+        return 0 if delay < 0 else delay / 1000
+
+    def _delay_for_lower_refresh(self, expire_date: float):
+        return (
+            expire_date
+            - self._config.get_lower_refresh_bound_millis()
+            - (datetime.now(timezone.utc).timestamp() * 1000)
+        )
+
+    def _delay_for_ratio_refresh(self, expire_date: float, issue_date: float):
+        token_ttl = expire_date - issue_date
+        refresh_before = token_ttl - (
+            token_ttl * self._config.get_expiration_refresh_ratio()
+        )
+
+        return (
+            expire_date
+            - refresh_before
+            - (datetime.now(timezone.utc).timestamp() * 1000)
+        )
+
+    def _renew_token(
+        self, skip_initial: bool = False, init_event: asyncio.Event = None
+    ):
+        """
+        Task to renew token from identity provider.
+        Schedules renewal tasks based on token TTL.
+        """
+
+        try:
+            token_res = self.acquire_token(force_refresh=True)
+            delay = self._calculate_renewal_delay(
+                token_res.get_token().get_expires_at_ms(),
+                token_res.get_token().get_received_at_ms(),
+            )
+
+            if token_res.get_token().is_expired():
+                raise TokenRenewalErr("Requested token is expired")
+
+            if self._listener.on_next is None:
+                logger.warning(
+                    "No registered callback for token renewal task. Renewal cancelled"
+                )
+                return
+
+            if not skip_initial:
+                try:
+                    self._listener.on_next(token_res.get_token())
+                except Exception as e:
+                    raise TokenRenewalErr(e)
+
+            if delay <= 0:
+                return
+
+            loop = asyncio.get_running_loop()
+            self._next_timer = loop.call_later(delay, self._renew_token)
+            logger.info(f"Next token renewal scheduled in {delay} seconds")
+            return token_res
+        except Exception as e:
+            if self._listener.on_error is None:
+                raise e
+
+            self._listener.on_error(e)
+        finally:
+            if init_event:
+                init_event.set()
+
+    async def _renew_token_async(
+        self, skip_initial: bool = False, init_event: asyncio.Event = None
+    ):
+        """
+        Async task to renew tokens from identity provider.
+        Schedules renewal tasks based on token TTL.
+        """
+
+        try:
+            token_res = await self.acquire_token_async(force_refresh=True)
+            delay = self._calculate_renewal_delay(
+                token_res.get_token().get_expires_at_ms(),
+                token_res.get_token().get_received_at_ms(),
+            )
+
+            if token_res.get_token().is_expired():
+                raise TokenRenewalErr("Requested token is expired")
+
+            if self._listener.on_next is None:
+                logger.warning(
+                    "No registered callback for token renewal task. Renewal cancelled"
+                )
+                return
+
+            if not skip_initial:
+                try:
+                    await self._listener.on_next(token_res.get_token())
+                except Exception as e:
+                    raise TokenRenewalErr(e)
+
+            if delay <= 0:
+                return
+
+            loop = asyncio.get_running_loop()
+            wrapped = _async_to_sync_wrapper(loop, self._renew_token_async)
+            logger.info(f"Next token renewal scheduled in {delay} seconds")
+            loop.call_later(delay, wrapped)
+        except Exception as e:
+            if self._listener.on_error is None:
+                raise e
+
+            await self._listener.on_error(e)
+        finally:
+            if init_event:
+                init_event.set()
+
+
+def _async_to_sync_wrapper(loop, coro_func, *args, **kwargs):
+    """
+    Wraps an asynchronous function so it can be used with loop.call_later.
+
+    :param loop: The event loop in which the coroutine will be executed.
+    :param coro_func: The coroutine function to wrap.
+    :param args: Positional arguments to pass to the coroutine function.
+    :param kwargs: Keyword arguments to pass to the coroutine function.
+    :return: A regular function suitable for loop.call_later.
+    """
+
+    def wrapped():
+        # Schedule the coroutine in the event loop
+        asyncio.ensure_future(coro_func(*args, **kwargs), loop=loop)
+
+    return wrapped
+
+
+def _start_event_loop_in_thread(event_loop: asyncio.AbstractEventLoop):
+    """
+    Starts event loop in a thread.
+    Used to be able to schedule tasks using loop.call_later.
+
+    :param event_loop:
+    :return:
+    """
+    asyncio.set_event_loop(event_loop)
+    event_loop.run_forever()

redis/client.py

@@ -27,6 +27,13 @@ from redis.connection import (
     UnixDomainSocketConnection,
 )
 from redis.credentials import CredentialProvider
+from redis.event import (
+    AfterPooledConnectionsInstantiationEvent,
+    AfterPubSubConnectionInstantiationEvent,
+    AfterSingleConnectionInstantiationEvent,
+    ClientType,
+    EventDispatcher,
+)
 from redis.exceptions import (
     ConnectionError,
     ExecAbortError,
@@ -213,6 +220,7 @@ class Redis(RedisModuleCommands, CoreCommands, SentinelCommands):
         protocol: Optional[int] = 2,
         cache: Optional[CacheInterface] = None,
         cache_config: Optional[CacheConfig] = None,
+        event_dispatcher: Optional[EventDispatcher] = None,
     ) -> None:
         """
         Initialize a new Redis client.
@@ -227,6 +235,10 @@ class Redis(RedisModuleCommands, CoreCommands, SentinelCommands):
             if `True`, connection pool is not used. In that case `Redis`
             instance use is not thread safe.
         """
+        if event_dispatcher is None:
+            self._event_dispatcher = EventDispatcher()
+        else:
+            self._event_dispatcher = event_dispatcher
         if not connection_pool:
             if charset is not None:
                 warnings.warn(
@@ -313,9 +325,19 @@ class Redis(RedisModuleCommands, CoreCommands, SentinelCommands):
                         }
                     )
             connection_pool = ConnectionPool(**kwargs)
+            self._event_dispatcher.dispatch(
+                AfterPooledConnectionsInstantiationEvent(
+                    [connection_pool], ClientType.SYNC, credential_provider
+                )
+            )
             self.auto_close_connection_pool = True
         else:
             self.auto_close_connection_pool = False
+            self._event_dispatcher.dispatch(
+                AfterPooledConnectionsInstantiationEvent(
+                    [connection_pool], ClientType.SYNC, credential_provider
+                )
+            )
 
         self.connection_pool = connection_pool
 
@@ -325,9 +347,16 @@ class Redis(RedisModuleCommands, CoreCommands, SentinelCommands):
         ]:
             raise RedisError("Client caching is only supported with RESP version 3")
 
+        self.single_connection_lock = threading.Lock()
         self.connection = None
-        if single_connection_client:
+        self._single_connection_client = single_connection_client
+        if self._single_connection_client:
             self.connection = self.connection_pool.get_connection("_")
+            self._event_dispatcher.dispatch(
+                AfterSingleConnectionInstantiationEvent(
+                    self.connection, ClientType.SYNC, self.single_connection_lock
+                )
+            )
 
         self.response_callbacks = CaseInsensitiveDict(_RedisCallbacks)
 
@@ -500,7 +529,9 @@ class Redis(RedisModuleCommands, CoreCommands, SentinelCommands):
         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)
@@ -563,6 +594,9 @@ class Redis(RedisModuleCommands, CoreCommands, SentinelCommands):
         pool = self.connection_pool
         command_name = args[0]
         conn = self.connection or pool.get_connection(command_name, **options)
+
+        if self._single_connection_client:
+            self.single_connection_lock.acquire()
         try:
             return conn.retry.call_with_retry(
                 lambda: self._send_command_parse_response(
@@ -571,6 +605,8 @@ class Redis(RedisModuleCommands, CoreCommands, SentinelCommands):
                 lambda error: self._disconnect_raise(conn, error),
             )
         finally:
+            if self._single_connection_client:
+                self.single_connection_lock.release()
             if not self.connection:
                 pool.release(conn)
 
@@ -691,6 +727,7 @@ class PubSub:
         ignore_subscribe_messages: bool = False,
         encoder: Optional["Encoder"] = None,
         push_handler_func: Union[None, Callable[[str], None]] = None,
+        event_dispatcher: Optional["EventDispatcher"] = None,
     ):
         self.connection_pool = connection_pool
         self.shard_hint = shard_hint
@@ -701,6 +738,11 @@ class PubSub:
         # to lookup channel and pattern names for callback handlers.
         self.encoder = encoder
         self.push_handler_func = push_handler_func
+        if event_dispatcher is None:
+            self._event_dispatcher = EventDispatcher()
+        else:
+            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)
@@ -791,11 +833,17 @@ class PubSub:
             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:
         """

redis/cluster.py

@@ -15,6 +15,12 @@ from redis.commands import READ_COMMANDS, RedisClusterCommands
 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 (
+    AfterPooledConnectionsInstantiationEvent,
+    AfterPubSubConnectionInstantiationEvent,
+    ClientType,
+    EventDispatcher,
+)
 from redis.exceptions import (
     AskError,
     AuthenticationError,
@@ -505,6 +511,7 @@ class RedisCluster(AbstractRedisCluster, RedisClusterCommands):
         address_remap: Optional[Callable[[Tuple[str, int]], Tuple[str, int]]] = None,
         cache: Optional[CacheInterface] = None,
         cache_config: Optional[CacheConfig] = None,
+        event_dispatcher: Optional[EventDispatcher] = None,
         **kwargs,
     ):
         """
@@ -638,6 +645,10 @@ class RedisCluster(AbstractRedisCluster, RedisClusterCommands):
         self.read_from_replicas = read_from_replicas
         self.reinitialize_counter = 0
         self.reinitialize_steps = reinitialize_steps
+        if event_dispatcher is None:
+            self._event_dispatcher = EventDispatcher()
+        else:
+            self._event_dispatcher = event_dispatcher
         self.nodes_manager = NodesManager(
             startup_nodes=startup_nodes,
             from_url=from_url,
@@ -646,6 +657,7 @@ class RedisCluster(AbstractRedisCluster, RedisClusterCommands):
             address_remap=address_remap,
             cache=cache,
             cache_config=cache_config,
+            event_dispatcher=self._event_dispatcher,
             **kwargs,
         )
 
@@ -1332,6 +1344,7 @@ class NodesManager:
         cache: Optional[CacheInterface] = None,
         cache_config: Optional[CacheConfig] = None,
         cache_factory: Optional[CacheFactoryInterface] = None,
+        event_dispatcher: Optional[EventDispatcher] = None,
         **kwargs,
     ):
         self.nodes_cache = {}
@@ -1353,6 +1366,13 @@ class NodesManager:
         if lock is None:
             lock = threading.Lock()
         self._lock = lock
+        if event_dispatcher is None:
+            self._event_dispatcher = EventDispatcher()
+        else:
+            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 +1499,19 @@ class NodesManager:
         """
         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:
@@ -1698,6 +1726,7 @@ class ClusterPubSub(PubSub):
         host=None,
         port=None,
         push_handler_func=None,
+        event_dispatcher: Optional["EventDispatcher"] = None,
         **kwargs,
     ):
         """
@@ -1723,10 +1752,15 @@ class ClusterPubSub(PubSub):
         self.cluster = redis_cluster
         self.node_pubsub_mapping = {}
         self._pubsubs_generator = self._pubsubs_generator()
+        if event_dispatcher is None:
+            self._event_dispatcher = EventDispatcher()
+        else:
+            self._event_dispatcher = event_dispatcher
         super().__init__(
             connection_pool=connection_pool,
             encoder=redis_cluster.encoder,
             push_handler_func=push_handler_func,
+            event_dispatcher=self._event_dispatcher,
             **kwargs,
         )
 
@@ -1813,6 +1847,11 @@ class ClusterPubSub(PubSub):
             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
         self._execute(connection, connection.send_command, *args)