Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/redis/asyncio/client.py: 22%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1import asyncio
2import copy
3import inspect
4import math
5import re
6import time
7import warnings
8from typing import (
9 TYPE_CHECKING,
10 Any,
11 AsyncIterator,
12 Awaitable,
13 Callable,
14 Dict,
15 Iterable,
16 List,
17 Literal,
18 Mapping,
19 MutableMapping,
20 Optional,
21 Protocol,
22 Set,
23 Tuple,
24 Type,
25 TypedDict,
26 TypeVar,
27 Union,
28 cast,
29)
31from redis._defaults import (
32 DEFAULT_RETRY_BASE,
33 DEFAULT_RETRY_CAP,
34 DEFAULT_RETRY_COUNT,
35 DEFAULT_SOCKET_CONNECT_TIMEOUT,
36 DEFAULT_SOCKET_READ_SIZE,
37 DEFAULT_SOCKET_TIMEOUT,
38)
39from redis._parsers.helpers import bool_ok, get_response_callbacks
40from redis.asyncio import _himport_exec
41from redis.asyncio.connection import (
42 Connection,
43 ConnectionPool,
44 SSLConnection,
45 UnixDomainSocketConnection,
46)
47from redis.asyncio.lock import Lock
48from redis.asyncio.observability.recorder import (
49 record_error_count,
50 record_operation_duration,
51 record_pubsub_message,
52)
53from redis.asyncio.retry import Retry
54from redis.backoff import ExponentialWithJitterBackoff
55from redis.client import (
56 EMPTY_RESPONSE,
57 NEVER_DECODE,
58 AbstractRedis,
59 CaseInsensitiveDict,
60)
61from redis.commands import (
62 AsyncCoreCommands,
63 AsyncRedisModuleCommands,
64 AsyncSentinelCommands,
65 list_or_args,
66)
67from redis.commands.helpers import parse_pubsub_subscriptions, pubsub_subscription_args
68from redis.credentials import CredentialProvider
69from redis.driver_info import DriverInfo, resolve_driver_info
70from redis.event import (
71 AfterPooledConnectionsInstantiationEvent,
72 AfterPubSubConnectionInstantiationEvent,
73 AfterSingleConnectionInstantiationEvent,
74 ClientType,
75 EventDispatcher,
76)
77from redis.exceptions import (
78 ConnectionError,
79 ExecAbortError,
80 PubSubError,
81 RedisError,
82 ResponseError,
83 WatchError,
84)
85from redis.himport import HImportRegistry, parse_himport_set_args
86from redis.maint_notifications import MaintNotificationsConfig
87from redis.observability.attributes import PubSubDirection
88from redis.typing import (
89 ChannelT,
90 EncodableT,
91 FieldT,
92 KeyT,
93 PubSubHandler,
94 Subscription,
95)
96from redis.utils import (
97 SENTINEL,
98 SSL_AVAILABLE,
99 _set_info_logger,
100 check_protocol_version,
101 deprecated_args,
102 deprecated_function,
103 experimental_method,
104 safe_str,
105 str_if_bytes,
106 truncate_text,
107)
109if TYPE_CHECKING and SSL_AVAILABLE:
110 from ssl import TLSVersion, VerifyFlags, VerifyMode
111else:
112 TLSVersion = None
113 VerifyMode = None
114 VerifyFlags = None
116_KeyT = TypeVar("_KeyT", bound=KeyT)
117_ArgT = TypeVar("_ArgT", KeyT, EncodableT)
118_RedisT = TypeVar("_RedisT", bound="Redis")
119_NormalizeKeysT = TypeVar("_NormalizeKeysT", bound=Mapping[ChannelT, object])
120if TYPE_CHECKING:
121 from redis.asyncio.keyspace_notifications import AsyncKeyspaceNotifications
122 from redis.commands.core import Script
125class ResponseCallbackProtocol(Protocol):
126 def __call__(self, response: Any, **kwargs): ...
129class AsyncResponseCallbackProtocol(Protocol):
130 async def __call__(self, response: Any, **kwargs): ...
133ResponseCallbackT = Union[ResponseCallbackProtocol, AsyncResponseCallbackProtocol]
136class Redis(
137 AbstractRedis, AsyncRedisModuleCommands, AsyncCoreCommands, AsyncSentinelCommands
138):
139 """
140 Implementation of the Redis protocol.
142 This abstract class provides a Python interface to all Redis commands
143 and an implementation of the Redis protocol.
145 Pipelines derive from this, implementing how
146 the commands are sent and received to the Redis server. Based on
147 configuration, an instance will either use a ConnectionPool, or
148 Connection object to talk to redis.
149 """
151 # Type discrimination marker for @overload self-type pattern
152 _is_async_client: Literal[True] = True
154 response_callbacks: MutableMapping[Union[str, bytes], ResponseCallbackT]
156 @classmethod
157 def from_url(
158 cls: Type["Redis"],
159 url: str,
160 single_connection_client: bool = False,
161 auto_close_connection_pool: Optional[bool] = None,
162 **kwargs,
163 ) -> "Redis":
164 """
165 Return a Redis client object configured from the given URL
167 For example::
169 redis://[[username]:[password]]@localhost:6379/0
170 rediss://[[username]:[password]]@localhost:6379/0
171 unix://[username@]/path/to/socket.sock?db=0[&password=password]
173 Three URL schemes are supported:
175 - `redis://` creates a TCP socket connection. See more at:
176 <https://www.iana.org/assignments/uri-schemes/prov/redis>
177 - `rediss://` creates a SSL wrapped TCP socket connection. See more at:
178 <https://www.iana.org/assignments/uri-schemes/prov/rediss>
179 - ``unix://``: creates a Unix Domain Socket connection.
181 The username, password, hostname and path are passed through
182 urllib.parse.unquote in order to replace any percent-encoded values
183 with their corresponding characters. Querystring values are decoded
184 by urllib.parse.parse_qs and are not unquoted again.
186 There are several ways to specify a database number. The first value
187 found will be used:
189 1. A ``db`` querystring option, e.g. redis://localhost?db=0
191 2. If using the redis:// or rediss:// schemes, the path argument
192 of the url, e.g. redis://localhost/0
194 3. A ``db`` keyword argument to this function.
196 If none of these options are specified, the default db=0 is used.
198 All querystring options are cast to their appropriate Python types.
199 Boolean arguments can be specified with string values "True"/"False"
200 or "Yes"/"No". Values that cannot be properly cast cause a
201 ``ValueError`` to be raised. Once parsed, the querystring arguments
202 and keyword arguments are passed to the ``ConnectionPool``'s
203 class initializer. In the case of conflicting arguments, querystring
204 arguments always win.
206 """
207 connection_pool = ConnectionPool.from_url(url, **kwargs)
208 client = cls(
209 connection_pool=connection_pool,
210 single_connection_client=single_connection_client,
211 )
212 if auto_close_connection_pool is not None:
213 warnings.warn(
214 DeprecationWarning(
215 '"auto_close_connection_pool" is deprecated '
216 "since version 5.0.1. "
217 "Please create a ConnectionPool explicitly and "
218 "provide to the Redis() constructor instead."
219 )
220 )
221 else:
222 auto_close_connection_pool = True
223 client.auto_close_connection_pool = auto_close_connection_pool
224 return client
226 @classmethod
227 def from_pool(
228 cls: Type["Redis"],
229 connection_pool: ConnectionPool,
230 ) -> "Redis":
231 """
232 Return a Redis client from the given connection pool.
233 The Redis client will take ownership of the connection pool and
234 close it when the Redis client is closed.
236 Because the client closes (disconnects all connections in) the pool
237 when it is closed or garbage-collected, the pool must not be shared
238 with other clients. Constructing multiple clients from the same pool
239 via ``from_pool`` -- for example one per request across tasks -- is
240 not safe: when one client is closed it will disconnect connections
241 still in use by the others.
243 To share a single pool across clients, construct the pool explicitly
244 and manage its lifecycle instead. Unlike ``from_pool``, the plain
245 ``Redis(connection_pool=pool)`` constructor does not take ownership of
246 the pool and will not close it, so a pool created this way can be
247 safely shared across clients. ``ConnectionPool`` supports the async
248 context manager protocol for this::
250 async with ConnectionPool.from_url(url) as pool:
251 r = Redis(connection_pool=pool)
252 """
253 client = cls(
254 connection_pool=connection_pool,
255 )
256 client.auto_close_connection_pool = True
257 return client
259 @deprecated_args(
260 args_to_warn=["retry_on_timeout"],
261 reason="TimeoutError is included by default.",
262 version="6.0.0",
263 )
264 @deprecated_args(
265 args_to_warn=["lib_name", "lib_version"],
266 reason="Use 'driver_info' parameter instead. "
267 "lib_name and lib_version will be removed in a future version.",
268 )
269 def __init__(
270 self,
271 *,
272 host: str = "localhost",
273 port: int = 6379,
274 db: str | int = 0,
275 password: str | None = None,
276 socket_timeout: float | None = DEFAULT_SOCKET_TIMEOUT,
277 socket_connect_timeout: float | None = DEFAULT_SOCKET_CONNECT_TIMEOUT,
278 socket_read_size: int = DEFAULT_SOCKET_READ_SIZE,
279 socket_keepalive: bool | None = True,
280 socket_keepalive_options: Mapping[int, int | bytes] | object | None = SENTINEL,
281 connection_pool: ConnectionPool | None = None,
282 unix_socket_path: str | None = None,
283 encoding: str = "utf-8",
284 encoding_errors: str = "strict",
285 decode_responses: bool = False,
286 retry_on_timeout: bool = False,
287 retry: Retry = Retry(
288 backoff=ExponentialWithJitterBackoff(
289 base=DEFAULT_RETRY_BASE, cap=DEFAULT_RETRY_CAP
290 ),
291 retries=DEFAULT_RETRY_COUNT,
292 ),
293 retry_on_error: list | None = None,
294 ssl: bool = False,
295 ssl_keyfile: str | None = None,
296 ssl_certfile: str | None = None,
297 ssl_cert_reqs: "str | VerifyMode" = "required",
298 ssl_include_verify_flags: List["VerifyFlags"] | None = None,
299 ssl_exclude_verify_flags: List["VerifyFlags"] | None = None,
300 ssl_ca_certs: str | None = None,
301 ssl_ca_data: str | None = None,
302 ssl_ca_path: str | None = None,
303 ssl_check_hostname: bool = True,
304 ssl_min_version: "TLSVersion | None" = None,
305 ssl_ciphers: str | None = None,
306 ssl_password: str | None = None,
307 max_connections: int | None = None,
308 single_connection_client: bool = False,
309 health_check_interval: int = 0,
310 client_name: str | None = None,
311 lib_name: str | object | None = SENTINEL,
312 lib_version: str | object | None = SENTINEL,
313 driver_info: DriverInfo | object | None = SENTINEL,
314 username: str | None = None,
315 auto_close_connection_pool: bool | None = None,
316 redis_connect_func=None,
317 credential_provider: CredentialProvider | None = None,
318 protocol: int | None = None,
319 legacy_responses: bool = True,
320 event_dispatcher: EventDispatcher | None = None,
321 maint_notifications_config: MaintNotificationsConfig | None = None,
322 ):
323 """
324 Initialize a new Redis client.
326 To specify a retry policy for specific errors, you have two options:
328 1. Set the `retry_on_error` to a list of the error/s to retry on, and
329 you can also set `retry` to a valid `Retry` object(in case the default
330 one is not appropriate) - with this approach the retries will be triggered
331 on the default errors specified in the Retry object enriched with the
332 errors specified in `retry_on_error`.
334 2. Define a `Retry` object with configured 'supported_errors' and set
335 it to the `retry` parameter - with this approach you completely redefine
336 the errors on which retries will happen.
338 `retry_on_timeout` is deprecated - please include the TimeoutError
339 either in the Retry object or in the `retry_on_error` list.
341 When 'connection_pool' is provided - the retry configuration of the
342 provided pool will be used.
344 Args:
346 socket_keepalive:
347 if `True`, TCP keepalive is enabled for TCP socket connections.
348 Argument is ignored when connection_pool is provided.
349 socket_keepalive_options:
350 mapping of TCP keepalive socket option constants to values, for
351 example `{socket.TCP_KEEPIDLE: 30}`. If left unspecified, redis-py
352 uses TCP keepalive defaults when `socket_keepalive` is enabled:
353 idle 30 seconds, interval 5 seconds, and 3 probes. Platform-specific
354 options that are not available are skipped. Pass `None` or `{}` to
355 avoid setting additional TCP keepalive options. Argument is ignored
356 when connection_pool is provided.
357 maint_notifications_config:
358 configures the pool to support maintenance notifications - see
359 `redis.maint_notifications.MaintNotificationsConfig` for details.
360 Only supported with RESP3
361 If not provided and protocol is RESP3, the maintenance notifications
362 will be enabled by default (logic is included in the connection pool
363 initialization).
364 Argument is ignored when connection_pool is provided.
365 """
366 kwargs: Dict[str, Any]
367 if event_dispatcher is None:
368 self._event_dispatcher = EventDispatcher()
369 else:
370 self._event_dispatcher = event_dispatcher
371 # auto_close_connection_pool only has an effect if connection_pool is
372 # None. It is assumed that if connection_pool is not None, the user
373 # wants to manage the connection pool themselves.
374 if auto_close_connection_pool is not None:
375 warnings.warn(
376 DeprecationWarning(
377 '"auto_close_connection_pool" is deprecated '
378 "since version 5.0.1. "
379 "Please create a ConnectionPool explicitly and "
380 "provide to the Redis() constructor instead."
381 )
382 )
383 else:
384 auto_close_connection_pool = True
386 if not connection_pool:
387 # Create internal connection pool, expected to be closed by Redis instance
388 if not retry_on_error:
389 retry_on_error = []
391 # Handle driver_info: if provided, use it; otherwise create from lib_name/lib_version.
392 computed_driver_info = resolve_driver_info(
393 driver_info, lib_name, lib_version
394 )
396 kwargs = {
397 "db": db,
398 "username": username,
399 "password": password,
400 "credential_provider": credential_provider,
401 "socket_timeout": socket_timeout,
402 "socket_read_size": socket_read_size,
403 "encoding": encoding,
404 "encoding_errors": encoding_errors,
405 "decode_responses": decode_responses,
406 "retry_on_error": retry_on_error,
407 "retry": copy.deepcopy(retry),
408 "max_connections": max_connections,
409 "health_check_interval": health_check_interval,
410 "client_name": client_name,
411 "driver_info": computed_driver_info,
412 "redis_connect_func": redis_connect_func,
413 "protocol": protocol,
414 "legacy_responses": legacy_responses,
415 }
416 # based on input, setup appropriate connection args
417 if unix_socket_path is not None:
418 if (
419 maint_notifications_config
420 and maint_notifications_config.enabled is True
421 ):
422 raise RedisError(
423 "Maintenance notifications are not supported with Unix "
424 "domain socket connections"
425 )
426 kwargs.update(
427 {
428 "path": unix_socket_path,
429 "connection_class": UnixDomainSocketConnection,
430 "maint_notifications_config": MaintNotificationsConfig(
431 enabled=False
432 ),
433 }
434 )
435 else:
436 # TCP specific options
437 kwargs.update(
438 {
439 "host": host,
440 "port": port,
441 "socket_connect_timeout": socket_connect_timeout,
442 "socket_keepalive": socket_keepalive,
443 "socket_keepalive_options": socket_keepalive_options,
444 }
445 )
447 if ssl:
448 kwargs.update(
449 {
450 "connection_class": SSLConnection,
451 "ssl_keyfile": ssl_keyfile,
452 "ssl_certfile": ssl_certfile,
453 "ssl_cert_reqs": ssl_cert_reqs,
454 "ssl_include_verify_flags": ssl_include_verify_flags,
455 "ssl_exclude_verify_flags": ssl_exclude_verify_flags,
456 "ssl_ca_certs": ssl_ca_certs,
457 "ssl_ca_data": ssl_ca_data,
458 "ssl_ca_path": ssl_ca_path,
459 "ssl_check_hostname": ssl_check_hostname,
460 "ssl_min_version": ssl_min_version,
461 "ssl_ciphers": ssl_ciphers,
462 "ssl_password": ssl_password,
463 }
464 )
465 maint_notifications_enabled = (
466 maint_notifications_config and maint_notifications_config.enabled
467 )
468 if maint_notifications_enabled and not check_protocol_version(protocol, 3):
469 raise RedisError(
470 "Maintenance notifications handlers on connection are only supported with RESP version 3"
471 )
472 if maint_notifications_config:
473 kwargs.update(
474 {
475 "maint_notifications_config": maint_notifications_config,
476 }
477 )
478 # This arg only used if no pool is passed in
479 self.auto_close_connection_pool = auto_close_connection_pool
480 connection_pool = ConnectionPool(**kwargs)
481 self._event_dispatcher.dispatch(
482 AfterPooledConnectionsInstantiationEvent(
483 [connection_pool], ClientType.ASYNC, credential_provider
484 )
485 )
486 else:
487 # If a pool is passed in, do not close it
488 self.auto_close_connection_pool = False
489 self._event_dispatcher.dispatch(
490 AfterPooledConnectionsInstantiationEvent(
491 [connection_pool], ClientType.ASYNC, credential_provider
492 )
493 )
495 self.connection_pool = connection_pool
496 self.single_connection_client = single_connection_client
497 self.connection: Optional[Connection] = None
499 connection_kwargs = self.connection_pool.connection_kwargs
500 self.response_callbacks = CaseInsensitiveDict(
501 get_response_callbacks(
502 user_protocol=connection_kwargs.get("protocol"),
503 legacy_responses=connection_kwargs.get("legacy_responses", True),
504 )
505 )
507 # If using a single connection client, we need to lock creation-of and use-of
508 # the client in order to avoid race conditions such as using asyncio.gather
509 # on a set of redis commands
510 self._single_conn_lock = asyncio.Lock()
512 # When used as an async context manager, we need to increment and decrement
513 # a usage counter so that we can close the connection pool when no one is
514 # using the client.
515 self._usage_counter = 0
516 self._usage_lock = asyncio.Lock()
518 def __repr__(self):
519 return (
520 f"<{self.__class__.__module__}.{self.__class__.__name__}"
521 f"({self.connection_pool!r})>"
522 )
524 def __await__(self):
525 return self.initialize().__await__()
527 async def initialize(self: _RedisT) -> _RedisT:
528 if self.single_connection_client:
529 async with self._single_conn_lock:
530 if self.connection is None:
531 self.connection = await self.connection_pool.get_connection()
533 self._event_dispatcher.dispatch(
534 AfterSingleConnectionInstantiationEvent(
535 self.connection, ClientType.ASYNC, self._single_conn_lock
536 )
537 )
538 return self
540 def set_response_callback(self, command: str, callback: ResponseCallbackT):
541 """Set a custom Response Callback"""
542 self.response_callbacks[command] = callback
544 def get_encoder(self):
545 """Get the connection pool's encoder"""
546 return self.connection_pool.get_encoder()
548 def get_connection_kwargs(self):
549 """Get the connection's key-word arguments"""
550 return self.connection_pool.connection_kwargs
552 @property
553 def himport_registry(self) -> HImportRegistry:
554 """The client's HIMPORT fieldset registry (empty if none was declared).
556 Read-only: the registry is mutated only through the HIMPORT command methods.
557 """
558 return self.connection_pool.himport_registry
560 def get_retry(self) -> Optional[Retry]:
561 return self.get_connection_kwargs().get("retry")
563 def set_retry(self, retry: Retry) -> None:
564 self.get_connection_kwargs().update({"retry": retry})
565 self.connection_pool.set_retry(retry)
567 def load_external_module(self, funcname, func):
568 """
569 This function can be used to add externally defined redis modules,
570 and their namespaces to the redis client.
572 funcname - A string containing the name of the function to create
573 func - The function, being added to this class.
575 ex: Assume that one has a custom redis module named foomod that
576 creates command named 'foo.dothing' and 'foo.anotherthing' in redis.
577 To load function functions into this namespace:
579 from redis import Redis
580 from foomodule import F
581 r = Redis()
582 r.load_external_module("foo", F)
583 r.foo().dothing('your', 'arguments')
585 For a concrete example see the reimport of the redisjson module in
586 tests/test_connection.py::test_loading_external_modules
587 """
588 setattr(self, funcname, func)
590 def pipeline(
591 self, transaction: bool = True, shard_hint: Optional[str] = None
592 ) -> "Pipeline":
593 """
594 Return a new pipeline object that can queue multiple commands for
595 later execution. ``transaction`` indicates whether all commands
596 should be executed atomically. Apart from making a group of operations
597 atomic, pipelines are useful for reducing the back-and-forth overhead
598 between the client and server.
599 """
600 return Pipeline(
601 self.connection_pool, self.response_callbacks, transaction, shard_hint
602 )
604 async def transaction(
605 self,
606 func: Callable[["Pipeline"], Union[Any, Awaitable[Any]]],
607 *watches: KeyT,
608 shard_hint: Optional[str] = None,
609 value_from_callable: bool = False,
610 watch_delay: Optional[float] = None,
611 ):
612 """
613 Convenience method for executing the callable `func` as a transaction
614 while watching all keys specified in `watches`. The 'func' callable
615 should expect a single argument which is a Pipeline object.
616 """
617 pipe: Pipeline
618 async with self.pipeline(True, shard_hint) as pipe:
619 while True:
620 try:
621 if watches:
622 await pipe.watch(*watches)
623 func_value = func(pipe)
624 if inspect.isawaitable(func_value):
625 func_value = await func_value
626 exec_value = await pipe.execute()
627 return func_value if value_from_callable else exec_value
628 except WatchError:
629 if watch_delay is not None and watch_delay > 0:
630 await asyncio.sleep(watch_delay)
631 continue
633 def lock(
634 self,
635 name: KeyT,
636 timeout: Optional[float] = None,
637 sleep: float = 0.1,
638 blocking: bool = True,
639 blocking_timeout: Optional[float] = None,
640 lock_class: Optional[Type[Lock]] = None,
641 thread_local: bool = True,
642 raise_on_release_error: bool = True,
643 ) -> Lock:
644 """
645 Return a new Lock object using key ``name`` that mimics
646 the behavior of threading.Lock.
648 If specified, ``timeout`` indicates a maximum life for the lock.
649 By default, it will remain locked until release() is called.
651 ``sleep`` indicates the amount of time to sleep per loop iteration
652 when the lock is in blocking mode and another client is currently
653 holding the lock.
655 ``blocking`` indicates whether calling ``acquire`` should block until
656 the lock has been acquired or to fail immediately, causing ``acquire``
657 to return False and the lock not being acquired. Defaults to True.
658 Note this value can be overridden by passing a ``blocking``
659 argument to ``acquire``.
661 ``blocking_timeout`` indicates the maximum amount of time in seconds to
662 spend trying to acquire the lock. A value of ``None`` indicates
663 continue trying forever. ``blocking_timeout`` can be specified as a
664 float or integer, both representing the number of seconds to wait.
666 ``lock_class`` forces the specified lock implementation. Note that as
667 of redis-py 3.0, the only lock class we implement is ``Lock`` (which is
668 a Lua-based lock). So, it's unlikely you'll need this parameter, unless
669 you have created your own custom lock class.
671 ``thread_local`` indicates whether the lock token is placed in
672 thread-local storage. By default, the token is placed in thread local
673 storage so that a thread only sees its token, not a token set by
674 another thread. Consider the following timeline:
676 time: 0, thread-1 acquires `my-lock`, with a timeout of 5 seconds.
677 thread-1 sets the token to "abc"
678 time: 1, thread-2 blocks trying to acquire `my-lock` using the
679 Lock instance.
680 time: 5, thread-1 has not yet completed. redis expires the lock
681 key.
682 time: 5, thread-2 acquired `my-lock` now that it's available.
683 thread-2 sets the token to "xyz"
684 time: 6, thread-1 finishes its work and calls release(). if the
685 token is *not* stored in thread local storage, then
686 thread-1 would see the token value as "xyz" and would be
687 able to successfully release the thread-2's lock.
689 ``raise_on_release_error`` indicates whether to raise an exception when
690 the lock is no longer owned when exiting the context manager. By default,
691 this is True, meaning an exception will be raised. If False, the warning
692 will be logged and the exception will be suppressed.
694 In some use cases it's necessary to disable thread local storage. For
695 example, if you have code where one thread acquires a lock and passes
696 that lock instance to a worker thread to release later. If thread
697 local storage isn't disabled in this case, the worker thread won't see
698 the token set by the thread that acquired the lock. Our assumption
699 is that these cases aren't common and as such default to using
700 thread local storage."""
701 if lock_class is None:
702 lock_class = Lock
703 return lock_class(
704 self,
705 name,
706 timeout=timeout,
707 sleep=sleep,
708 blocking=blocking,
709 blocking_timeout=blocking_timeout,
710 thread_local=thread_local,
711 raise_on_release_error=raise_on_release_error,
712 )
714 def pubsub(self, **kwargs) -> "PubSub":
715 """
716 Return a Publish/Subscribe object. With this object, you can
717 subscribe to channels and listen for messages that get published to
718 them.
719 """
720 return PubSub(
721 self.connection_pool, event_dispatcher=self._event_dispatcher, **kwargs
722 )
724 def keyspace_notifications(
725 self,
726 key_prefix: Union[str, bytes, None] = None,
727 ignore_subscribe_messages: bool = True,
728 ) -> "AsyncKeyspaceNotifications":
729 """
730 Return an :class:`~redis.asyncio.keyspace_notifications.AsyncKeyspaceNotifications`
731 object for subscribing to keyspace and keyevent notifications.
733 Note: Keyspace notifications must be enabled on the Redis server via
734 the ``notify-keyspace-events`` configuration option.
736 Args:
737 key_prefix: Optional prefix to filter and strip from keys in
738 notifications.
739 ignore_subscribe_messages: If True, subscribe/unsubscribe
740 confirmations are not returned by
741 get_message/listen.
742 """
743 from redis.asyncio.keyspace_notifications import AsyncKeyspaceNotifications
745 return AsyncKeyspaceNotifications(
746 self,
747 key_prefix=key_prefix,
748 ignore_subscribe_messages=ignore_subscribe_messages,
749 )
751 def monitor(self) -> "Monitor":
752 return Monitor(self.connection_pool)
754 def client(self) -> "Redis":
755 return self.__class__(
756 connection_pool=self.connection_pool, single_connection_client=True
757 )
759 async def __aenter__(self: _RedisT) -> _RedisT:
760 """
761 Async context manager entry. Increments a usage counter so that the
762 connection pool is only closed (via aclose()) when no context is using
763 the client.
764 """
765 await self._increment_usage()
766 try:
767 # Initialize the client (i.e. establish connection, etc.)
768 return await self.initialize()
769 except Exception:
770 # If initialization fails, decrement the counter to keep it in sync
771 await self._decrement_usage()
772 raise
774 async def _increment_usage(self) -> int:
775 """
776 Helper coroutine to increment the usage counter while holding the lock.
777 Returns the new value of the usage counter.
778 """
779 async with self._usage_lock:
780 self._usage_counter += 1
781 return self._usage_counter
783 async def _decrement_usage(self) -> int:
784 """
785 Helper coroutine to decrement the usage counter while holding the lock.
786 Returns the new value of the usage counter.
787 """
788 async with self._usage_lock:
789 self._usage_counter -= 1
790 return self._usage_counter
792 async def __aexit__(self, exc_type, exc_value, traceback):
793 """
794 Async context manager exit. Decrements a usage counter. If this is the
795 last exit (counter becomes zero), the client closes its connection pool.
796 """
797 current_usage = await asyncio.shield(self._decrement_usage())
798 if current_usage == 0:
799 # This was the last active context, so disconnect the pool.
800 await asyncio.shield(self.aclose())
802 _DEL_MESSAGE = "Unclosed Redis client"
804 # passing _warnings and _grl as argument default since they may be gone
805 # by the time __del__ is called at shutdown
806 def __del__(
807 self,
808 _warn: Any = warnings.warn,
809 _grl: Any = asyncio.get_running_loop,
810 ) -> None:
811 if hasattr(self, "connection") and (self.connection is not None):
812 _warn(f"Unclosed client session {self!r}", ResourceWarning, source=self)
813 try:
814 context = {"client": self, "message": self._DEL_MESSAGE}
815 _grl().call_exception_handler(context)
816 except RuntimeError:
817 pass
818 self.connection._close()
820 async def aclose(self, close_connection_pool: Optional[bool] = None) -> None:
821 """
822 Closes Redis client connection
824 Args:
825 close_connection_pool:
826 decides whether to close the connection pool used by this Redis client,
827 overriding Redis.auto_close_connection_pool.
828 By default, let Redis.auto_close_connection_pool decide
829 whether to close the connection pool.
830 """
831 conn = self.connection
832 if conn:
833 self.connection = None
834 await self.connection_pool.release(conn)
835 if close_connection_pool or (
836 close_connection_pool is None and self.auto_close_connection_pool
837 ):
838 await self.connection_pool.aclose()
840 @deprecated_function(version="5.0.1", reason="Use aclose() instead", name="close")
841 async def close(self, close_connection_pool: Optional[bool] = None) -> None:
842 """
843 Alias for aclose(), for backwards compatibility
844 """
845 await self.aclose(close_connection_pool)
847 async def _send_command_parse_response(self, conn, command_name, *args, **options):
848 """
849 Send a command and parse the response
850 """
851 # HIMPORT SET is the one command whose wire form depends on per-connection
852 # state: the fieldset must be PREPAREd on this connection first, and any
853 # fieldset discarded since this connection last reconciled must be dropped.
854 # Handling it here (rather than in himport_set) lets himport_set reuse the
855 # full execute_command machinery — retry, disconnect-on-error, pooling — so
856 # a failed HIMPORT SET disconnects the connection like any other command.
857 # This per-command branch in the hot dispatch path is deliberate and has no
858 # cleaner alternative: this is the only seam where the concrete borrowed
859 # connection is known, and connection-scoped session setup can only happen
860 # once that connection is chosen. The overhead is one string compare per
861 # command.
862 himport_set = parse_himport_set_args(args)
863 if himport_set is not None:
864 # ``args`` is an HIMPORT SET in either the joined ("HIMPORT SET", key,
865 # ...) or split ("HIMPORT", "SET", key, ...) raw form; the operands come
866 # back at the right offsets for the form. A command with too few operands
867 # returns None and falls through to the normal send path so the server
868 # returns its arity error instead of a client-side IndexError here.
869 key, fieldset_name, values = himport_set
870 return await self._himport_execute_set(conn, key, fieldset_name, values)
871 await conn.send_command(*args)
872 return await self.parse_response(conn, command_name, **options)
874 async def _himport_reconcile_discards(self, conn):
875 """Delegate to the shared async HIMPORT executor."""
876 return await _himport_exec.reconcile_discards(self, conn)
878 async def _himport_prepare_and_set(
879 self, conn, key, fieldset_name, values, fieldset
880 ):
881 """Delegate to the shared async HIMPORT executor."""
882 return await _himport_exec.prepare_and_set(
883 self, conn, key, fieldset_name, values, fieldset
884 )
886 async def _himport_execute_set(self, conn, key, fieldset_name, values):
887 """Delegate to the shared async HIMPORT executor."""
888 return await _himport_exec.execute_set(self, conn, key, fieldset_name, values)
890 async def _close_connection(
891 self,
892 conn: Connection,
893 error: Optional[BaseException] = None,
894 failure_count: Optional[int] = None,
895 start_time: Optional[float] = None,
896 command_name: Optional[str] = None,
897 ):
898 """
899 Close the connection before retrying.
901 The supported exceptions are already checked in the
902 retry object so we don't need to do it here.
904 After we disconnect the connection, it will try to reconnect and
905 do a health check as part of the send_command logic(on connection level).
906 """
907 if (
908 error
909 and failure_count is not None
910 and failure_count <= conn.retry.get_retries()
911 ):
912 await record_operation_duration(
913 command_name=command_name,
914 duration_seconds=time.monotonic() - start_time,
915 server_address=getattr(conn, "host", None),
916 server_port=getattr(conn, "port", None),
917 db_namespace=str(conn.db),
918 error=error,
919 retry_attempts=failure_count,
920 )
922 await conn.disconnect(error=error, failure_count=failure_count)
924 # COMMAND EXECUTION AND PROTOCOL PARSING
925 async def execute_command(self, *args, **options):
926 """Execute a command and return a parsed response"""
927 await self.initialize()
928 pool = self.connection_pool
929 command_name = args[0]
930 conn = self.connection or await pool.get_connection()
932 # Start timing for observability
933 start_time = time.monotonic()
934 # Track actual retry attempts for error reporting
935 actual_retry_attempts = 0
937 def failure_callback(error, failure_count):
938 nonlocal actual_retry_attempts
939 actual_retry_attempts = failure_count
940 return self._close_connection(
941 conn, error, failure_count, start_time, command_name
942 )
944 if self.single_connection_client:
945 await self._single_conn_lock.acquire()
946 try:
947 result = await conn.retry.call_with_retry(
948 lambda: self._send_command_parse_response(
949 conn, command_name, *args, **options
950 ),
951 failure_callback,
952 with_failure_count=True,
953 )
955 await record_operation_duration(
956 command_name=command_name,
957 duration_seconds=time.monotonic() - start_time,
958 server_address=getattr(conn, "host", None),
959 server_port=getattr(conn, "port", None),
960 db_namespace=str(conn.db),
961 )
962 return result
963 except Exception as e:
964 await record_error_count(
965 server_address=getattr(conn, "host", None),
966 server_port=getattr(conn, "port", None),
967 network_peer_address=getattr(conn, "host", None),
968 network_peer_port=getattr(conn, "port", None),
969 error_type=e,
970 retry_attempts=actual_retry_attempts,
971 is_internal=False,
972 )
973 raise
974 finally:
975 try:
976 if self.single_connection_client and conn and conn.should_reconnect():
977 await self._close_connection(conn)
978 await conn.connect()
979 finally:
980 if self.single_connection_client:
981 self._single_conn_lock.release()
982 if not self.connection:
983 await pool.release(conn)
985 async def parse_response(
986 self, connection: Connection, command_name: Union[str, bytes], **options
987 ):
988 """Parses a response from the Redis server"""
989 try:
990 if NEVER_DECODE in options:
991 response = await connection.read_response(disable_decoding=True)
992 options.pop(NEVER_DECODE)
993 else:
994 response = await connection.read_response()
995 except ResponseError:
996 if EMPTY_RESPONSE in options:
997 return options[EMPTY_RESPONSE]
998 raise
1000 if EMPTY_RESPONSE in options:
1001 options.pop(EMPTY_RESPONSE)
1003 # Remove keys entry, it needs only for cache.
1004 options.pop("keys", None)
1006 if command_name in self.response_callbacks:
1007 # Mypy bug: https://github.com/python/mypy/issues/10977
1008 command_name = cast(str, command_name)
1009 retval = self.response_callbacks[command_name](response, **options)
1010 return await retval if inspect.isawaitable(retval) else retval
1011 return response
1013 # HIMPORT orchestration (async mirror of redis.client.Redis). See
1014 # ``.agents/himport_client_support_spec.md``.
1016 @experimental_method()
1017 async def himport_prepare(
1018 self, fieldset_name: str, fields: Iterable[FieldT]
1019 ) -> bool:
1020 """Declare an HIMPORT fieldset for use by :meth:`himport_set`."""
1021 await self.initialize()
1022 fieldset = self.himport_registry.prepare(fieldset_name, fields)
1023 conn = self.connection
1024 if self.single_connection_client and conn is not None and conn.is_connected:
1025 await self.himport_prepare_internal(fieldset_name, fieldset.fields)
1026 conn._himport_prepared[fieldset_name] = fieldset.version
1027 return True
1029 @experimental_method()
1030 async def himport_discard(self, fieldset_name: str) -> int:
1031 """Remove a fieldset from the registry."""
1032 await self.initialize()
1033 removed = self.himport_registry.discard(fieldset_name)
1034 conn = self.connection
1035 if self.single_connection_client and conn is not None and conn.is_connected:
1036 if removed:
1037 await self.himport_discard_internal(fieldset_name)
1038 conn._himport_prepared.pop(fieldset_name, None)
1039 conn._himport_reconciled_revision = self.himport_registry.revision
1040 return 1 if removed else 0
1042 @experimental_method()
1043 async def himport_discard_all(self) -> int:
1044 """Remove all fieldsets from the registry."""
1045 await self.initialize()
1046 count = self.himport_registry.discard_all()
1047 conn = self.connection
1048 if self.single_connection_client and conn is not None and conn.is_connected:
1049 if count:
1050 await self.himport_discard_all_internal()
1051 conn._himport_prepared.clear()
1052 conn._himport_reconciled_revision = self.himport_registry.revision
1053 return count
1056StrictRedis = Redis
1059class MonitorCommandInfo(TypedDict):
1060 time: float
1061 db: int
1062 client_address: str
1063 client_port: str
1064 client_type: str
1065 command: str
1068class Monitor:
1069 """
1070 Monitor is useful for handling the MONITOR command to the redis server.
1071 next_command() method returns one command from monitor
1072 listen() method yields commands from monitor.
1073 """
1075 monitor_re = re.compile(r"\[(\d+) (.*?)\] (.*)")
1076 command_re = re.compile(r'"(.*?)(?<!\\)"')
1078 def __init__(self, connection_pool: ConnectionPool):
1079 self.connection_pool = connection_pool
1080 self.connection: Optional[Connection] = None
1082 async def connect(self):
1083 if self.connection is None:
1084 self.connection = await self.connection_pool.get_connection()
1086 async def __aenter__(self):
1087 await self.connect()
1088 await self.connection.send_command("MONITOR")
1089 # check that monitor returns 'OK', but don't return it to user
1090 response = await self.connection.read_response()
1091 if not bool_ok(response):
1092 raise RedisError(f"MONITOR failed: {response}")
1093 return self
1095 async def __aexit__(self, *args):
1096 await self.connection.disconnect()
1097 await self.connection_pool.release(self.connection)
1099 async def next_command(self) -> MonitorCommandInfo:
1100 """Parse the response from a monitor command"""
1101 await self.connect()
1102 response = await self.connection.read_response()
1103 if isinstance(response, bytes):
1104 response = self.connection.encoder.decode(response, force=True)
1105 command_time, command_data = response.split(" ", 1)
1106 m = self.monitor_re.match(command_data)
1107 db_id, client_info, command = m.groups()
1108 command = " ".join(self.command_re.findall(command))
1109 # Redis escapes double quotes because each piece of the command
1110 # string is surrounded by double quotes. We don't have that
1111 # requirement so remove the escaping and leave the quote.
1112 command = command.replace('\\"', '"')
1114 if client_info == "lua":
1115 client_address = "lua"
1116 client_port = ""
1117 client_type = "lua"
1118 elif client_info.startswith("unix"):
1119 client_address = "unix"
1120 client_port = client_info[5:]
1121 client_type = "unix"
1122 else:
1123 # use rsplit as ipv6 addresses contain colons
1124 client_address, client_port = client_info.rsplit(":", 1)
1125 client_type = "tcp"
1126 return {
1127 "time": float(command_time),
1128 "db": int(db_id),
1129 "client_address": client_address,
1130 "client_port": client_port,
1131 "client_type": client_type,
1132 "command": command,
1133 }
1135 async def listen(self) -> AsyncIterator[MonitorCommandInfo]:
1136 """Listen for commands coming to the server."""
1137 while True:
1138 yield await self.next_command()
1141class PubSub:
1142 """
1143 PubSub provides publish, subscribe and listen support to Redis channels.
1145 After subscribing to one or more channels, the listen() method will block
1146 until a message arrives on one of the subscribed channels. That message
1147 will be returned and it's safe to start listening again.
1148 """
1150 PUBLISH_MESSAGE_TYPES = ("message", "pmessage", "smessage")
1151 UNSUBSCRIBE_MESSAGE_TYPES = ("unsubscribe", "punsubscribe", "sunsubscribe")
1152 HEALTH_CHECK_MESSAGE = "redis-py-health-check"
1154 def __init__(
1155 self,
1156 connection_pool: ConnectionPool,
1157 shard_hint: Optional[str] = None,
1158 ignore_subscribe_messages: bool = False,
1159 encoder=None,
1160 push_handler_func: Optional[Callable] = None,
1161 event_dispatcher: Optional["EventDispatcher"] = None,
1162 ):
1163 if event_dispatcher is None:
1164 self._event_dispatcher = EventDispatcher()
1165 else:
1166 self._event_dispatcher = event_dispatcher
1167 self.connection_pool = connection_pool
1168 self.shard_hint = shard_hint
1169 self.ignore_subscribe_messages = ignore_subscribe_messages
1170 self.connection = None
1171 # we need to know the encoding options for this connection in order
1172 # to lookup channel and pattern names for callback handlers.
1173 self.encoder = encoder
1174 self.push_handler_func = push_handler_func
1175 if self.encoder is None:
1176 self.encoder = self.connection_pool.get_encoder()
1177 if self.encoder.decode_responses:
1178 self.health_check_response = [
1179 ["pong", self.HEALTH_CHECK_MESSAGE],
1180 self.HEALTH_CHECK_MESSAGE,
1181 ]
1182 else:
1183 self.health_check_response = [
1184 [b"pong", self.encoder.encode(self.HEALTH_CHECK_MESSAGE)],
1185 self.encoder.encode(self.HEALTH_CHECK_MESSAGE),
1186 ]
1187 if self.push_handler_func is None:
1188 _set_info_logger()
1189 self.channels = {}
1190 self.pending_unsubscribe_channels = set()
1191 self.patterns = {}
1192 self.pending_unsubscribe_patterns = set()
1193 self.shard_channels = {}
1194 self.pending_unsubscribe_shard_channels = set()
1195 self._lock = asyncio.Lock()
1197 async def __aenter__(self):
1198 return self
1200 async def __aexit__(self, exc_type, exc_value, traceback):
1201 await self.aclose()
1203 def __del__(self):
1204 if self.connection:
1205 self.connection.deregister_connect_callback(self.on_connect)
1207 async def aclose(self):
1208 # In case a connection property does not yet exist
1209 # (due to a crash earlier in the Redis() constructor), return
1210 # immediately as there is nothing to clean-up.
1211 if not hasattr(self, "connection"):
1212 return
1213 async with self._lock:
1214 if self.connection:
1215 # Use nowait=True to avoid awaiting StreamWriter.wait_closed(),
1216 # which can deadlock when a concurrent reader task (e.g. one
1217 # running pubsub.run() or get_message(block=True)) still holds
1218 # the transport. See https://github.com/redis/redis-py/issues/3941
1219 await self.connection.disconnect(nowait=True)
1220 self.connection.deregister_connect_callback(self.on_connect)
1221 await self.connection_pool.release(self.connection)
1222 self.connection = None
1223 self.channels = {}
1224 self.pending_unsubscribe_channels = set()
1225 self.patterns = {}
1226 self.pending_unsubscribe_patterns = set()
1227 self.shard_channels = {}
1228 self.pending_unsubscribe_shard_channels = set()
1230 @deprecated_function(version="5.0.1", reason="Use aclose() instead", name="close")
1231 async def close(self) -> None:
1232 """Alias for aclose(), for backwards compatibility"""
1233 await self.aclose()
1235 @deprecated_function(version="5.0.1", reason="Use aclose() instead", name="reset")
1236 async def reset(self) -> None:
1237 """Alias for aclose(), for backwards compatibility"""
1238 await self.aclose()
1240 async def _resubscribe(self, subscribed, subscribe_fn) -> None:
1241 # Replay handler-backed subscriptions as positional Subscription objects
1242 # so binary names never need to be decoded into keyword argument keys.
1243 subscriptions = pubsub_subscription_args(subscribed)
1244 if subscriptions:
1245 await subscribe_fn(*subscriptions)
1247 async def _resubscribe_shard_channels(self) -> None:
1248 await self._resubscribe(self.shard_channels, self.ssubscribe)
1250 async def on_connect(self, connection: Connection):
1251 """Re-subscribe to any channels and patterns previously subscribed to"""
1252 self.pending_unsubscribe_channels.clear()
1253 self.pending_unsubscribe_patterns.clear()
1254 self.pending_unsubscribe_shard_channels.clear()
1255 if self.channels:
1256 await self._resubscribe(self.channels, self.subscribe)
1257 if self.patterns:
1258 await self._resubscribe(self.patterns, self.psubscribe)
1259 if self.shard_channels:
1260 await self._resubscribe_shard_channels()
1262 @property
1263 def subscribed(self):
1264 """Indicates if there are subscriptions to any channels or patterns"""
1265 return bool(self.channels or self.patterns or self.shard_channels)
1267 async def execute_command(self, *args: EncodableT):
1268 """Execute a publish/subscribe command"""
1270 # NOTE: don't parse the response in this function -- it could pull a
1271 # legitimate message off the stack if the connection is already
1272 # subscribed to one or more channels
1274 await self.connect()
1275 connection = self.connection
1276 kwargs = {"check_health": not self.subscribed}
1277 await self._execute(connection, connection.send_command, *args, **kwargs)
1279 async def connect(self):
1280 """
1281 Ensure that the PubSub is connected
1282 """
1283 if self.connection is None:
1284 self.connection = await self.connection_pool.get_connection()
1285 # register a callback that re-subscribes to any channels we
1286 # were listening to when we were disconnected
1287 self.connection.register_connect_callback(self.on_connect)
1288 else:
1289 await self.connection.connect()
1290 if self.push_handler_func is not None:
1291 self.connection._parser.set_pubsub_push_handler(self.push_handler_func)
1293 self._event_dispatcher.dispatch(
1294 AfterPubSubConnectionInstantiationEvent(
1295 self.connection, self.connection_pool, ClientType.ASYNC, self._lock
1296 )
1297 )
1299 async def _reconnect(
1300 self,
1301 conn,
1302 error: Optional[BaseException] = None,
1303 failure_count: Optional[int] = None,
1304 start_time: Optional[float] = None,
1305 command_name: Optional[str] = None,
1306 ):
1307 """
1308 The supported exceptions are already checked in the
1309 retry object so we don't need to do it here.
1311 In this error handler we are trying to reconnect to the server.
1312 """
1313 if (
1314 error
1315 and failure_count is not None
1316 and failure_count <= conn.retry.get_retries()
1317 ):
1318 if command_name:
1319 await record_operation_duration(
1320 command_name=command_name,
1321 duration_seconds=time.monotonic() - start_time,
1322 server_address=getattr(conn, "host", None),
1323 server_port=getattr(conn, "port", None),
1324 db_namespace=str(conn.db),
1325 error=error,
1326 retry_attempts=failure_count,
1327 )
1328 await conn.disconnect(error=error, failure_count=failure_count)
1329 await conn.connect()
1331 async def _execute(self, conn, command, *args, **kwargs):
1332 """
1333 Connect manually upon disconnection. If the Redis server is down,
1334 this will fail and raise a ConnectionError as desired.
1335 After reconnection, the ``on_connect`` callback should have been
1336 called by the # connection to resubscribe us to any channels and
1337 patterns we were previously listening to
1338 """
1339 if not len(args) == 0:
1340 command_name = args[0]
1341 else:
1342 command_name = None
1344 # Start timing for observability
1345 start_time = time.monotonic()
1346 # Track actual retry attempts for error reporting
1347 actual_retry_attempts = 0
1349 def failure_callback(error, failure_count):
1350 nonlocal actual_retry_attempts
1351 actual_retry_attempts = failure_count
1352 return self._reconnect(conn, error, failure_count, start_time, command_name)
1354 try:
1355 response = await conn.retry.call_with_retry(
1356 lambda: command(*args, **kwargs),
1357 failure_callback,
1358 with_failure_count=True,
1359 )
1361 if command_name:
1362 await record_operation_duration(
1363 command_name=command_name,
1364 duration_seconds=time.monotonic() - start_time,
1365 server_address=getattr(conn, "host", None),
1366 server_port=getattr(conn, "port", None),
1367 db_namespace=str(conn.db),
1368 )
1370 return response
1371 except Exception as e:
1372 await record_error_count(
1373 server_address=getattr(conn, "host", None),
1374 server_port=getattr(conn, "port", None),
1375 network_peer_address=getattr(conn, "host", None),
1376 network_peer_port=getattr(conn, "port", None),
1377 error_type=e,
1378 retry_attempts=actual_retry_attempts,
1379 is_internal=False,
1380 )
1381 raise
1383 async def parse_response(self, block: bool = True, timeout: float = 0):
1384 """
1385 Parse the response from a publish/subscribe command.
1387 Args:
1388 block: If True, block indefinitely until a message is available.
1389 If False, return immediately if no message is available.
1390 Default: True
1391 timeout: The timeout in seconds for reading a response when block=False.
1392 This parameter is ignored when block=True.
1393 Default: 0 (return immediately if no data available)
1395 Returns:
1396 The parsed response from the server, or None if no message is available
1397 within the timeout period (when block=False).
1399 Important:
1400 The block and timeout parameters work together:
1401 - When block=True: timeout is IGNORED, method blocks indefinitely
1402 - When block=False: timeout is USED, method returns after timeout expires
1404 Typically, you should use get_message(timeout=X) instead of calling
1405 parse_response() directly. The get_message() method automatically sets
1406 block=False when a timeout is provided, and block=True when timeout=None.
1408 Example:
1409 # Block indefinitely (timeout is ignored)
1410 response = await pubsub.parse_response(block=True, timeout=0.1)
1412 # Non-blocking with 0.1 second timeout
1413 response = await pubsub.parse_response(block=False, timeout=0.1)
1415 # Non-blocking, return immediately
1416 response = await pubsub.parse_response(block=False, timeout=0)
1418 # Recommended: use get_message() instead
1419 msg = await pubsub.get_message(timeout=0.1) # automatically sets block=False
1420 msg = await pubsub.get_message(timeout=None) # automatically sets block=True
1421 """
1422 conn = self.connection
1423 if conn is None:
1424 raise RuntimeError(
1425 "pubsub connection not set: "
1426 "did you forget to call subscribe() or psubscribe()?"
1427 )
1429 await self.check_health()
1431 if not conn.is_connected:
1432 await conn.connect()
1434 # Block=True: signal "no timeout" to conn.read_response via
1435 # math.inf. The connection treats math.inf as the per-read
1436 # opt-in for blocking indefinitely without falling back to
1437 # self.socket_timeout. Reconnect/AUTH/HELLO/resubscribe
1438 # operations performed by the retry layer continue to honor
1439 # self.socket_timeout because they do not pass math.inf.
1440 #
1441 # TODO(next-major): when the async Connection.read_response
1442 # default for ``timeout`` is changed to SENTINEL, passing
1443 # ``timeout=None`` from this method will become the natural
1444 # "no timeout" signal and the math.inf hand-off can be
1445 # removed. That swap is a breaking change to the
1446 # Connection.read_response signature so it must wait for a
1447 # major release.
1448 read_timeout = math.inf if block else timeout
1449 response = await self._execute(
1450 conn,
1451 conn.read_response,
1452 timeout=read_timeout,
1453 disconnect_on_error=False,
1454 push_request=True,
1455 )
1457 if conn.health_check_interval and response in self.health_check_response:
1458 # ignore the health check message as user might not expect it
1459 return None
1460 return response
1462 async def check_health(self):
1463 conn = self.connection
1464 if conn is None:
1465 raise RuntimeError(
1466 "pubsub connection not set: "
1467 "did you forget to call subscribe() or psubscribe()?"
1468 )
1470 if (
1471 conn.health_check_interval
1472 and asyncio.get_running_loop().time() > conn.next_health_check
1473 ):
1474 await conn.send_command(
1475 "PING", self.HEALTH_CHECK_MESSAGE, check_health=False
1476 )
1478 def _normalize_keys(self, data: _NormalizeKeysT) -> _NormalizeKeysT:
1479 """
1480 normalize channel/pattern names to be either bytes or strings
1481 based on whether responses are automatically decoded. this saves us
1482 from coercing the value for each message coming in.
1483 """
1484 encode = self.encoder.encode
1485 decode = self.encoder.decode
1486 return {decode(encode(k)): v for k, v in data.items()} # type: ignore[return-value] # noqa: E501
1488 async def psubscribe(
1489 self, *args: ChannelT | Subscription, **kwargs: PubSubHandler
1490 ) -> None:
1491 """
1492 Subscribe to channel patterns.
1493 Patterns supplied as keyword arguments expect a pattern name as the
1494 key and a callable as the value.
1495 ``Subscription`` objects can also be supplied positionally with an
1496 optional handler.
1497 A pattern's callable will be invoked automatically
1498 when a message is received on that pattern rather than producing a
1499 message via ``listen()``.
1500 """
1501 new_patterns = parse_pubsub_subscriptions(args, kwargs)
1502 ret_val = await self.execute_command("PSUBSCRIBE", *new_patterns.keys())
1503 # update the patterns dict AFTER we send the command. we don't want to
1504 # subscribe twice to these patterns, once for the command and again
1505 # for the reconnection.
1506 new_patterns = self._normalize_keys(new_patterns)
1507 self.patterns.update(new_patterns)
1508 self.pending_unsubscribe_patterns.difference_update(new_patterns)
1509 return ret_val
1511 def punsubscribe(self, *args: ChannelT) -> Awaitable:
1512 """
1513 Unsubscribe from the supplied patterns. If empty, unsubscribe from
1514 all patterns.
1515 """
1516 patterns: Iterable[ChannelT]
1517 if args:
1518 parsed_args = list_or_args((args[0],), args[1:])
1519 patterns = self._normalize_keys(dict.fromkeys(parsed_args)).keys()
1520 else:
1521 parsed_args = []
1522 patterns = self.patterns
1523 self.pending_unsubscribe_patterns.update(patterns)
1524 return self.execute_command("PUNSUBSCRIBE", *parsed_args)
1526 async def subscribe(
1527 self, *args: ChannelT | Subscription, **kwargs: PubSubHandler
1528 ) -> None:
1529 """
1530 Subscribe to channels.
1531 Channels supplied as keyword arguments expect
1532 a channel name as the key and a callable as the value.
1533 ``Subscription`` objects can also be supplied positionally with an
1534 optional handler.
1535 A channel's callable will be invoked automatically
1536 when a message is received on that channel rather than producing a
1537 message via ``listen()`` or ``get_message()``.
1538 """
1539 new_channels = parse_pubsub_subscriptions(args, kwargs)
1540 ret_val = await self.execute_command("SUBSCRIBE", *new_channels.keys())
1541 # update the channels dict AFTER we send the command. we don't want to
1542 # subscribe twice to these channels, once for the command and again
1543 # for the reconnection.
1544 new_channels = self._normalize_keys(new_channels)
1545 self.channels.update(new_channels)
1546 self.pending_unsubscribe_channels.difference_update(new_channels)
1547 return ret_val
1549 def unsubscribe(self, *args) -> Awaitable:
1550 """
1551 Unsubscribe from the supplied channels. If empty, unsubscribe from
1552 all channels
1553 """
1554 if args:
1555 parsed_args = list_or_args(args[0], args[1:])
1556 channels = self._normalize_keys(dict.fromkeys(parsed_args))
1557 else:
1558 parsed_args = []
1559 channels = self.channels
1560 self.pending_unsubscribe_channels.update(channels)
1561 return self.execute_command("UNSUBSCRIBE", *parsed_args)
1563 async def ssubscribe(
1564 self,
1565 *args: ChannelT | Subscription,
1566 target_node: Any = None,
1567 **kwargs: PubSubHandler,
1568 ) -> None:
1569 """
1570 Subscribes the client to the specified shard channels.
1571 Channels supplied as keyword arguments expect a channel name as the key
1572 and a callable as the value.
1573 ``Subscription`` objects can also be supplied positionally
1574 with an optional handler.
1575 A channel's callable will be invoked automatically when a message
1576 is received on that channel rather than producing a message
1577 via ``listen()`` or ``get_sharded_message()``.
1578 """
1579 new_s_channels = parse_pubsub_subscriptions(args, kwargs)
1580 ret_val = await self.execute_command("SSUBSCRIBE", *new_s_channels.keys())
1581 # update the s_channels dict AFTER we send the command. we don't want to
1582 # subscribe twice to these channels, once for the command and again
1583 # for the reconnection.
1584 new_s_channels = self._normalize_keys(new_s_channels)
1585 self.shard_channels.update(new_s_channels)
1586 self.pending_unsubscribe_shard_channels.difference_update(new_s_channels)
1587 return ret_val
1589 def sunsubscribe(self, *args, target_node=None) -> Awaitable:
1590 """
1591 Unsubscribe from the supplied shard_channels. If empty, unsubscribe from
1592 all shard_channels
1593 """
1594 if args:
1595 args = list_or_args(args[0], args[1:])
1596 s_channels = self._normalize_keys(dict.fromkeys(args))
1597 else:
1598 s_channels = self.shard_channels
1599 self.pending_unsubscribe_shard_channels.update(s_channels)
1600 return self.execute_command("SUNSUBSCRIBE", *args)
1602 async def listen(self) -> AsyncIterator:
1603 """Listen for messages on channels this client has been subscribed to"""
1604 while self.subscribed:
1605 response = await self.handle_message(await self.parse_response(block=True))
1606 if response is not None:
1607 yield response
1609 async def get_message(
1610 self, ignore_subscribe_messages: bool = False, timeout: Optional[float] = 0.0
1611 ):
1612 """
1613 Get the next message if one is available, otherwise None.
1615 If timeout is specified, the system will wait for `timeout` seconds
1616 before returning. Timeout should be specified as a floating point
1617 number or None to wait indefinitely.
1618 """
1619 response = await self.parse_response(block=(timeout is None), timeout=timeout)
1620 if response:
1621 return await self.handle_message(response, ignore_subscribe_messages)
1622 return None
1624 def ping(self, message=None) -> Awaitable[bool]:
1625 """
1626 Ping the Redis server to test connectivity.
1628 Sends a PING command to the Redis server and returns True if the server
1629 responds with "PONG".
1630 """
1631 args = ["PING", message] if message is not None else ["PING"]
1632 return self.execute_command(*args)
1634 async def handle_message(self, response, ignore_subscribe_messages=False):
1635 """
1636 Parses a pub/sub message. If the channel or pattern was subscribed to
1637 with a message handler, the handler is invoked instead of a parsed
1638 message being returned.
1639 """
1640 if response is None:
1641 return None
1642 if isinstance(response, bytes):
1643 response = [b"pong", response] if response != b"PONG" else [b"pong", b""]
1644 message_type = str_if_bytes(response[0])
1645 if message_type == "pmessage":
1646 message = {
1647 "type": message_type,
1648 "pattern": response[1],
1649 "channel": response[2],
1650 "data": response[3],
1651 }
1652 elif message_type == "pong":
1653 message = {
1654 "type": message_type,
1655 "pattern": None,
1656 "channel": None,
1657 "data": response[1],
1658 }
1659 else:
1660 message = {
1661 "type": message_type,
1662 "pattern": None,
1663 "channel": response[1],
1664 "data": response[2],
1665 }
1667 if message_type in ["message", "pmessage"]:
1668 channel = str_if_bytes(message["channel"])
1669 await record_pubsub_message(
1670 direction=PubSubDirection.RECEIVE,
1671 channel=channel,
1672 )
1673 elif message_type == "smessage":
1674 channel = str_if_bytes(message["channel"])
1675 await record_pubsub_message(
1676 direction=PubSubDirection.RECEIVE,
1677 channel=channel,
1678 sharded=True,
1679 )
1681 # if this is an unsubscribe message, remove it from memory
1682 if message_type in self.UNSUBSCRIBE_MESSAGE_TYPES:
1683 if message_type == "punsubscribe":
1684 pattern = response[1]
1685 if pattern in self.pending_unsubscribe_patterns:
1686 self.pending_unsubscribe_patterns.remove(pattern)
1687 self.patterns.pop(pattern, None)
1688 elif message_type == "sunsubscribe":
1689 s_channel = response[1]
1690 if s_channel in self.pending_unsubscribe_shard_channels:
1691 self.pending_unsubscribe_shard_channels.remove(s_channel)
1692 self.shard_channels.pop(s_channel, None)
1693 else:
1694 channel = response[1]
1695 if channel in self.pending_unsubscribe_channels:
1696 self.pending_unsubscribe_channels.remove(channel)
1697 self.channels.pop(channel, None)
1699 if message_type in self.PUBLISH_MESSAGE_TYPES:
1700 # if there's a message handler, invoke it
1701 if message_type == "pmessage":
1702 handler = self.patterns.get(message["pattern"], None)
1703 elif message_type == "smessage":
1704 handler = self.shard_channels.get(message["channel"], None)
1705 else:
1706 handler = self.channels.get(message["channel"], None)
1707 if handler:
1708 if inspect.iscoroutinefunction(handler):
1709 await handler(message)
1710 else:
1711 handler(message)
1712 return None
1713 elif message_type != "pong":
1714 # this is a subscribe/unsubscribe message. ignore if we don't
1715 # want them
1716 if ignore_subscribe_messages or self.ignore_subscribe_messages:
1717 return None
1719 return message
1721 async def run(
1722 self,
1723 *,
1724 exception_handler: Optional["PSWorkerThreadExcHandlerT"] = None,
1725 poll_timeout: float = 1.0,
1726 pubsub=None,
1727 ) -> None:
1728 """Process pub/sub messages using registered callbacks.
1730 This is the equivalent of :py:meth:`redis.PubSub.run_in_thread` in
1731 redis-py, but it is a coroutine. To launch it as a separate task, use
1732 ``asyncio.create_task``:
1734 >>> task = asyncio.create_task(pubsub.run())
1736 To shut it down, use asyncio cancellation:
1738 >>> task.cancel()
1739 >>> await task
1740 """
1741 for channel, handler in self.channels.items():
1742 if handler is None:
1743 raise PubSubError(f"Channel: '{channel}' has no handler registered")
1744 for pattern, handler in self.patterns.items():
1745 if handler is None:
1746 raise PubSubError(f"Pattern: '{pattern}' has no handler registered")
1748 await self.connect()
1749 while True:
1750 try:
1751 if pubsub is None:
1752 await self.get_message(
1753 ignore_subscribe_messages=True, timeout=poll_timeout
1754 )
1755 else:
1756 await pubsub.get_message(
1757 ignore_subscribe_messages=True, timeout=poll_timeout
1758 )
1759 except asyncio.CancelledError:
1760 raise
1761 except BaseException as e:
1762 if exception_handler is None:
1763 raise
1764 res = exception_handler(e, self)
1765 if inspect.isawaitable(res):
1766 await res
1767 # Ensure that other tasks on the event loop get a chance to run
1768 # if we didn't have to block for I/O anywhere.
1769 await asyncio.sleep(0)
1772class PubsubWorkerExceptionHandler(Protocol):
1773 def __call__(self, e: BaseException, pubsub: PubSub): ...
1776class AsyncPubsubWorkerExceptionHandler(Protocol):
1777 async def __call__(self, e: BaseException, pubsub: PubSub): ...
1780PSWorkerThreadExcHandlerT = Union[
1781 PubsubWorkerExceptionHandler, AsyncPubsubWorkerExceptionHandler
1782]
1785CommandT = Tuple[Tuple[Union[str, bytes], ...], Mapping[str, Any]]
1786CommandStackT = List[CommandT]
1789class Pipeline(Redis): # lgtm [py/init-calls-subclass]
1790 """
1791 Pipelines provide a way to transmit multiple commands to the Redis server
1792 in one transmission. This is convenient for batch processing, such as
1793 saving all the values in a list to Redis.
1795 All commands executed within a pipeline(when running in transactional mode,
1796 which is the default behavior) are wrapped with MULTI and EXEC
1797 calls. This guarantees all commands executed in the pipeline will be
1798 executed atomically.
1800 Any command raising an exception does *not* halt the execution of
1801 subsequent commands in the pipeline. Instead, the exception is caught
1802 and its instance is placed into the response list returned by execute().
1803 Code iterating over the response list should be able to deal with an
1804 instance of an exception as a potential value. In general, these will be
1805 ResponseError exceptions, such as those raised when issuing a command
1806 on a key of a different datatype.
1807 """
1809 UNWATCH_COMMANDS = {"DISCARD", "EXEC", "UNWATCH"}
1811 def __init__(
1812 self,
1813 connection_pool: ConnectionPool,
1814 response_callbacks: MutableMapping[Union[str, bytes], ResponseCallbackT],
1815 transaction: bool,
1816 shard_hint: Optional[str],
1817 ):
1818 self.connection_pool = connection_pool
1819 self.connection = None
1820 self.response_callbacks = response_callbacks
1821 self.is_transaction = transaction
1822 self.shard_hint = shard_hint
1823 self.watching = False
1824 self.command_stack: CommandStackT = []
1825 self.scripts: Set[Script] = set()
1826 self.explicit_transaction = False
1828 async def __aenter__(self: _RedisT) -> _RedisT:
1829 return self
1831 async def __aexit__(self, exc_type, exc_value, traceback):
1832 await self.reset()
1834 def __await__(self):
1835 return self._async_self().__await__()
1837 _DEL_MESSAGE = "Unclosed Pipeline client"
1839 def __len__(self):
1840 return len(self.command_stack)
1842 def __bool__(self):
1843 """Pipeline instances should always evaluate to True"""
1844 return True
1846 async def _async_self(self):
1847 return self
1849 async def reset(self):
1850 self.command_stack = []
1851 self.scripts = set()
1852 try:
1853 # make sure to reset the connection state in the event that we were
1854 # watching something
1855 if self.watching and self.connection:
1856 try:
1857 # call this manually since our unwatch or
1858 # immediate_execute_command methods can call reset()
1859 await self.connection.send_command("UNWATCH")
1860 await self.connection.read_response()
1861 except ConnectionError:
1862 # disconnect will also remove any previous WATCHes
1863 if self.connection:
1864 await self.connection.disconnect()
1865 except asyncio.CancelledError:
1866 # Disconnect so any unread UNWATCH reply does not get
1867 # served to the next caller that takes the connection.
1868 if self.connection:
1869 await self.connection.disconnect()
1870 raise
1871 finally:
1872 self.watching = False
1873 self.explicit_transaction = False
1874 # We can safely return the connection to the pool here since we're
1875 # sure we're no longer WATCHing anything. Detach self.connection
1876 # before awaiting release: if a second cancel aborts the await,
1877 # the pipeline must not be left holding a reference to a
1878 # connection that is being returned to the pool. Shield the
1879 # release itself so a second cancel cannot split the pool's
1880 # internal in-use/available bookkeeping mid-update.
1881 if self.connection:
1882 connection, self.connection = self.connection, None
1883 await asyncio.shield(self.connection_pool.release(connection))
1885 async def aclose(self) -> None:
1886 """Alias for reset(), a standard method name for cleanup"""
1887 await self.reset()
1889 def multi(self):
1890 """
1891 Start a transactional block of the pipeline after WATCH commands
1892 are issued. End the transactional block with `execute`.
1893 """
1894 if self.explicit_transaction:
1895 raise RedisError("Cannot issue nested calls to MULTI")
1896 if self.command_stack:
1897 raise RedisError(
1898 "Commands without an initial WATCH have already been issued"
1899 )
1900 self.explicit_transaction = True
1902 def execute_command(
1903 self, *args, **kwargs
1904 ) -> Union["Pipeline", Awaitable["Pipeline"]]:
1905 if (self.watching or args[0] == "WATCH") and not self.explicit_transaction:
1906 return self.immediate_execute_command(*args, **kwargs)
1907 return self.pipeline_execute_command(*args, **kwargs)
1909 async def _disconnect_reset_raise_on_watching(
1910 self,
1911 conn: Connection,
1912 error: Exception,
1913 failure_count: Optional[int] = None,
1914 start_time: Optional[float] = None,
1915 command_name: Optional[str] = None,
1916 ) -> None:
1917 """
1918 Close the connection reset watching state and
1919 raise an exception if we were watching.
1921 The supported exceptions are already checked in the
1922 retry object so we don't need to do it here.
1924 After we disconnect the connection, it will try to reconnect and
1925 do a health check as part of the send_command logic(on connection level).
1926 """
1927 if (
1928 error
1929 and failure_count is not None
1930 and failure_count <= conn.retry.get_retries()
1931 ):
1932 await record_operation_duration(
1933 command_name=command_name,
1934 duration_seconds=time.monotonic() - start_time,
1935 server_address=getattr(conn, "host", None),
1936 server_port=getattr(conn, "port", None),
1937 db_namespace=str(conn.db),
1938 error=error,
1939 retry_attempts=failure_count,
1940 )
1941 await conn.disconnect(error=error, failure_count=failure_count)
1942 # if we were already watching a variable, the watch is no longer
1943 # valid since this connection has died. raise a WatchError, which
1944 # indicates the user should retry this transaction.
1945 if self.watching:
1946 await self.reset()
1947 raise WatchError(
1948 f"A {type(error).__name__} occurred while watching one or more keys"
1949 )
1951 async def immediate_execute_command(self, *args, **options):
1952 """
1953 Execute a command immediately, but don't auto-retry on the supported
1954 errors for retry if we're already WATCHing a variable.
1955 Used when issuing WATCH or subsequent commands retrieving their values but before
1956 MULTI is called.
1957 """
1958 command_name = args[0]
1959 conn = self.connection
1960 # if this is the first call, we need a connection
1961 if not conn:
1962 conn = await self.connection_pool.get_connection()
1963 self.connection = conn
1965 # Start timing for observability
1966 start_time = time.monotonic()
1967 # Track actual retry attempts for error reporting
1968 actual_retry_attempts = 0
1970 def failure_callback(error, failure_count):
1971 nonlocal actual_retry_attempts
1972 actual_retry_attempts = failure_count
1973 return self._disconnect_reset_raise_on_watching(
1974 conn, error, failure_count, start_time, command_name
1975 )
1977 try:
1978 response = await conn.retry.call_with_retry(
1979 lambda: self._send_command_parse_response(
1980 conn, command_name, *args, **options
1981 ),
1982 failure_callback,
1983 with_failure_count=True,
1984 )
1986 await record_operation_duration(
1987 command_name=command_name,
1988 duration_seconds=time.monotonic() - start_time,
1989 server_address=getattr(conn, "host", None),
1990 server_port=getattr(conn, "port", None),
1991 db_namespace=str(conn.db),
1992 )
1994 return response
1995 except Exception as e:
1996 await record_error_count(
1997 server_address=getattr(conn, "host", None),
1998 server_port=getattr(conn, "port", None),
1999 network_peer_address=getattr(conn, "host", None),
2000 network_peer_port=getattr(conn, "port", None),
2001 error_type=e,
2002 retry_attempts=actual_retry_attempts,
2003 is_internal=False,
2004 )
2005 raise
2007 def pipeline_execute_command(self, *args, **options):
2008 """
2009 Stage a command to be executed when execute() is next called
2011 Returns the current Pipeline object back so commands can be
2012 chained together, such as:
2014 pipe = pipe.set('foo', 'bar').incr('baz').decr('bang')
2016 At some other point, you can then run: pipe.execute(),
2017 which will execute all commands queued in the pipe.
2018 """
2019 self.command_stack.append((args, options))
2020 return self
2022 async def _himport_prepare_pipeline(self, conn, commands):
2023 """Delegate to the shared async HIMPORT executor."""
2024 await _himport_exec.prepare_pipeline(self, conn, [args for args, _ in commands])
2026 async def _execute_transaction( # noqa: C901
2027 self, connection: Connection, commands: CommandStackT, raise_on_error
2028 ):
2029 # Ensure fieldsets referenced by buffered HIMPORT SETs are prepared on this
2030 # connection before the MULTI/EXEC block (session state, not transactional).
2031 await self._himport_prepare_pipeline(connection, commands)
2032 pre: CommandT = (("MULTI",), {})
2033 post: CommandT = (("EXEC",), {})
2034 cmds = (pre, *commands, post)
2035 all_cmds = connection.pack_commands(
2036 args for args, options in cmds if EMPTY_RESPONSE not in options
2037 )
2038 await connection.send_packed_command(all_cmds)
2039 errors = []
2041 # parse off the response for MULTI
2042 # NOTE: we need to handle ResponseErrors here and continue
2043 # so that we read all the additional command messages from
2044 # the socket
2045 try:
2046 await self.parse_response(connection, "_")
2047 except ResponseError as err:
2048 errors.append((0, err))
2050 # and all the other commands
2051 for i, command in enumerate(commands):
2052 if EMPTY_RESPONSE in command[1]:
2053 errors.append((i, command[1][EMPTY_RESPONSE]))
2054 else:
2055 try:
2056 await self.parse_response(connection, "_")
2057 except ResponseError as err:
2058 self.annotate_exception(err, i + 1, command[0])
2059 errors.append((i, err))
2061 # parse the EXEC.
2062 try:
2063 response = await self.parse_response(connection, "_")
2064 except ExecAbortError as err:
2065 if errors:
2066 raise errors[0][1] from err
2067 raise
2069 # EXEC clears any watched keys
2070 self.watching = False
2072 if response is None:
2073 raise WatchError("Watched variable changed.") from None
2075 # put any parse errors into the response
2076 for i, e in errors:
2077 response.insert(i, e)
2079 if len(response) != len(commands):
2080 if self.connection:
2081 await self.connection.disconnect()
2082 raise ResponseError(
2083 "Wrong number of response items from pipeline execution"
2084 ) from None
2086 # find any errors in the response and raise if necessary
2087 if raise_on_error:
2088 self.raise_first_error(commands, response)
2090 # We have to run response callbacks manually
2091 data = []
2092 for r, cmd in zip(response, commands):
2093 if not isinstance(r, Exception):
2094 args, options = cmd
2095 command_name = args[0]
2097 # Remove keys entry, it needs only for cache.
2098 options.pop("keys", None)
2100 if command_name in self.response_callbacks:
2101 r = self.response_callbacks[command_name](r, **options)
2102 if inspect.isawaitable(r):
2103 r = await r
2104 data.append(r)
2105 return data
2107 async def _execute_pipeline(
2108 self, connection: Connection, commands: CommandStackT, raise_on_error: bool
2109 ):
2110 # Fold any first-use HIMPORT PREPAREs for referenced fieldsets into the same
2111 # packed write as the queued commands, so a pipeline that lands on a fresh or
2112 # reconnected connection stays a single round trip (the batched write bypasses
2113 # the per-command lazy PREPARE path). Deferred-discard reconciliation happens
2114 # inside pipeline_prepares and only touches the socket when discards are
2115 # actually pending.
2116 fieldsets = await _himport_exec.pipeline_prepares(
2117 self, connection, [args for args, _ in commands]
2118 )
2119 preflight = _himport_exec.prepare_wire_commands(fieldsets)
2120 # build up all commands into a single request to increase network perf
2121 all_cmds = connection.pack_commands(preflight + [args for args, _ in commands])
2122 await connection.send_packed_command(all_cmds)
2124 # Drain the leading PREPARE replies (bookkeeping + capture the first error)
2125 # before the queued replies. Everything on the wire is read before raising so
2126 # the pooled socket never desyncs.
2127 prep_error = await _himport_exec.drain_pipeline_prepares(
2128 self, connection, fieldsets
2129 )
2131 response = []
2132 for args, options in commands:
2133 try:
2134 response.append(
2135 await self.parse_response(connection, args[0], **options)
2136 )
2137 except ResponseError as e:
2138 response.append(e)
2140 # A PREPARE failure (rare: an invalid fieldset definition) is a hard error,
2141 # raised regardless of raise_on_error as it was before folding -- only now
2142 # every reply has already been drained.
2143 if prep_error is not None:
2144 raise prep_error
2145 if raise_on_error:
2146 self.raise_first_error(commands, response)
2147 return response
2149 def raise_first_error(self, commands: CommandStackT, response: Iterable[Any]):
2150 for i, r in enumerate(response):
2151 if isinstance(r, ResponseError):
2152 self.annotate_exception(r, i + 1, commands[i][0])
2153 raise r
2155 def annotate_exception(
2156 self, exception: Exception, number: int, command: Iterable[object]
2157 ) -> None:
2158 cmd = " ".join(map(safe_str, command))
2159 msg = (
2160 f"Command # {number} ({truncate_text(cmd)}) "
2161 f"of pipeline caused error: {exception.args}"
2162 )
2163 exception.args = (msg,) + exception.args[1:]
2165 async def parse_response(
2166 self, connection: Connection, command_name: Union[str, bytes], **options
2167 ):
2168 result = await super().parse_response(connection, command_name, **options)
2169 if command_name in self.UNWATCH_COMMANDS:
2170 self.watching = False
2171 elif command_name == "WATCH":
2172 self.watching = True
2173 return result
2175 async def load_scripts(self):
2176 # make sure all scripts that are about to be run on this pipeline exist
2177 scripts = list(self.scripts)
2178 immediate = self.immediate_execute_command
2179 shas = [s.sha for s in scripts]
2180 # we can't use the normal script_* methods because they would just
2181 # get buffered in the pipeline.
2182 exists = await immediate("SCRIPT EXISTS", *shas)
2183 if not all(exists):
2184 for s, exist in zip(scripts, exists):
2185 if not exist:
2186 s.sha = await immediate("SCRIPT LOAD", s.script)
2188 async def _disconnect_raise_on_watching(
2189 self,
2190 conn: Connection,
2191 error: Exception,
2192 failure_count: Optional[int] = None,
2193 start_time: Optional[float] = None,
2194 command_name: Optional[str] = None,
2195 ):
2196 """
2197 Close the connection, raise an exception if we were watching.
2199 The supported exceptions are already checked in the
2200 retry object so we don't need to do it here.
2202 After we disconnect the connection, it will try to reconnect and
2203 do a health check as part of the send_command logic(on connection level).
2204 """
2205 if (
2206 error
2207 and failure_count is not None
2208 and failure_count <= conn.retry.get_retries()
2209 ):
2210 await record_operation_duration(
2211 command_name=command_name,
2212 duration_seconds=time.monotonic() - start_time,
2213 server_address=getattr(conn, "host", None),
2214 server_port=getattr(conn, "port", None),
2215 db_namespace=str(conn.db),
2216 error=error,
2217 retry_attempts=failure_count,
2218 )
2219 await conn.disconnect(error=error, failure_count=failure_count)
2220 # if we were watching a variable, the watch is no longer valid
2221 # since this connection has died. raise a WatchError, which
2222 # indicates the user should retry this transaction.
2223 if self.watching:
2224 raise WatchError(
2225 f"A {type(error).__name__} occurred while watching one or more keys"
2226 )
2228 async def execute(self, raise_on_error: bool = True) -> List[Any]:
2229 """Execute all the commands in the current pipeline"""
2230 stack = self.command_stack
2231 if not stack and not self.watching:
2232 return []
2233 if self.scripts:
2234 await self.load_scripts()
2235 if self.is_transaction or self.explicit_transaction:
2236 execute = self._execute_transaction
2237 operation_name = "MULTI"
2238 else:
2239 execute = self._execute_pipeline
2240 operation_name = "PIPELINE"
2242 conn = self.connection
2243 if not conn:
2244 conn = await self.connection_pool.get_connection()
2245 # assign to self.connection so reset() releases the connection
2246 # back to the pool after we're done
2247 self.connection = conn
2248 conn = cast(Connection, conn)
2250 # Start timing for observability
2251 start_time = time.monotonic()
2252 # Track actual retry attempts for error reporting
2253 actual_retry_attempts = 0
2255 def failure_callback(error, failure_count):
2256 nonlocal actual_retry_attempts
2257 actual_retry_attempts = failure_count
2258 return self._disconnect_raise_on_watching(
2259 conn, error, failure_count, start_time, operation_name
2260 )
2262 try:
2263 response = await conn.retry.call_with_retry(
2264 lambda: execute(conn, stack, raise_on_error),
2265 failure_callback,
2266 with_failure_count=True,
2267 )
2269 await record_operation_duration(
2270 command_name=operation_name,
2271 duration_seconds=time.monotonic() - start_time,
2272 server_address=getattr(conn, "host", None),
2273 server_port=getattr(conn, "port", None),
2274 db_namespace=str(conn.db),
2275 )
2276 return response
2277 except Exception as e:
2278 await record_error_count(
2279 server_address=getattr(conn, "host", None),
2280 server_port=getattr(conn, "port", None),
2281 network_peer_address=getattr(conn, "host", None),
2282 network_peer_port=getattr(conn, "port", None),
2283 error_type=e,
2284 retry_attempts=actual_retry_attempts,
2285 is_internal=False,
2286 )
2287 raise
2288 finally:
2289 await self.reset()
2291 async def discard(self):
2292 """Flushes all previously queued commands
2293 See: https://redis.io/commands/DISCARD
2294 """
2295 await self.execute_command("DISCARD")
2297 async def watch(self, *names: KeyT):
2298 """Watches the values at keys ``names``"""
2299 if self.explicit_transaction:
2300 raise RedisError("Cannot issue a WATCH after a MULTI")
2301 return await self.execute_command("WATCH", *names)
2303 async def unwatch(self):
2304 """Unwatches all previously specified keys"""
2305 return self.watching and await self.execute_command("UNWATCH") or True