from __future__ import annotations
from contextlib import asynccontextmanager
from typing import Any, ClassVar, cast
from anyio import AsyncContextManagerMixin
from deprecated.sphinx import versionchanged
from coredis._telemetry import get_telemetry_provider
from coredis._utils import nativestr
from coredis.client import Client, RedisCluster
from coredis.cluster._node import ClusterNodeLocation
from coredis.commands import CommandRequest, CommandResponseT
from coredis.commands._routing import ExplicitSlotStrategy, RandomStrategy
from coredis.commands.constants import CommandName
from coredis.commands.script import Script
from coredis.connection._base import BaseConnection
from coredis.connection._request import Request
from coredis.exceptions import (
AskError,
ClusterCrossSlotError,
ClusterDownError,
ClusterTransactionError,
ConnectionError,
ExecAbortError,
MovedError,
RedisClusterError,
RedisError,
ResponseError,
TryAgainError,
WatchError,
)
from coredis.pool import ClusterConnectionPool
from coredis.response._callbacks import (
AnyStrCallback,
BoolsCallback,
NoopCallback,
SimpleStringCallback,
)
from coredis.retry import ConstantRetryPolicy, retryable
from coredis.typing import (
AnyStr,
AsyncGenerator,
Awaitable,
Callable,
ExecutionParameters,
Generator,
Iterable,
Key,
KeyT,
ParamSpec,
RedisValueT,
ResponseType,
Self,
T_co,
TypeAdapter,
TypeVar,
ValueT,
)
P = ParamSpec("P")
R = TypeVar("R")
T = TypeVar("T")
ERRORS_ALLOW_RETRY = (
MovedError,
AskError,
TryAgainError,
)
UNWATCH_COMMANDS = {CommandName.DISCARD, CommandName.EXEC, CommandName.UNWATCH}
class PipelineResult(Awaitable[T]):
__slots__ = ("_result",)
def __init__(self, result: T) -> None:
self._result = result
@property
def value(self) -> T:
if isinstance(self._result, Exception):
raise self._result
else:
return self._result
def __await__(self) -> Generator[Any, None, T]:
async def _coro() -> T:
return self.value
return _coro().__await__()
class PipelineCommandRequest(CommandRequest[CommandResponseT]):
"""
Command request returned by a pipeline command
"""
def __await__(self) -> Generator[None, None, CommandResponseT]:
if hasattr(self, "_response"):
return self._response.__await__()
raise RuntimeError("You can't await a pipeline command before it completes executing")
class ClusterPipelineCommandRequest(CommandRequest[CommandResponseT]):
"""
Command request for cluster pipelines, tracks position and result for cluster routing.
"""
def __init__(
self,
name: bytes,
*arguments: ValueT | Key,
callback: Callable[..., CommandResponseT],
execution_parameters: ExecutionParameters,
type_adapter: TypeAdapter,
) -> None:
super().__init__(
name,
*arguments,
callback=callback,
execution_parameters=execution_parameters,
type_adapter=type_adapter,
)
self.position: int = 0
self.result: Any = None
def __await__(self) -> Generator[None, None, CommandResponseT]:
if hasattr(self, "_response"):
return self._response.__await__()
raise RuntimeError("You can't await a pipeline command before it completes executing")
class NodeCommands(AsyncContextManagerMixin):
"""
Helper for grouping and executing commands on a single cluster node, handling transactions if needed.
"""
connection: BaseConnection
def __init__(
self,
client: RedisCluster[AnyStr],
node: ClusterNodeLocation,
connection: BaseConnection | None = None,
in_transaction: bool = False,
timeout: float | None = None,
raise_on_error: bool = True,
):
self.client: RedisCluster[Any] = client
self.node = node
self._connection = connection
self.commands: list[ClusterPipelineCommandRequest[Any]] = []
self.in_transaction = in_transaction
self.timeout = timeout
self._raise_on_error = raise_on_error
self.multi_cmd: Request | None = None
self.exec_cmd: Request | None = None
self.request_batch: Awaitable[list[ResponseType | BaseException | None]] | None = None
def extend(self, c: list[ClusterPipelineCommandRequest[Any]]) -> None:
self.commands.extend(c)
def append(self, c: ClusterPipelineCommandRequest[Any]) -> None:
self.commands.append(c)
@asynccontextmanager
async def __asynccontextmanager__(self) -> AsyncGenerator[None]:
if not self._connection:
async with self.client.connection_pool.acquire(node=self.node) as self.connection:
yield
else:
self.connection = self._connection
yield
def write(self) -> None:
connection = self.connection
commands = self.commands
# Reset results for all commands before writing.
for c in commands:
c.result = None
# Batch all commands into a single request for efficiency.
try:
if self.in_transaction:
self.multi_cmd = connection.create_request(CommandName.MULTI, timeout=self.timeout)
self.request_batch = connection.create_request_batch(
commands,
timeout=self.timeout,
)
if self.in_transaction:
self.exec_cmd = connection.create_request(CommandName.EXEC, timeout=self.timeout)
except (ConnectionError, TimeoutError) as e:
for c in commands:
c.result = e
async def read(self) -> None:
success = True
multi_result = None
if self.multi_cmd:
multi_result = await self.multi_cmd
success = multi_result in {b"OK", "OK"}
if self.request_batch:
responses = await self.request_batch
for command, response in zip(self.commands, responses):
command.result = response
command._response = PipelineResult(response)
if isinstance(response, (ConnectionError, TimeoutError, RedisError)):
success = False
if self.in_transaction and self.exec_cmd:
if success:
res = await self.exec_cmd
if res:
transaction_result = cast(list[ResponseType], res)
else:
raise WatchError("Watched variable changed.")
for idx, c in enumerate(
[
_c
for _c in sorted(self.commands, key=lambda x: x.position)
if _c.name not in {CommandName.MULTI, CommandName.EXEC}
]
):
c.result = c.callback(
transaction_result[idx],
)
c._response = PipelineResult(c.result)
elif isinstance(multi_result, BaseException):
raise multi_result
[docs]
@versionchanged(
version="6.0.0",
reason="Pipelines are no longer awaitable. They support the async context manager protocol and must always be used as such",
)
class Pipeline(Client[AnyStr]):
"""
Pipeline for batching multiple commands to a Redis server.
Supports transactions and command stacking.
All commands executed within a pipeline are wrapped with MULTI and EXEC
calls when :paramref:`transaction` is ``True``.
Any command raising an exception does **not** halt the execution of
subsequent commands in the pipeline, however the first exception encountered
will be raised when exiting the pipeline if :paramref:`raise_on_error` is ``True``.
If not the exception is caught and will be returned when awaiting the command that failed.
"""
QUEUED_RESPONSES: ClassVar[set[bytes | str]] = {b"QUEUED", "QUEUED"}
def __init__(
self,
client: Client[AnyStr],
transaction: bool | None,
raise_on_error: bool = True,
timeout: float | None = None,
) -> None:
"""
:param transaction: Whether to wrap the commands in the pipeline in a ``MULTI``, ``EXEC``
:param raise_on_error: Whether to raise the first error encounterd in the pipeline after
executing it
:param timeout: Time in seconds to wait for the pipeline results to return
"""
self.client: Client[AnyStr] = client
self._connection: BaseConnection | None = None
self._transaction = transaction
self._raise_on_error = raise_on_error
self._watching: bool = False
self.command_stack: list[PipelineCommandRequest[Any]] = []
self.cache = None
self.scripts: set[Script[AnyStr]] = set()
self.timeout = timeout
self.type_adapter = client.type_adapter
self._results: tuple[Any] | None = None
@asynccontextmanager
async def __asynccontextmanager__(self) -> AsyncGenerator[Self]:
try:
yield self
await self._execute()
finally:
self._unwatch()
if self._connection:
self.client.connection_pool.release(self._connection)
self._connection = None
def __repr__(self) -> str:
return f"{type(self).__name__}<{self._connection!r}>"
def create_request(
self,
name: bytes,
*arguments: ValueT | Key,
callback: Callable[..., T_co],
execution_parameters: ExecutionParameters | None = None,
) -> CommandRequest[T_co]:
"""
:meta private:
"""
command = PipelineCommandRequest(
name,
*arguments,
callback=callback,
execution_parameters=execution_parameters or {},
type_adapter=self.type_adapter,
)
self.command_stack.append(command)
return command
[docs]
@asynccontextmanager
async def watch(self, *keys: KeyT) -> AsyncGenerator[None]:
"""
The given keys will be watched for changes within this context and the
commands stacked within the context will be automatically executed when
the context exits. On abort, watches are cleared without executing.
"""
if self.command_stack:
raise WatchError("Unable to add a watch after pipeline commands have been added")
if not self._connection:
self._connection = await self.client.connection_pool.get_connection()
await self._immediate_execute_command(
self.client.create_request(
CommandName.WATCH,
*[Key(key) for key in keys],
callback=SimpleStringCallback(),
)
)
self._watching = True
try:
yield
await self._execute()
finally:
self._unwatch()
@property
def results(self) -> tuple[Any, ...] | None:
"""
The results of the pipeline execution which can be accessed
after the pipeline has completed.
"""
if self.command_stack:
raise RuntimeError("Pipeline results are not available before it completes execution")
return self._results
def execute_command(
self,
command: CommandRequest[R],
) -> Awaitable[R]:
raise NotImplementedError
def _unwatch(self) -> None:
if not self._watching:
return
try:
if self._connection and self._connection.usable:
# The response is not awaited to make this function safe for finally blocks.
# noreply is intentionally not used since this doesn't work for dragonfly :(
self._connection.create_request(CommandName.UNWATCH, decode=False)
finally:
self._watching = False
async def _clear(self) -> None:
"""
Clear the pipeline and reset state.
"""
self.command_stack.clear()
self.scripts.clear()
self._unwatch()
async def _immediate_execute_command(
self,
command: CommandRequest[R],
) -> R:
"""
Executes a command immediately, but don't auto-retry on a
ConnectionError if we're already WATCHing a variable. Used when
issuing WATCH or subsequent commands retrieving their values but before
MULTI is called.
:meta private:
"""
assert self._connection
request = self._connection.create_request(
command.name,
*command.serialized_arguments,
decode=command.execution_parameters.get("decode"),
)
return command.callback(await request)
async def _execute_transaction(
self,
connection: BaseConnection,
commands: list[PipelineCommandRequest[Any]],
) -> None:
multi_request = connection.create_request(CommandName.MULTI, timeout=self.timeout)
queued_batch = connection.create_request_batch(commands, timeout=self.timeout)
exec_request = connection.create_request(CommandName.EXEC, timeout=self.timeout)
errors: list[tuple[int, RedisError | TimeoutError | None]] = []
try:
await multi_request
except (RedisError, TimeoutError) as e:
errors.append((0, e))
for i, (cmd, queued_response) in enumerate(zip(commands, await queued_batch)):
if queued_response not in self.QUEUED_RESPONSES:
if isinstance(queued_response, (RedisError, TimeoutError)):
self._annotate_exception(
queued_response, i + 1, cmd.name, cmd.serialized_arguments
)
errors.append((i + 1, queued_response))
elif isinstance(queued_response, BaseException):
raise queued_response
else:
raise ResponseError(
f"Abnormal response in pipeline for command {cmd.name!r}: {queued_response!r}"
)
try:
response = cast(list[ResponseType] | None, await exec_request)
except Exception as exec_err:
if errors and errors[0][1]:
raise errors[0][1] from exec_err
raise
if response is None:
raise WatchError("Watched variable changed.")
for i, batch_err in errors:
response.insert(i, cast(ResponseType, batch_err))
if len(response) != len(commands):
raise ResponseError("Wrong number of response items from pipeline execution")
data: list[Any] = []
for r, cmd in zip(response, commands):
if not isinstance(r, Exception):
r = cmd.callback(r, **cmd.execution_parameters)
cmd._response = PipelineResult(r)
data.append(r)
self._results = tuple(data)
if self._raise_on_error:
self._raise_first_error(commands, response)
async def _execute_pipeline(
self, connection: BaseConnection, commands: list[PipelineCommandRequest[Any]]
) -> None:
request_batch = connection.create_request_batch(commands, timeout=self.timeout)
results: list[Any] = []
for cmd, response in zip(commands, await request_batch):
try:
if isinstance(response, BaseException):
raise response
resp = cmd.callback(
response,
**cmd.execution_parameters,
)
cmd._response = PipelineResult(resp)
results.append(resp)
except (ResponseError, TimeoutError) as re:
cmd._response = PipelineResult(re)
results.append(re)
self._results = tuple(results)
if self._raise_on_error:
self._raise_first_error(commands, results)
def _raise_first_error(
self, commands: list[PipelineCommandRequest[Any]], response: ResponseType
) -> None:
assert isinstance(response, list)
for i, r in enumerate(response):
if isinstance(r, (RedisError, TimeoutError)):
self._annotate_exception(
r, i + 1, commands[i].name, commands[i].serialized_arguments
)
raise r
def _annotate_exception(
self,
exception: RedisError | TimeoutError | None,
number: int,
command: bytes,
args: Iterable[RedisValueT],
) -> None:
if exception:
cmd = command.decode("latin-1")
args = " ".join(map(str, args))
msg = (
f"Command # {number} ({cmd} {args}) of pipeline caused error: {exception.args[0]!s}"
)
exception.args = (msg,) + exception.args[1:]
async def _load_scripts(self) -> None:
# make sure all scripts that are about to be run on this pipeline exist
scripts = list(self.scripts)
shas = [s.sha for s in scripts]
exists = await self._immediate_execute_command(
self.client.create_request(CommandName.SCRIPT_EXISTS, *shas, callback=BoolsCallback())
)
if not all(exists):
for s, exist in zip(scripts, exists):
if not exist:
s.sha = await self._immediate_execute_command(
self.client.create_request(
CommandName.SCRIPT_LOAD,
s.script,
callback=AnyStrCallback[AnyStr](),
)
)
async def _execute(self) -> None:
"""
Execute all queued commands in the pipeline.
"""
if not self.command_stack:
return None
if not self._connection:
self._connection = await self.client.connection_pool.get_connection()
with get_telemetry_provider().start_span(
self.command_stack,
self._connection,
self.client.connection_pool,
name="MULTI" if self._transaction else "PIPELINE",
):
if self.scripts:
await self._load_scripts()
if self._transaction or self._watching:
exec = self._execute_transaction
else:
exec = self._execute_pipeline
try:
return await exec(self._connection, self.command_stack)
except (ConnectionError, TimeoutError) as e:
if self._watching:
raise WatchError(
"A connection error occurred while watching one or more keys"
) from e
raise
finally:
await self._clear()
[docs]
@versionchanged(
version="6.0.0",
reason="Cluster Pipelines are no longer awaitable. They support the async context manager protocol and must always be used as such",
)
class ClusterPipeline(Client[AnyStr]):
"""
Pipeline for batching multiple commands to a Redis Cluster
Supports transactions only when all keys map to the same shard, and therefore
:paramref:`transactions` is set to ``False`` by default due to the limited scope.
Any command raising an exception does **not** halt the execution of
subsequent commands in the pipeline, however the first exception encountered
will be raised when exiting the pipeline if :paramref:`raise_on_error` is ``True``.
If not the exception is caught and will be returned when awaiting the command that failed.
"""
client: RedisCluster[AnyStr]
connection_pool: ClusterConnectionPool
command_stack: list[ClusterPipelineCommandRequest[Any]]
def __init__(
self,
client: RedisCluster[AnyStr],
raise_on_error: bool = True,
transaction: bool = False,
timeout: float | None = None,
) -> None:
"""
:param transaction: Whether to wrap the commands in the pipeline in a ``MULTI``, ``EXEC``
:param raise_on_error: Whether to raise the first error encounterd in the pipeline after
executing it
:param timeout: Time in seconds to wait for the pipeline results to return
"""
self.command_stack = []
self.client = client
self.connection_pool = client.connection_pool
self._raise_on_error = raise_on_error
self._transaction = transaction
self._watched_node: ClusterNodeLocation | None = None
self._watched_connection: BaseConnection | None = None
self._watching: bool = False
self.cache = None
self.scripts: set[Script[AnyStr]] = set()
self.timeout = timeout
self.type_adapter = client.type_adapter
self._results: tuple[Any] | None = None
@asynccontextmanager
async def __asynccontextmanager__(self) -> AsyncGenerator[Self]:
try:
yield self
await self._execute()
finally:
self._unwatch()
def create_request(
self,
name: bytes,
*arguments: ValueT | Key,
callback: Callable[..., T_co],
execution_parameters: ExecutionParameters | None = None,
) -> CommandRequest[T_co]:
"""
:meta private:
"""
command = ClusterPipelineCommandRequest(
name,
*arguments,
callback=callback,
execution_parameters=execution_parameters or {},
type_adapter=self.type_adapter,
)
command.position = len(self.command_stack)
self.command_stack.append(command)
return command
[docs]
@asynccontextmanager
async def watch(self, *keys: KeyT) -> AsyncGenerator[None]:
"""
The given keys will be watched for changes within this context and the
commands stacked within the context will be automatically executed when
the context exits. On abort, watches are cleared without executing.
"""
if self.command_stack:
raise WatchError("Unable to add a watch after pipeline commands have been added")
self._watched_node = self.connection_pool.cluster_layout.node_for_request(
self.client.create_request(
b"WATCH", *[Key(key) for key in keys], callback=NoopCallback()
)
)
async with self.connection_pool.acquire(
node=self._watched_node
) as self._watched_connection:
try:
await self._watched_connection.create_request(CommandName.WATCH, *keys)
self._watching = True
yield
await self._execute()
finally:
self._unwatch()
self._watched_connection = None
self._watched_node = None
@property
def results(self) -> tuple[Any, ...] | None:
"""
The results of the pipeline execution which can be accessed
after the pipeline has completed.
"""
if self.command_stack:
raise RuntimeError("Pipeline results are not available before it completes")
return self._results
def execute_command(
self,
command: CommandRequest[R],
) -> Awaitable[R]:
raise NotImplementedError
def _unwatch(self) -> None:
if not self._watching:
return
try:
if self._watched_connection and self._watched_connection.usable:
# The response is not awaited to make this function safe for finally blocks.
# noreply is intentionally not used since this doesn't work for dragonfly :(
self._watched_connection.create_request(CommandName.UNWATCH, decode=False)
finally:
self._watching = False
async def _clear(self) -> None:
"""
Clear the pipeline and reset state.
"""
self.command_stack = []
self.scripts.clear()
self._unwatch()
def _raise_first_error(self) -> None:
for c in self.command_stack:
r = c.result
if isinstance(r, (RedisError, TimeoutError)):
self._annotate_exception(r, c.position + 1, c.name, c.serialized_arguments)
raise r
def _annotate_exception(
self,
exception: RedisError | TimeoutError | None,
number: int,
command: bytes,
args: Iterable[RedisValueT],
) -> None:
if exception:
cmd = command.decode("latin-1")
args = " ".join(str(x) for x in args)
msg = f"Command # {number} ({cmd} {args}) of pipeline caused error: {exception.args[0]}"
exception.args = (msg,) + exception.args[1:]
async def _execute(self) -> None:
"""
Execute all queued commands in the cluster pipeline. Returns a tuple of results.
"""
if not self.command_stack:
return
with get_telemetry_provider().start_span(
self.command_stack,
self.client.connection_pool,
name="MULTI" if self._transaction else "PIPELINE",
):
if self.scripts:
await self._load_scripts()
use_primary = not (
self.client.connection_pool.read_from_replicas
and all(cmd.readonly for cmd in self.command_stack)
)
if self._transaction or self._watching:
execute = self._send_cluster_transaction
else:
execute = self._send_cluster_commands
try:
await execute(self._raise_on_error, use_primary)
finally:
await self._clear()
def _get_slot_for_command(self, command: ClusterPipelineCommandRequest[Any]) -> int | None:
affected_slots = command.affected_slots
if len(affected_slots) == 1:
return affected_slots[0]
match command.routing_strategy:
case ExplicitSlotStrategy():
return command.routing_strategy.slot
case RandomStrategy():
return None
if not affected_slots:
raise RedisClusterError(
f"No way to dispatch {nativestr(command.name)} to Redis Cluster. Missing key"
)
else:
raise ClusterCrossSlotError(command=command.name, keys=command.keys)
@retryable(policy=ConstantRetryPolicy((ClusterDownError,), retries=3, delay=0.1))
async def _send_cluster_transaction(
self, raise_on_error: bool = True, use_primary: bool = False
) -> None:
"""
:meta private:
"""
attempt = sorted(self.command_stack, key=lambda x: x.position)
slots: set[int] = set()
for c in attempt:
if (slot := self._get_slot_for_command(c)) is not None:
slots.add(slot)
if len(slots) > 1:
raise ClusterTransactionError("Multiple slots involved in transaction")
if not slots:
raise ClusterTransactionError("No slots found for transaction")
node = self.connection_pool.cluster_layout.node_for_slot(slots.pop(), use_primary)
if self._watched_node and node != self._watched_node:
raise ClusterTransactionError("Multiple slots involved in transaction")
node_commands = NodeCommands(
self.client,
node,
in_transaction=True,
timeout=self.timeout,
connection=self._watched_connection,
raise_on_error=self._raise_on_error,
)
node_commands.extend(attempt)
async with node_commands:
node_commands.write()
try:
await node_commands.read()
except ExecAbortError:
await node_commands.connection.create_request(CommandName.DISCARD)
# If at least one watched key is modified before EXEC, the transaction aborts and EXEC returns null.
if node_commands.exec_cmd:
exec_result = await node_commands.exec_cmd
if exec_result is None:
raise WatchError("Watched variable changed.")
self._results = tuple(
n.result
for n in node_commands.commands
if n.name not in {CommandName.MULTI, CommandName.EXEC}
)
if raise_on_error:
self._raise_first_error()
@retryable(policy=ConstantRetryPolicy((ClusterDownError,), retries=3, delay=0.1))
async def _send_cluster_commands(
self, raise_on_error: bool = True, use_primary: bool = False
) -> None:
"""
Execute all queued commands in the cluster pipeline, handling redirections
and retries as needed.
:meta private:
"""
attempt: dict[ClusterPipelineCommandRequest[Any], ClusterNodeLocation] = {}
nodes: dict[str, NodeCommands] = {}
for c in sorted(self.command_stack, key=lambda x: x.position):
if (slot := self._get_slot_for_command(c)) is not None:
node = self.connection_pool.cluster_layout.node_for_slot(slot, use_primary)
else:
node = node or self.connection_pool.cluster_layout.random_node(use_primary)
if node.name not in nodes:
nodes[node.name] = NodeCommands(
self.client,
node,
timeout=self.timeout,
raise_on_error=self._raise_on_error,
)
nodes[node.name].append(c)
attempt[c] = node
# Write to all nodes, then read from all nodes in sequence.
for n in nodes.values():
async with n:
n.write()
await n.read()
# Retry MOVED/ASK/connection errors one by one if allowed.
attempt = dict(
sorted(
(c for c in attempt.items() if isinstance(c[0].result, ERRORS_ALLOW_RETRY)),
key=lambda x: x[0].position,
)
)
if attempt:
for c, node in attempt.items():
self.connection_pool.cluster_layout.report_errors(node, c.result)
try:
c.result = await self.client.execute_command(c.raw())
except (RedisError, TimeoutError) as e:
c.result = e
# Flatten results to match the original command order.
response = []
for c in sorted(self.command_stack, key=lambda x: x.position):
r = c.result
if not isinstance(c.result, (RedisError, TimeoutError)):
r = c.callback(c.result)
c._response = PipelineResult(r)
response.append(r)
self._results = tuple(response)
if raise_on_error:
self._raise_first_error()
async def _load_scripts(self) -> None:
shas = [s.sha for s in self.scripts]
exists = await self.client.script_exists(shas)
if not all(exists):
for s, exist in zip(self.scripts, exists):
if not exist:
s.sha = await self.client.script_load(s.script)