coredis 4.23.1__py3-none-any.whl → 5.0.0__py3-none-any.whl

coredis/client/__init__.py

@@ -2,6 +2,5 @@ from __future__ import annotations
 
 from .basic import Client, Redis
 from .cluster import RedisCluster
-from .keydb import KeyDB, KeyDBCluster
 
-__all__ = ["Client", "Redis", "RedisCluster", "KeyDB", "KeyDBCluster"]
+__all__ = ["Client", "Redis", "RedisCluster"]

coredis/client/basic.py

@@ -4,6 +4,7 @@ import asyncio
 import contextlib
 import contextvars
 import functools
+import random
 import warnings
 from collections import defaultdict
 from ssl import SSLContext
@@ -14,7 +15,8 @@ from packaging import version
 from packaging.version import InvalidVersion, Version
 
 from coredis._utils import EncodingInsensitiveDict, nativestr
-from coredis.cache import AbstractCache, SupportsClientTracking
+from coredis.cache import AbstractCache
+from coredis.commands import CommandRequest
 from coredis.commands._key_spec import KeySpec
 from coredis.commands.constants import CommandFlag, CommandName
 from coredis.commands.core import CoreCommands
@@ -31,6 +33,7 @@ from coredis.connection import (
 )
 from coredis.credentials import AbstractCredentialProvider
 from coredis.exceptions import (
+    AuthenticationError,
     ConnectionError,
     PersistenceError,
     RedisError,
@@ -39,7 +42,7 @@ from coredis.exceptions import (
     UnknownCommandError,
     WatchError,
 )
-from coredis.globals import COMMAND_FLAGS, READONLY_COMMANDS
+from coredis.globals import CACHEABLE_COMMANDS, COMMAND_FLAGS, READONLY_COMMANDS
 from coredis.modules import ModuleMixin
 from coredis.pool import ConnectionPool
 from coredis.response._callbacks import (
@@ -55,6 +58,7 @@ from coredis.typing import (
     AsyncIterator,
     Callable,
     Coroutine,
+    ExecutionParameters,
     Generator,
     Generic,
     Iterator,
@@ -63,9 +67,14 @@ from coredis.typing import (
     Mapping,
     Parameters,
     ParamSpec,
+    RedisCommandP,
+    RedisValueT,
     ResponseType,
     StringT,
+    T_co,
+    TypeAdapter,
     TypeVar,
+    Unpack,
     ValueT,
 )
 
@@ -75,7 +84,6 @@ R = TypeVar("R")
 if TYPE_CHECKING:
     import coredis.pipeline
 
-
 ClientT = TypeVar("ClientT", bound="Client[Any]")
 RedisT = TypeVar("RedisT", bound="Redis[Any]")
 
@@ -93,6 +101,7 @@ class Client(
     protocol_version: Literal[2, 3]
     server_version: Version | None
     callback_storage: dict[type[ResponseCallback[Any, Any, Any]], dict[str, Any]]
+    type_adapter: TypeAdapter
 
     def __init__(
         self,
@@ -126,6 +135,7 @@ class Client(
         retry_policy: RetryPolicy = NoRetryPolicy(),
         noevict: bool = False,
         notouch: bool = False,
+        type_adapter: TypeAdapter | None = None,
         **kwargs: Any,
     ):
         if not connection_pool:
@@ -198,6 +208,30 @@ class Client(
         self.retry_policy = retry_policy
         self._module_info: dict[str, version.Version] | None = None
         self.callback_storage = defaultdict(dict)
+        self.type_adapter = type_adapter or TypeAdapter()
+
+    def create_request(
+        self,
+        name: bytes,
+        *arguments: ValueT,
+        callback: Callable[..., T_co],
+        execution_parameters: ExecutionParameters | None = None,
+    ) -> CommandRequest[T_co]:
+        """
+        Factory method to create a command request awaitable.
+        Subclasses of :class:`coredis.client.Client` can override this method
+        if custom behavior is required. See :class:`~coredis.commands.CommandRequest`
+        for details.
+
+        :param name: The name of the command
+        :param arguments: all arguments sent to the command
+        :param callback: a callback that takes the RESP response and converts it
+         into a shape to be returned
+        :return: An instance of a command request bound to this client.
+        """
+        return CommandRequest(
+            self, name, *arguments, callback=callback, execution_parameters=execution_parameters
+        )
 
     @property
     def noreply(self) -> bool:
@@ -220,9 +254,7 @@ class Client(
             return False
         return True
 
-    async def get_server_module_version(self, module: str) -> version.Version | None:
-        if self._module_info is None:
-            await self._populate_module_versions()
+    def get_server_module_version(self, module: str) -> version.Version | None:
         return (self._module_info or {}).get(module)
 
     def _ensure_server_version(self, version: str | None) -> None:
@@ -246,7 +278,7 @@ class Client(
                 self.server_version = None
 
     async def _ensure_wait(
-        self, command: bytes, connection: BaseConnection
+        self, command: RedisCommandP, connection: BaseConnection
     ) -> asyncio.Future[None]:
         maybe_wait: asyncio.Future[None] = asyncio.get_running_loop().create_future()
         wait = self._waitcontext.get()
@@ -257,7 +289,7 @@ class Client(
                 if exc:
                     maybe_wait.set_exception(exc)
                 elif not cast(int, response.result()) >= wait[0]:
-                    maybe_wait.set_exception(ReplicationError(command, wait[0], wait[1]))
+                    maybe_wait.set_exception(ReplicationError(command.name, wait[0], wait[1]))
                 else:
                     maybe_wait.set_result(None)
 
@@ -268,7 +300,7 @@ class Client(
         return maybe_wait
 
     async def _ensure_persistence(
-        self, command: bytes, connection: BaseConnection
+        self, command: RedisCommandP, connection: BaseConnection
     ) -> asyncio.Future[None]:
         maybe_wait: asyncio.Future[None] = asyncio.get_running_loop().create_future()
         waitaof = self._waitaof_context.get()
@@ -283,7 +315,7 @@ class Client(
                 else:
                     res = cast(tuple[int, int], response.result())
                     if not (res[0] >= waitaof[0] and res[1] >= waitaof[1]):
-                        maybe_wait.set_exception(PersistenceError(command, *waitaof))
+                        maybe_wait.set_exception(PersistenceError(command.name, *waitaof))
                     else:
                         maybe_wait.set_result(None)
 
@@ -294,7 +326,7 @@ class Client(
         return maybe_wait
 
     async def _populate_module_versions(self) -> None:
-        if self.noreply:
+        if self.noreply or getattr(self, "_module_info", None) is not None:
             return
         try:
             modules = await self.module_list()
@@ -307,11 +339,12 @@ class Client(
                 ver, minor = divmod(ver, 100)
                 ver, major = divmod(ver, 100)
                 self._module_info[name] = version.Version(f"{major}.{minor}.{patch}")
-        except UnknownCommandError:
+        except (UnknownCommandError, AuthenticationError):
             self._module_info = {}
 
     async def initialize(self: ClientT) -> ClientT:
         await self.connection_pool.initialize()
+        await self._populate_module_versions()
         return self
 
     def __await__(self: ClientT) -> Generator[Any, None, ClientT]:
@@ -422,7 +455,7 @@ class Client(
             for item in data:
                 yield item
 
-    def register_script(self, script: ValueT) -> Script[AnyStr]:
+    def register_script(self, script: RedisValueT) -> Script[AnyStr]:
         """
         Registers a Lua :paramref:`script`
 
@@ -538,8 +571,8 @@ class Client(
         finally:
             self._waitaof_context.set(None)
 
-    def should_quick_release(self, command: bytes) -> bool:
-        return CommandFlag.BLOCKING not in COMMAND_FLAGS[command]
+    def should_quick_release(self, command: RedisCommandP) -> bool:
+        return CommandFlag.BLOCKING not in COMMAND_FLAGS[command.name]
 
 
 class Redis(Client[AnyStr]):
@@ -580,6 +613,7 @@ class Redis(Client[AnyStr]):
         noevict: bool = ...,
         notouch: bool = ...,
         retry_policy: RetryPolicy = ...,
+        type_adapter: TypeAdapter | None = ...,
         **kwargs: Any,
     ) -> None: ...
 
@@ -618,6 +652,7 @@ class Redis(Client[AnyStr]):
         noevict: bool = ...,
         notouch: bool = ...,
         retry_policy: RetryPolicy = ...,
+        type_adapter: TypeAdapter | None = ...,
         **kwargs: Any,
     ) -> None: ...
 
@@ -655,6 +690,7 @@ class Redis(Client[AnyStr]):
         noevict: bool = False,
         notouch: bool = False,
         retry_policy: RetryPolicy = ConstantRetryPolicy((ConnectionError, TimeoutError), 2, 0.01),
+        type_adapter: TypeAdapter | None = None,
         **kwargs: Any,
     ) -> None:
         """
@@ -775,6 +811,8 @@ class Redis(Client[AnyStr]):
         :param notouch: Ensures that commands sent by the client will not alter the LRU/LFU of
          the keys they access.
         :param retry_policy: The retry policy to use when interacting with the redis server
+        :param type_adapter: The adapter to use for serializing / deserializing customs types
+         when interacting with redis commands.
 
         """
         super().__init__(
@@ -796,7 +834,7 @@ class Redis(Client[AnyStr]):
             ssl_keyfile=ssl_keyfile,
             ssl_certfile=ssl_certfile,
             ssl_cert_reqs=ssl_cert_reqs,
-            ssh_check_hostname=ssl_check_hostname,
+            ssl_check_hostname=ssl_check_hostname,
             ssl_ca_certs=ssl_ca_certs,
             max_connections=max_connections,
             max_idle_time=max_idle_time,
@@ -808,6 +846,7 @@ class Redis(Client[AnyStr]):
             noevict=noevict,
             notouch=notouch,
             retry_policy=retry_policy,
+            type_adapter=type_adapter,
             **kwargs,
         )
         self.cache = cache
@@ -867,6 +906,7 @@ class Redis(Client[AnyStr]):
         noevict: bool = False,
         notouch: bool = False,
         retry_policy: RetryPolicy = ConstantRetryPolicy((ConnectionError, TimeoutError), 2, 0.01),
+        type_adapter: TypeAdapter | None = None,
         cache: AbstractCache | None = None,
         **kwargs: Any,
     ) -> RedisT:
@@ -892,6 +932,7 @@ class Redis(Client[AnyStr]):
                 verify_version=verify_version,
                 noreply=noreply,
                 retry_policy=retry_policy,
+                type_adapter=type_adapter,
                 cache=cache,
                 connection_pool=ConnectionPool.from_url(
                     url,
@@ -911,6 +952,7 @@ class Redis(Client[AnyStr]):
                 verify_version=verify_version,
                 noreply=noreply,
                 retry_policy=retry_policy,
+                type_adapter=type_adapter,
                 cache=cache,
                 connection_pool=ConnectionPool.from_url(
                     url,
@@ -933,66 +975,97 @@ class Redis(Client[AnyStr]):
 
     async def execute_command(
         self,
-        command: bytes,
-        *args: ValueT,
+        command: RedisCommandP,
         callback: Callable[..., R] = NoopCallback(),
-        **options: ValueT | None,
+        **options: Unpack[ExecutionParameters],
     ) -> R:
         """
         Executes a command with configured retries and returns
         the parsed response
         """
         return await self.retry_policy.call_with_retries(
-            lambda: self._execute_command(command, *args, callback=callback, **options),
+            lambda: self._execute_command(command, callback=callback, **options),
             before_hook=self.initialize,
         )
 
     async def _execute_command(
         self,
-        command: bytes,
-        *args: ValueT,
+        command: RedisCommandP,
         callback: Callable[..., R] = NoopCallback(),
-        **options: ValueT | None,
+        **options: Unpack[ExecutionParameters],
     ) -> R:
         pool = self.connection_pool
         quick_release = self.should_quick_release(command)
         connection = await pool.get_connection(
-            command,
-            *args,
+            command.name,
+            *command.arguments,
             acquire=not quick_release or self.requires_wait or self.requires_waitaof,
         )
-        if (
-            self.cache
-            and isinstance(self.cache, SupportsClientTracking)
-            and connection.tracking_client_id != self.cache.get_client_id(connection)
-        ):
-            self.cache.reset()
-            await connection.update_tracking_client(True, self.cache.get_client_id(connection))
         try:
-            if self.cache and command not in READONLY_COMMANDS:
-                self.cache.invalidate(*KeySpec.extract_keys(command, *args))
-            request = await connection.create_request(
-                command,
-                *args,
-                noreply=self.noreply,
-                decode=options.get("decode", self._decodecontext.get()),
-                encoding=self._encodingcontext.get(),
-            )
-            maybe_wait = [
-                await self._ensure_wait(command, connection),
-                await self._ensure_persistence(command, connection),
-            ]
-            reply = await request
-            await asyncio.gather(*maybe_wait)
-            if self.noreply:
-                return None  # type: ignore
-            if isinstance(callback, AsyncPreProcessingCallback):
-                await callback.pre_process(self, reply, version=self.protocol_version, **options)
-            return callback(
-                reply,
-                version=self.protocol_version,
-                **options,
+            keys = KeySpec.extract_keys(command.name, *command.arguments)
+            cacheable = (
+                command.name in CACHEABLE_COMMANDS
+                and len(keys) == 1
+                and not self.noreply
+                and self._decodecontext.get() is None
             )
+            cached_reply = None
+            cache_hit = False
+            use_cached = False
+            reply = None
+            if self.cache:
+                if connection.tracking_client_id != self.cache.get_client_id(connection):
+                    self.cache.reset()
+                    await connection.update_tracking_client(
+                        True, self.cache.get_client_id(connection)
+                    )
+                if command.name not in READONLY_COMMANDS:
+                    self.cache.invalidate(*keys)
+                elif cacheable:
+                    try:
+                        cached_reply = cast(
+                            R,
+                            self.cache.get(
+                                command.name,
+                                keys[0],
+                                *command.arguments,
+                            ),
+                        )
+                        use_cached = random.random() * 100.0 < min(100.0, self.cache.confidence)
+                        cache_hit = True
+                    except KeyError:
+                        pass
+            if not (use_cached and cached_reply):
+                request = await connection.create_request(
+                    command.name,
+                    *command.arguments,
+                    noreply=self.noreply,
+                    decode=options.get("decode", self._decodecontext.get()),
+                    encoding=self._encodingcontext.get(),
+                )
+                maybe_wait = [
+                    await self._ensure_wait(command, connection),
+                    await self._ensure_persistence(command, connection),
+                ]
+                reply = await request
+                await asyncio.gather(*maybe_wait)
+                if self.noreply:
+                    return None  # type: ignore
+                if isinstance(callback, AsyncPreProcessingCallback):
+                    await callback.pre_process(self, reply)
+            if self.cache and cacheable:
+                if cache_hit and not use_cached:
+                    self.cache.feedback(
+                        command.name, keys[0], *command.arguments, match=cached_reply == reply
+                    )
+                if not cache_hit:
+                    self.cache.put(
+                        command.name,
+                        keys[0],
+                        *command.arguments,
+                        value=reply,
+                    )
+            return callback(cached_reply if cache_hit else reply, version=self.protocol_version)
         except RedisError:
             connection.disconnect()
             raise
@@ -1121,7 +1194,7 @@ class Redis(Client[AnyStr]):
         """
         from coredis.pipeline import Pipeline
 
-        return Pipeline[AnyStr].proxy(self, transaction, watches, timeout)
+        return Pipeline[AnyStr](self, transaction, watches, timeout)
 
     async def transaction(
         self,