redis 5.2.1__tar.gz → 5.3.0b3__tar.gz

{redis-5.2.1/redis.egg-info → redis-5.3.0b3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: redis
-Version: 5.2.1
+Version: 5.3.0b3
 Summary: Python client for Redis database and key-value store
 Home-page: https://github.com/redis/redis-py
 Author: Redis Inc.

{redis-5.2.1 → redis-5.3.0b3}/redis/asyncio/client.py

@@ -53,6 +53,13 @@ from redis.commands import (
     list_or_args,
 )
 from redis.credentials import CredentialProvider
+from redis.event import (
+    AfterPooledConnectionsInstantiationEvent,
+    AfterPubSubConnectionInstantiationEvent,
+    AfterSingleConnectionInstantiationEvent,
+    ClientType,
+    EventDispatcher,
+)
 from redis.exceptions import (
     ConnectionError,
     ExecAbortError,
@@ -233,6 +240,7 @@ class Redis(
         redis_connect_func=None,
         credential_provider: Optional[CredentialProvider] = None,
         protocol: Optional[int] = 2,
+        event_dispatcher: Optional[EventDispatcher] = None,
     ):
         """
         Initialize a new Redis client.
@@ -242,6 +250,10 @@ class Redis(
         To retry on TimeoutError, `retry_on_timeout` can also be set to `True`.
         """
         kwargs: Dict[str, Any]
+        if event_dispatcher is None:
+            self._event_dispatcher = EventDispatcher()
+        else:
+            self._event_dispatcher = event_dispatcher
         # auto_close_connection_pool only has an effect if connection_pool is
         # None. It is assumed that if connection_pool is not None, the user
         # wants to manage the connection pool themselves.
@@ -320,9 +332,19 @@ class Redis(
             # This arg only used if no pool is passed in
             self.auto_close_connection_pool = auto_close_connection_pool
             connection_pool = ConnectionPool(**kwargs)
+            self._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
+            self._event_dispatcher.dispatch(
+                AfterPooledConnectionsInstantiationEvent(
+                    [connection_pool], ClientType.ASYNC, credential_provider
+                )
+            )
 
         self.connection_pool = connection_pool
         self.single_connection_client = single_connection_client
@@ -354,6 +376,12 @@ class Redis(
             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):
@@ -521,7 +549,9 @@ class Redis(
         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)
@@ -759,7 +789,12 @@ class PubSub:
         ignore_subscribe_messages: bool = False,
         encoder=None,
         push_handler_func: Optional[Callable] = None,
+        event_dispatcher: Optional["EventDispatcher"] = None,
     ):
+        if event_dispatcher is None:
+            self._event_dispatcher = EventDispatcher()
+        else:
+            self._event_dispatcher = event_dispatcher
         self.connection_pool = connection_pool
         self.shard_hint = shard_hint
         self.ignore_subscribe_messages = ignore_subscribe_messages
@@ -876,6 +911,12 @@ class PubSub:
         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

{redis-5.2.1 → redis-5.3.0b3}/redis/asyncio/cluster.py

@@ -29,6 +29,7 @@ from redis.asyncio.client import ResponseCallbackT
 from redis.asyncio.connection import Connection, DefaultParser, SSLConnection, parse_url
 from redis.asyncio.lock import Lock
 from redis.asyncio.retry import Retry
+from redis.auth.token import TokenInterface
 from redis.backoff import default_backoff
 from redis.client import EMPTY_RESPONSE, NEVER_DECODE, AbstractRedis
 from redis.cluster import (
@@ -45,6 +46,7 @@ from redis.cluster import (
 from redis.commands import READ_COMMANDS, AsyncRedisClusterCommands
 from redis.crc import REDIS_CLUSTER_HASH_SLOTS, key_slot
 from redis.credentials import CredentialProvider
+from redis.event import AfterAsyncClusterInstantiationEvent, EventDispatcher
 from redis.exceptions import (
     AskError,
     BusyLoadingError,
@@ -57,6 +59,7 @@ from redis.exceptions import (
     MaxConnectionsError,
     MovedError,
     RedisClusterException,
+    RedisError,
     ResponseError,
     SlotNotCoveredError,
     TimeoutError,
@@ -270,6 +273,7 @@ class RedisCluster(AbstractRedis, AbstractRedisCluster, AsyncRedisClusterCommand
         ssl_ciphers: Optional[str] = None,
         protocol: Optional[int] = 2,
         address_remap: Optional[Callable[[Tuple[str, int]], Tuple[str, int]]] = None,
+        event_dispatcher: Optional[EventDispatcher] = None,
     ) -> None:
         if db:
             raise RedisClusterException(
@@ -366,11 +370,17 @@ class RedisCluster(AbstractRedis, AbstractRedisCluster, AsyncRedisClusterCommand
         if host and port:
             startup_nodes.append(ClusterNode(host, port, **self.connection_kwargs))
 
+        if event_dispatcher is None:
+            self._event_dispatcher = EventDispatcher()
+        else:
+            self._event_dispatcher = event_dispatcher
+
         self.nodes_manager = NodesManager(
             startup_nodes,
             require_full_coverage,
             kwargs,
             address_remap=address_remap,
+            event_dispatcher=self._event_dispatcher,
         )
         self.encoder = Encoder(encoding, encoding_errors, decode_responses)
         self.read_from_replicas = read_from_replicas
@@ -929,6 +939,8 @@ class ClusterNode:
     __slots__ = (
         "_connections",
         "_free",
+        "_lock",
+        "_event_dispatcher",
         "connection_class",
         "connection_kwargs",
         "host",
@@ -966,6 +978,9 @@ class ClusterNode:
 
         self._connections: List[Connection] = []
         self._free: Deque[Connection] = collections.deque(maxlen=self.max_connections)
+        self._event_dispatcher = self.connection_kwargs.get("event_dispatcher", None)
+        if self._event_dispatcher is None:
+            self._event_dispatcher = EventDispatcher()
 
     def __repr__(self) -> str:
         return (
@@ -1082,10 +1097,38 @@ class ClusterNode:
 
         return ret
 
+    async def re_auth_callback(self, token: TokenInterface):
+        tmp_queue = collections.deque()
+        while self._free:
+            conn = self._free.popleft()
+            await conn.retry.call_with_retry(
+                lambda: conn.send_command(
+                    "AUTH", token.try_get("oid"), token.get_value()
+                ),
+                lambda error: self._mock(error),
+            )
+            await conn.retry.call_with_retry(
+                lambda: conn.read_response(), lambda error: self._mock(error)
+            )
+            tmp_queue.append(conn)
+
+        while tmp_queue:
+            conn = tmp_queue.popleft()
+            self._free.append(conn)
+
+    async def _mock(self, error: RedisError):
+        """
+        Dummy functions, needs to be passed as error callback to retry object.
+        :param error:
+        :return:
+        """
+        pass
+
 
 class NodesManager:
     __slots__ = (
         "_moved_exception",
+        "_event_dispatcher",
         "connection_kwargs",
         "default_node",
         "nodes_cache",
@@ -1102,6 +1145,7 @@ class NodesManager:
         require_full_coverage: bool,
         connection_kwargs: Dict[str, Any],
         address_remap: Optional[Callable[[Tuple[str, int]], Tuple[str, int]]] = None,
+        event_dispatcher: Optional[EventDispatcher] = None,
     ) -> None:
         self.startup_nodes = {node.name: node for node in startup_nodes}
         self.require_full_coverage = require_full_coverage
@@ -1113,6 +1157,10 @@ class NodesManager:
         self.slots_cache: Dict[int, List["ClusterNode"]] = {}
         self.read_load_balancer = LoadBalancer()
         self._moved_exception: MovedError = None
+        if event_dispatcher is None:
+            self._event_dispatcher = EventDispatcher()
+        else:
+            self._event_dispatcher = event_dispatcher
 
     def get_node(
         self,
@@ -1230,6 +1278,12 @@ class NodesManager:
             try:
                 # Make sure cluster mode is enabled on this node
                 try:
+                    self._event_dispatcher.dispatch(
+                        AfterAsyncClusterInstantiationEvent(
+                            self.nodes_cache,
+                            self.connection_kwargs.get("credential_provider", None),
+                        )
+                    )
                     cluster_slots = await startup_node.execute_command("CLUSTER SLOTS")
                 except ResponseError:
                     raise RedisClusterException(

{redis-5.2.1 → redis-5.3.0b3}/redis/asyncio/connection.py

@@ -27,6 +27,8 @@ from typing import (
 )
 from urllib.parse import ParseResult, parse_qs, unquote, urlparse
 
+from ..auth.token import TokenInterface
+from ..event import AsyncAfterConnectionReleasedEvent, EventDispatcher
 from ..utils import format_error_message
 
 # the functionality is available in 3.11.x but has a major issue before
@@ -148,6 +150,7 @@ class AbstractConnection:
         encoder_class: Type[Encoder] = Encoder,
         credential_provider: Optional[CredentialProvider] = None,
         protocol: Optional[int] = 2,
+        event_dispatcher: Optional[EventDispatcher] = None,
     ):
         if (username or password) and credential_provider is not None:
             raise DataError(
@@ -156,6 +159,10 @@ class AbstractConnection:
                 "1. 'password' and (optional) 'username'\n"
                 "2. 'credential_provider'"
             )
+        if event_dispatcher is None:
+            self._event_dispatcher = EventDispatcher()
+        else:
+            self._event_dispatcher = event_dispatcher
         self.db = db
         self.client_name = client_name
         self.lib_name = lib_name
@@ -195,6 +202,8 @@ class AbstractConnection:
         self.set_parser(parser_class)
         self._connect_callbacks: List[weakref.WeakMethod[ConnectCallbackT]] = []
         self._buffer_cutoff = 6000
+        self._re_auth_token: Optional[TokenInterface] = None
+
         try:
             p = int(protocol)
         except TypeError:
@@ -327,6 +336,9 @@ class AbstractConnection:
     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)
@@ -339,7 +351,8 @@ class AbstractConnection:
                 self.credential_provider
                 or UsernamePasswordCredentialProvider(self.username, self.password)
             )
-            auth_args = cred_provider.get_credentials()
+            auth_args = await cred_provider.get_credentials_async()
+
             # 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"]:
@@ -661,6 +674,19 @@ class AbstractConnection:
         while not self._socket_is_empty():
             await self.read_response(push_request=True)
 
+    def set_re_auth_token(self, token: TokenInterface):
+        self._re_auth_token = token
+
+    async def re_auth(self):
+        if self._re_auth_token is not None:
+            await self.send_command(
+                "AUTH",
+                self._re_auth_token.try_get("oid"),
+                self._re_auth_token.get_value(),
+            )
+            await self.read_response()
+            self._re_auth_token = None
+
 
 class Connection(AbstractConnection):
     "Manages TCP communication to and from a Redis server"
@@ -1039,6 +1065,10 @@ class ConnectionPool:
         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()
+        self._event_dispatcher = self.connection_kwargs.get("event_dispatcher", None)
+        if self._event_dispatcher is None:
+            self._event_dispatcher = EventDispatcher()
 
     def __repr__(self):
         return (
@@ -1058,13 +1088,14 @@ class ConnectionPool:
         )
 
     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
 
@@ -1114,6 +1145,9 @@ class ConnectionPool:
         # not doing so is an error that will cause an exception here.
         self._in_use_connections.remove(connection)
         self._available_connections.append(connection)
+        await self._event_dispatcher.dispatch_async(
+            AsyncAfterConnectionReleasedEvent(connection)
+        )
 
     async def disconnect(self, inuse_connections: bool = True):
         """
@@ -1147,6 +1181,29 @@ class ConnectionPool:
         for conn in self._in_use_connections:
             conn.retry = retry
 
+    async def re_auth_callback(self, token: TokenInterface):
+        async with self._lock:
+            for conn in self._available_connections:
+                await conn.retry.call_with_retry(
+                    lambda: conn.send_command(
+                        "AUTH", token.try_get("oid"), token.get_value()
+                    ),
+                    lambda error: self._mock(error),
+                )
+                await conn.retry.call_with_retry(
+                    lambda: conn.read_response(), lambda error: self._mock(error)
+                )
+            for conn in self._in_use_connections:
+                conn.set_re_auth_token(token)
+
+    async def _mock(self, error: RedisError):
+        """
+        Dummy functions, needs to be passed as error callback to retry object.
+        :param error:
+        :return:
+        """
+        pass
+
 
 class BlockingConnectionPool(ConnectionPool):
     """

redis-5.3.0b3/redis/auth/err.py

@@ -0,0 +1,31 @@
+from typing import Iterable
+
+
+class RequestTokenErr(Exception):
+    """
+    Represents an exception during token request.
+    """
+
+    def __init__(self, *args):
+        super().__init__(*args)
+
+
+class InvalidTokenSchemaErr(Exception):
+    """
+    Represents an exception related to invalid token schema.
+    """
+
+    def __init__(self, missing_fields: Iterable[str] = []):
+        super().__init__(
+            "Unexpected token schema. Following fields are missing: "
+            + ", ".join(missing_fields)
+        )
+
+
+class TokenRenewalErr(Exception):
+    """
+    Represents an exception during token renewal process.
+    """
+
+    def __init__(self, *args):
+        super().__init__(*args)

redis-5.3.0b3/redis/auth/idp.py

@@ -0,0 +1,28 @@
+from abc import ABC, abstractmethod
+
+from redis.auth.token import TokenInterface
+
+"""
+This interface is the facade of an identity provider
+"""
+
+
+class IdentityProviderInterface(ABC):
+    """
+    Receive a token from the identity provider.
+    Receiving a token only works when being authenticated.
+    """
+
+    @abstractmethod
+    def request_token(self, force_refresh=False) -> TokenInterface:
+        pass
+
+
+class IdentityProviderConfigInterface(ABC):
+    """
+    Configuration class that provides a configured identity provider.
+    """
+
+    @abstractmethod
+    def get_provider(self) -> IdentityProviderInterface:
+        pass

redis-5.3.0b3/redis/auth/token.py

@@ -0,0 +1,126 @@
+from abc import ABC, abstractmethod
+from datetime import datetime, timezone
+
+import jwt
+from redis.auth.err import InvalidTokenSchemaErr
+
+
+class TokenInterface(ABC):
+    @abstractmethod
+    def is_expired(self) -> bool:
+        pass
+
+    @abstractmethod
+    def ttl(self) -> float:
+        pass
+
+    @abstractmethod
+    def try_get(self, key: str) -> str:
+        pass
+
+    @abstractmethod
+    def get_value(self) -> str:
+        pass
+
+    @abstractmethod
+    def get_expires_at_ms(self) -> float:
+        pass
+
+    @abstractmethod
+    def get_received_at_ms(self) -> float:
+        pass
+
+
+class TokenResponse:
+    def __init__(self, token: TokenInterface):
+        self._token = token
+
+    def get_token(self) -> TokenInterface:
+        return self._token
+
+    def get_ttl_ms(self) -> float:
+        return self._token.get_expires_at_ms() - self._token.get_received_at_ms()
+
+
+class SimpleToken(TokenInterface):
+    def __init__(
+        self, value: str, expires_at_ms: float, received_at_ms: float, claims: dict
+    ) -> None:
+        self.value = value
+        self.expires_at = expires_at_ms
+        self.received_at = received_at_ms
+        self.claims = claims
+
+    def ttl(self) -> float:
+        if self.expires_at == -1:
+            return -1
+
+        return self.expires_at - (datetime.now(timezone.utc).timestamp() * 1000)
+
+    def is_expired(self) -> bool:
+        if self.expires_at == -1:
+            return False
+
+        return self.ttl() <= 0
+
+    def try_get(self, key: str) -> str:
+        return self.claims.get(key)
+
+    def get_value(self) -> str:
+        return self.value
+
+    def get_expires_at_ms(self) -> float:
+        return self.expires_at
+
+    def get_received_at_ms(self) -> float:
+        return self.received_at
+
+
+class JWToken(TokenInterface):
+
+    REQUIRED_FIELDS = {"exp"}
+
+    def __init__(self, token: str):
+        self._value = token
+        self._decoded = jwt.decode(
+            self._value,
+            options={"verify_signature": False},
+            algorithms=[jwt.get_unverified_header(self._value).get("alg")],
+        )
+        self._validate_token()
+
+    def is_expired(self) -> bool:
+        exp = self._decoded["exp"]
+        if exp == -1:
+            return False
+
+        return (
+            self._decoded["exp"] * 1000 <= datetime.now(timezone.utc).timestamp() * 1000
+        )
+
+    def ttl(self) -> float:
+        exp = self._decoded["exp"]
+        if exp == -1:
+            return -1
+
+        return (
+            self._decoded["exp"] * 1000 - datetime.now(timezone.utc).timestamp() * 1000
+        )
+
+    def try_get(self, key: str) -> str:
+        return self._decoded.get(key)
+
+    def get_value(self) -> str:
+        return self._value
+
+    def get_expires_at_ms(self) -> float:
+        return float(self._decoded["exp"] * 1000)
+
+    def get_received_at_ms(self) -> float:
+        return datetime.now(timezone.utc).timestamp() * 1000
+
+    def _validate_token(self):
+        actual_fields = {x for x in self._decoded.keys()}
+
+        if len(self.REQUIRED_FIELDS - actual_fields) != 0:
+            raise InvalidTokenSchemaErr(self.REQUIRED_FIELDS - actual_fields)