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.connection import (
41 Connection,
42 ConnectionPool,
43 SSLConnection,
44 UnixDomainSocketConnection,
45)
46from redis.asyncio.lock import Lock
47from redis.asyncio.observability.recorder import (
48 record_error_count,
49 record_operation_duration,
50 record_pubsub_message,
51)
52from redis.asyncio.retry import Retry
53from redis.backoff import ExponentialWithJitterBackoff
54from redis.client import (
55 EMPTY_RESPONSE,
56 NEVER_DECODE,
57 AbstractRedis,
58 CaseInsensitiveDict,
59)
60from redis.commands import (
61 AsyncCoreCommands,
62 AsyncRedisModuleCommands,
63 AsyncSentinelCommands,
64 list_or_args,
65)
66from redis.commands.helpers import parse_pubsub_subscriptions, pubsub_subscription_args
67from redis.credentials import CredentialProvider
68from redis.driver_info import DriverInfo, resolve_driver_info
69from redis.event import (
70 AfterPooledConnectionsInstantiationEvent,
71 AfterPubSubConnectionInstantiationEvent,
72 AfterSingleConnectionInstantiationEvent,
73 ClientType,
74 EventDispatcher,
75)
76from redis.exceptions import (
77 ConnectionError,
78 ExecAbortError,
79 PubSubError,
80 RedisError,
81 ResponseError,
82 WatchError,
83)
84from redis.observability.attributes import PubSubDirection
85from redis.typing import ChannelT, EncodableT, KeyT, PubSubHandler, Subscription
86from redis.utils import (
87 SENTINEL,
88 SSL_AVAILABLE,
89 _set_info_logger,
90 deprecated_args,
91 deprecated_function,
92 safe_str,
93 str_if_bytes,
94 truncate_text,
95)
97if TYPE_CHECKING and SSL_AVAILABLE:
98 from ssl import TLSVersion, VerifyFlags, VerifyMode
99else:
100 TLSVersion = None
101 VerifyMode = None
102 VerifyFlags = None
104_KeyT = TypeVar("_KeyT", bound=KeyT)
105_ArgT = TypeVar("_ArgT", KeyT, EncodableT)
106_RedisT = TypeVar("_RedisT", bound="Redis")
107_NormalizeKeysT = TypeVar("_NormalizeKeysT", bound=Mapping[ChannelT, object])
108if TYPE_CHECKING:
109 from redis.asyncio.keyspace_notifications import AsyncKeyspaceNotifications
110 from redis.commands.core import Script
113class ResponseCallbackProtocol(Protocol):
114 def __call__(self, response: Any, **kwargs): ...
117class AsyncResponseCallbackProtocol(Protocol):
118 async def __call__(self, response: Any, **kwargs): ...
121ResponseCallbackT = Union[ResponseCallbackProtocol, AsyncResponseCallbackProtocol]
124class Redis(
125 AbstractRedis, AsyncRedisModuleCommands, AsyncCoreCommands, AsyncSentinelCommands
126):
127 """
128 Implementation of the Redis protocol.
130 This abstract class provides a Python interface to all Redis commands
131 and an implementation of the Redis protocol.
133 Pipelines derive from this, implementing how
134 the commands are sent and received to the Redis server. Based on
135 configuration, an instance will either use a ConnectionPool, or
136 Connection object to talk to redis.
137 """
139 # Type discrimination marker for @overload self-type pattern
140 _is_async_client: Literal[True] = True
142 response_callbacks: MutableMapping[Union[str, bytes], ResponseCallbackT]
144 @classmethod
145 def from_url(
146 cls: Type["Redis"],
147 url: str,
148 single_connection_client: bool = False,
149 auto_close_connection_pool: Optional[bool] = None,
150 **kwargs,
151 ) -> "Redis":
152 """
153 Return a Redis client object configured from the given URL
155 For example::
157 redis://[[username]:[password]]@localhost:6379/0
158 rediss://[[username]:[password]]@localhost:6379/0
159 unix://[username@]/path/to/socket.sock?db=0[&password=password]
161 Three URL schemes are supported:
163 - `redis://` creates a TCP socket connection. See more at:
164 <https://www.iana.org/assignments/uri-schemes/prov/redis>
165 - `rediss://` creates a SSL wrapped TCP socket connection. See more at:
166 <https://www.iana.org/assignments/uri-schemes/prov/rediss>
167 - ``unix://``: creates a Unix Domain Socket connection.
169 The username, password, hostname, path and all querystring values
170 are passed through urllib.parse.unquote in order to replace any
171 percent-encoded values with their corresponding characters.
173 There are several ways to specify a database number. The first value
174 found will be used:
176 1. A ``db`` querystring option, e.g. redis://localhost?db=0
178 2. If using the redis:// or rediss:// schemes, the path argument
179 of the url, e.g. redis://localhost/0
181 3. A ``db`` keyword argument to this function.
183 If none of these options are specified, the default db=0 is used.
185 All querystring options are cast to their appropriate Python types.
186 Boolean arguments can be specified with string values "True"/"False"
187 or "Yes"/"No". Values that cannot be properly cast cause a
188 ``ValueError`` to be raised. Once parsed, the querystring arguments
189 and keyword arguments are passed to the ``ConnectionPool``'s
190 class initializer. In the case of conflicting arguments, querystring
191 arguments always win.
193 """
194 connection_pool = ConnectionPool.from_url(url, **kwargs)
195 client = cls(
196 connection_pool=connection_pool,
197 single_connection_client=single_connection_client,
198 )
199 if auto_close_connection_pool is not None:
200 warnings.warn(
201 DeprecationWarning(
202 '"auto_close_connection_pool" is deprecated '
203 "since version 5.0.1. "
204 "Please create a ConnectionPool explicitly and "
205 "provide to the Redis() constructor instead."
206 )
207 )
208 else:
209 auto_close_connection_pool = True
210 client.auto_close_connection_pool = auto_close_connection_pool
211 return client
213 @classmethod
214 def from_pool(
215 cls: Type["Redis"],
216 connection_pool: ConnectionPool,
217 ) -> "Redis":
218 """
219 Return a Redis client from the given connection pool.
220 The Redis client will take ownership of the connection pool and
221 close it when the Redis client is closed.
223 Because the client closes (disconnects all connections in) the pool
224 when it is closed or garbage-collected, the pool must not be shared
225 with other clients. Constructing multiple clients from the same pool
226 via ``from_pool`` -- for example one per request across tasks -- is
227 not safe: when one client is closed it will disconnect connections
228 still in use by the others.
230 To share a single pool across clients, construct the pool explicitly
231 and manage its lifecycle instead. Unlike ``from_pool``, the plain
232 ``Redis(connection_pool=pool)`` constructor does not take ownership of
233 the pool and will not close it, so a pool created this way can be
234 safely shared across clients. ``ConnectionPool`` supports the async
235 context manager protocol for this::
237 async with ConnectionPool.from_url(url) as pool:
238 r = Redis(connection_pool=pool)
239 """
240 client = cls(
241 connection_pool=connection_pool,
242 )
243 client.auto_close_connection_pool = True
244 return client
246 @deprecated_args(
247 args_to_warn=["retry_on_timeout"],
248 reason="TimeoutError is included by default.",
249 version="6.0.0",
250 )
251 @deprecated_args(
252 args_to_warn=["lib_name", "lib_version"],
253 reason="Use 'driver_info' parameter instead. "
254 "lib_name and lib_version will be removed in a future version.",
255 )
256 def __init__(
257 self,
258 *,
259 host: str = "localhost",
260 port: int = 6379,
261 db: str | int = 0,
262 password: str | None = None,
263 socket_timeout: float | None = DEFAULT_SOCKET_TIMEOUT,
264 socket_connect_timeout: float | None = DEFAULT_SOCKET_CONNECT_TIMEOUT,
265 socket_read_size: int = DEFAULT_SOCKET_READ_SIZE,
266 socket_keepalive: bool | None = True,
267 socket_keepalive_options: Mapping[int, int | bytes] | object | None = SENTINEL,
268 connection_pool: ConnectionPool | None = None,
269 unix_socket_path: str | None = None,
270 encoding: str = "utf-8",
271 encoding_errors: str = "strict",
272 decode_responses: bool = False,
273 retry_on_timeout: bool = False,
274 retry: Retry = Retry(
275 backoff=ExponentialWithJitterBackoff(
276 base=DEFAULT_RETRY_BASE, cap=DEFAULT_RETRY_CAP
277 ),
278 retries=DEFAULT_RETRY_COUNT,
279 ),
280 retry_on_error: list | None = None,
281 ssl: bool = False,
282 ssl_keyfile: str | None = None,
283 ssl_certfile: str | None = None,
284 ssl_cert_reqs: "str | VerifyMode" = "required",
285 ssl_include_verify_flags: List["VerifyFlags"] | None = None,
286 ssl_exclude_verify_flags: List["VerifyFlags"] | None = None,
287 ssl_ca_certs: str | None = None,
288 ssl_ca_data: str | None = None,
289 ssl_ca_path: str | None = None,
290 ssl_check_hostname: bool = True,
291 ssl_min_version: "TLSVersion | None" = None,
292 ssl_ciphers: str | None = None,
293 ssl_password: str | None = None,
294 max_connections: int | None = None,
295 single_connection_client: bool = False,
296 health_check_interval: int = 0,
297 client_name: str | None = None,
298 lib_name: str | object | None = SENTINEL,
299 lib_version: str | object | None = SENTINEL,
300 driver_info: DriverInfo | object | None = SENTINEL,
301 username: str | None = None,
302 auto_close_connection_pool: bool | None = None,
303 redis_connect_func=None,
304 credential_provider: CredentialProvider | None = None,
305 protocol: int | None = None,
306 legacy_responses: bool = True,
307 event_dispatcher: EventDispatcher | None = None,
308 ):
309 """
310 Initialize a new Redis client.
312 To specify a retry policy for specific errors, you have two options:
314 1. Set the `retry_on_error` to a list of the error/s to retry on, and
315 you can also set `retry` to a valid `Retry` object(in case the default
316 one is not appropriate) - with this approach the retries will be triggered
317 on the default errors specified in the Retry object enriched with the
318 errors specified in `retry_on_error`.
320 2. Define a `Retry` object with configured 'supported_errors' and set
321 it to the `retry` parameter - with this approach you completely redefine
322 the errors on which retries will happen.
324 `retry_on_timeout` is deprecated - please include the TimeoutError
325 either in the Retry object or in the `retry_on_error` list.
327 When 'connection_pool' is provided - the retry configuration of the
328 provided pool will be used.
330 Args:
332 socket_keepalive:
333 if `True`, TCP keepalive is enabled for TCP socket connections.
334 Argument is ignored when connection_pool is provided.
335 socket_keepalive_options:
336 mapping of TCP keepalive socket option constants to values, for
337 example `{socket.TCP_KEEPIDLE: 30}`. If left unspecified, redis-py
338 uses TCP keepalive defaults when `socket_keepalive` is enabled:
339 idle 30 seconds, interval 5 seconds, and 3 probes. Platform-specific
340 options that are not available are skipped. Pass `None` or `{}` to
341 avoid setting additional TCP keepalive options. Argument is ignored
342 when connection_pool is provided.
343 """
344 kwargs: Dict[str, Any]
345 if event_dispatcher is None:
346 self._event_dispatcher = EventDispatcher()
347 else:
348 self._event_dispatcher = event_dispatcher
349 # auto_close_connection_pool only has an effect if connection_pool is
350 # None. It is assumed that if connection_pool is not None, the user
351 # wants to manage the connection pool themselves.
352 if auto_close_connection_pool is not None:
353 warnings.warn(
354 DeprecationWarning(
355 '"auto_close_connection_pool" is deprecated '
356 "since version 5.0.1. "
357 "Please create a ConnectionPool explicitly and "
358 "provide to the Redis() constructor instead."
359 )
360 )
361 else:
362 auto_close_connection_pool = True
364 if not connection_pool:
365 # Create internal connection pool, expected to be closed by Redis instance
366 if not retry_on_error:
367 retry_on_error = []
369 # Handle driver_info: if provided, use it; otherwise create from lib_name/lib_version.
370 computed_driver_info = resolve_driver_info(
371 driver_info, lib_name, lib_version
372 )
374 kwargs = {
375 "db": db,
376 "username": username,
377 "password": password,
378 "credential_provider": credential_provider,
379 "socket_timeout": socket_timeout,
380 "socket_read_size": socket_read_size,
381 "encoding": encoding,
382 "encoding_errors": encoding_errors,
383 "decode_responses": decode_responses,
384 "retry_on_error": retry_on_error,
385 "retry": copy.deepcopy(retry),
386 "max_connections": max_connections,
387 "health_check_interval": health_check_interval,
388 "client_name": client_name,
389 "driver_info": computed_driver_info,
390 "redis_connect_func": redis_connect_func,
391 "protocol": protocol,
392 "legacy_responses": legacy_responses,
393 }
394 # based on input, setup appropriate connection args
395 if unix_socket_path is not None:
396 kwargs.update(
397 {
398 "path": unix_socket_path,
399 "connection_class": UnixDomainSocketConnection,
400 }
401 )
402 else:
403 # TCP specific options
404 kwargs.update(
405 {
406 "host": host,
407 "port": port,
408 "socket_connect_timeout": socket_connect_timeout,
409 "socket_keepalive": socket_keepalive,
410 "socket_keepalive_options": socket_keepalive_options,
411 }
412 )
414 if ssl:
415 kwargs.update(
416 {
417 "connection_class": SSLConnection,
418 "ssl_keyfile": ssl_keyfile,
419 "ssl_certfile": ssl_certfile,
420 "ssl_cert_reqs": ssl_cert_reqs,
421 "ssl_include_verify_flags": ssl_include_verify_flags,
422 "ssl_exclude_verify_flags": ssl_exclude_verify_flags,
423 "ssl_ca_certs": ssl_ca_certs,
424 "ssl_ca_data": ssl_ca_data,
425 "ssl_ca_path": ssl_ca_path,
426 "ssl_check_hostname": ssl_check_hostname,
427 "ssl_min_version": ssl_min_version,
428 "ssl_ciphers": ssl_ciphers,
429 "ssl_password": ssl_password,
430 }
431 )
432 # This arg only used if no pool is passed in
433 self.auto_close_connection_pool = auto_close_connection_pool
434 connection_pool = ConnectionPool(**kwargs)
435 self._event_dispatcher.dispatch(
436 AfterPooledConnectionsInstantiationEvent(
437 [connection_pool], ClientType.ASYNC, credential_provider
438 )
439 )
440 else:
441 # If a pool is passed in, do not close it
442 self.auto_close_connection_pool = False
443 self._event_dispatcher.dispatch(
444 AfterPooledConnectionsInstantiationEvent(
445 [connection_pool], ClientType.ASYNC, credential_provider
446 )
447 )
449 self.connection_pool = connection_pool
450 self.single_connection_client = single_connection_client
451 self.connection: Optional[Connection] = None
453 connection_kwargs = self.connection_pool.connection_kwargs
454 self.response_callbacks = CaseInsensitiveDict(
455 get_response_callbacks(
456 user_protocol=connection_kwargs.get("protocol"),
457 legacy_responses=connection_kwargs.get("legacy_responses", True),
458 )
459 )
461 # If using a single connection client, we need to lock creation-of and use-of
462 # the client in order to avoid race conditions such as using asyncio.gather
463 # on a set of redis commands
464 self._single_conn_lock = asyncio.Lock()
466 # When used as an async context manager, we need to increment and decrement
467 # a usage counter so that we can close the connection pool when no one is
468 # using the client.
469 self._usage_counter = 0
470 self._usage_lock = asyncio.Lock()
472 def __repr__(self):
473 return (
474 f"<{self.__class__.__module__}.{self.__class__.__name__}"
475 f"({self.connection_pool!r})>"
476 )
478 def __await__(self):
479 return self.initialize().__await__()
481 async def initialize(self: _RedisT) -> _RedisT:
482 if self.single_connection_client:
483 async with self._single_conn_lock:
484 if self.connection is None:
485 self.connection = await self.connection_pool.get_connection()
487 self._event_dispatcher.dispatch(
488 AfterSingleConnectionInstantiationEvent(
489 self.connection, ClientType.ASYNC, self._single_conn_lock
490 )
491 )
492 return self
494 def set_response_callback(self, command: str, callback: ResponseCallbackT):
495 """Set a custom Response Callback"""
496 self.response_callbacks[command] = callback
498 def get_encoder(self):
499 """Get the connection pool's encoder"""
500 return self.connection_pool.get_encoder()
502 def get_connection_kwargs(self):
503 """Get the connection's key-word arguments"""
504 return self.connection_pool.connection_kwargs
506 def get_retry(self) -> Optional[Retry]:
507 return self.get_connection_kwargs().get("retry")
509 def set_retry(self, retry: Retry) -> None:
510 self.get_connection_kwargs().update({"retry": retry})
511 self.connection_pool.set_retry(retry)
513 def load_external_module(self, funcname, func):
514 """
515 This function can be used to add externally defined redis modules,
516 and their namespaces to the redis client.
518 funcname - A string containing the name of the function to create
519 func - The function, being added to this class.
521 ex: Assume that one has a custom redis module named foomod that
522 creates command named 'foo.dothing' and 'foo.anotherthing' in redis.
523 To load function functions into this namespace:
525 from redis import Redis
526 from foomodule import F
527 r = Redis()
528 r.load_external_module("foo", F)
529 r.foo().dothing('your', 'arguments')
531 For a concrete example see the reimport of the redisjson module in
532 tests/test_connection.py::test_loading_external_modules
533 """
534 setattr(self, funcname, func)
536 def pipeline(
537 self, transaction: bool = True, shard_hint: Optional[str] = None
538 ) -> "Pipeline":
539 """
540 Return a new pipeline object that can queue multiple commands for
541 later execution. ``transaction`` indicates whether all commands
542 should be executed atomically. Apart from making a group of operations
543 atomic, pipelines are useful for reducing the back-and-forth overhead
544 between the client and server.
545 """
546 return Pipeline(
547 self.connection_pool, self.response_callbacks, transaction, shard_hint
548 )
550 async def transaction(
551 self,
552 func: Callable[["Pipeline"], Union[Any, Awaitable[Any]]],
553 *watches: KeyT,
554 shard_hint: Optional[str] = None,
555 value_from_callable: bool = False,
556 watch_delay: Optional[float] = None,
557 ):
558 """
559 Convenience method for executing the callable `func` as a transaction
560 while watching all keys specified in `watches`. The 'func' callable
561 should expect a single argument which is a Pipeline object.
562 """
563 pipe: Pipeline
564 async with self.pipeline(True, shard_hint) as pipe:
565 while True:
566 try:
567 if watches:
568 await pipe.watch(*watches)
569 func_value = func(pipe)
570 if inspect.isawaitable(func_value):
571 func_value = await func_value
572 exec_value = await pipe.execute()
573 return func_value if value_from_callable else exec_value
574 except WatchError:
575 if watch_delay is not None and watch_delay > 0:
576 await asyncio.sleep(watch_delay)
577 continue
579 def lock(
580 self,
581 name: KeyT,
582 timeout: Optional[float] = None,
583 sleep: float = 0.1,
584 blocking: bool = True,
585 blocking_timeout: Optional[float] = None,
586 lock_class: Optional[Type[Lock]] = None,
587 thread_local: bool = True,
588 raise_on_release_error: bool = True,
589 ) -> Lock:
590 """
591 Return a new Lock object using key ``name`` that mimics
592 the behavior of threading.Lock.
594 If specified, ``timeout`` indicates a maximum life for the lock.
595 By default, it will remain locked until release() is called.
597 ``sleep`` indicates the amount of time to sleep per loop iteration
598 when the lock is in blocking mode and another client is currently
599 holding the lock.
601 ``blocking`` indicates whether calling ``acquire`` should block until
602 the lock has been acquired or to fail immediately, causing ``acquire``
603 to return False and the lock not being acquired. Defaults to True.
604 Note this value can be overridden by passing a ``blocking``
605 argument to ``acquire``.
607 ``blocking_timeout`` indicates the maximum amount of time in seconds to
608 spend trying to acquire the lock. A value of ``None`` indicates
609 continue trying forever. ``blocking_timeout`` can be specified as a
610 float or integer, both representing the number of seconds to wait.
612 ``lock_class`` forces the specified lock implementation. Note that as
613 of redis-py 3.0, the only lock class we implement is ``Lock`` (which is
614 a Lua-based lock). So, it's unlikely you'll need this parameter, unless
615 you have created your own custom lock class.
617 ``thread_local`` indicates whether the lock token is placed in
618 thread-local storage. By default, the token is placed in thread local
619 storage so that a thread only sees its token, not a token set by
620 another thread. Consider the following timeline:
622 time: 0, thread-1 acquires `my-lock`, with a timeout of 5 seconds.
623 thread-1 sets the token to "abc"
624 time: 1, thread-2 blocks trying to acquire `my-lock` using the
625 Lock instance.
626 time: 5, thread-1 has not yet completed. redis expires the lock
627 key.
628 time: 5, thread-2 acquired `my-lock` now that it's available.
629 thread-2 sets the token to "xyz"
630 time: 6, thread-1 finishes its work and calls release(). if the
631 token is *not* stored in thread local storage, then
632 thread-1 would see the token value as "xyz" and would be
633 able to successfully release the thread-2's lock.
635 ``raise_on_release_error`` indicates whether to raise an exception when
636 the lock is no longer owned when exiting the context manager. By default,
637 this is True, meaning an exception will be raised. If False, the warning
638 will be logged and the exception will be suppressed.
640 In some use cases it's necessary to disable thread local storage. For
641 example, if you have code where one thread acquires a lock and passes
642 that lock instance to a worker thread to release later. If thread
643 local storage isn't disabled in this case, the worker thread won't see
644 the token set by the thread that acquired the lock. Our assumption
645 is that these cases aren't common and as such default to using
646 thread local storage."""
647 if lock_class is None:
648 lock_class = Lock
649 return lock_class(
650 self,
651 name,
652 timeout=timeout,
653 sleep=sleep,
654 blocking=blocking,
655 blocking_timeout=blocking_timeout,
656 thread_local=thread_local,
657 raise_on_release_error=raise_on_release_error,
658 )
660 def pubsub(self, **kwargs) -> "PubSub":
661 """
662 Return a Publish/Subscribe object. With this object, you can
663 subscribe to channels and listen for messages that get published to
664 them.
665 """
666 return PubSub(
667 self.connection_pool, event_dispatcher=self._event_dispatcher, **kwargs
668 )
670 def keyspace_notifications(
671 self,
672 key_prefix: Union[str, bytes, None] = None,
673 ignore_subscribe_messages: bool = True,
674 ) -> "AsyncKeyspaceNotifications":
675 """
676 Return an :class:`~redis.asyncio.keyspace_notifications.AsyncKeyspaceNotifications`
677 object for subscribing to keyspace and keyevent notifications.
679 Note: Keyspace notifications must be enabled on the Redis server via
680 the ``notify-keyspace-events`` configuration option.
682 Args:
683 key_prefix: Optional prefix to filter and strip from keys in
684 notifications.
685 ignore_subscribe_messages: If True, subscribe/unsubscribe
686 confirmations are not returned by
687 get_message/listen.
688 """
689 from redis.asyncio.keyspace_notifications import AsyncKeyspaceNotifications
691 return AsyncKeyspaceNotifications(
692 self,
693 key_prefix=key_prefix,
694 ignore_subscribe_messages=ignore_subscribe_messages,
695 )
697 def monitor(self) -> "Monitor":
698 return Monitor(self.connection_pool)
700 def client(self) -> "Redis":
701 return self.__class__(
702 connection_pool=self.connection_pool, single_connection_client=True
703 )
705 async def __aenter__(self: _RedisT) -> _RedisT:
706 """
707 Async context manager entry. Increments a usage counter so that the
708 connection pool is only closed (via aclose()) when no context is using
709 the client.
710 """
711 await self._increment_usage()
712 try:
713 # Initialize the client (i.e. establish connection, etc.)
714 return await self.initialize()
715 except Exception:
716 # If initialization fails, decrement the counter to keep it in sync
717 await self._decrement_usage()
718 raise
720 async def _increment_usage(self) -> int:
721 """
722 Helper coroutine to increment the usage counter while holding the lock.
723 Returns the new value of the usage counter.
724 """
725 async with self._usage_lock:
726 self._usage_counter += 1
727 return self._usage_counter
729 async def _decrement_usage(self) -> int:
730 """
731 Helper coroutine to decrement the usage counter while holding the lock.
732 Returns the new value of the usage counter.
733 """
734 async with self._usage_lock:
735 self._usage_counter -= 1
736 return self._usage_counter
738 async def __aexit__(self, exc_type, exc_value, traceback):
739 """
740 Async context manager exit. Decrements a usage counter. If this is the
741 last exit (counter becomes zero), the client closes its connection pool.
742 """
743 current_usage = await asyncio.shield(self._decrement_usage())
744 if current_usage == 0:
745 # This was the last active context, so disconnect the pool.
746 await asyncio.shield(self.aclose())
748 _DEL_MESSAGE = "Unclosed Redis client"
750 # passing _warnings and _grl as argument default since they may be gone
751 # by the time __del__ is called at shutdown
752 def __del__(
753 self,
754 _warn: Any = warnings.warn,
755 _grl: Any = asyncio.get_running_loop,
756 ) -> None:
757 if hasattr(self, "connection") and (self.connection is not None):
758 _warn(f"Unclosed client session {self!r}", ResourceWarning, source=self)
759 try:
760 context = {"client": self, "message": self._DEL_MESSAGE}
761 _grl().call_exception_handler(context)
762 except RuntimeError:
763 pass
764 self.connection._close()
766 async def aclose(self, close_connection_pool: Optional[bool] = None) -> None:
767 """
768 Closes Redis client connection
770 Args:
771 close_connection_pool:
772 decides whether to close the connection pool used by this Redis client,
773 overriding Redis.auto_close_connection_pool.
774 By default, let Redis.auto_close_connection_pool decide
775 whether to close the connection pool.
776 """
777 conn = self.connection
778 if conn:
779 self.connection = None
780 await self.connection_pool.release(conn)
781 if close_connection_pool or (
782 close_connection_pool is None and self.auto_close_connection_pool
783 ):
784 await self.connection_pool.disconnect()
786 @deprecated_function(version="5.0.1", reason="Use aclose() instead", name="close")
787 async def close(self, close_connection_pool: Optional[bool] = None) -> None:
788 """
789 Alias for aclose(), for backwards compatibility
790 """
791 await self.aclose(close_connection_pool)
793 async def _send_command_parse_response(self, conn, command_name, *args, **options):
794 """
795 Send a command and parse the response
796 """
797 await conn.send_command(*args)
798 return await self.parse_response(conn, command_name, **options)
800 async def _close_connection(
801 self,
802 conn: Connection,
803 error: Optional[BaseException] = None,
804 failure_count: Optional[int] = None,
805 start_time: Optional[float] = None,
806 command_name: Optional[str] = None,
807 ):
808 """
809 Close the connection before retrying.
811 The supported exceptions are already checked in the
812 retry object so we don't need to do it here.
814 After we disconnect the connection, it will try to reconnect and
815 do a health check as part of the send_command logic(on connection level).
816 """
817 if (
818 error
819 and failure_count is not None
820 and failure_count <= conn.retry.get_retries()
821 ):
822 await record_operation_duration(
823 command_name=command_name,
824 duration_seconds=time.monotonic() - start_time,
825 server_address=getattr(conn, "host", None),
826 server_port=getattr(conn, "port", None),
827 db_namespace=str(conn.db),
828 error=error,
829 retry_attempts=failure_count,
830 )
832 await conn.disconnect(error=error, failure_count=failure_count)
834 # COMMAND EXECUTION AND PROTOCOL PARSING
835 async def execute_command(self, *args, **options):
836 """Execute a command and return a parsed response"""
837 await self.initialize()
838 pool = self.connection_pool
839 command_name = args[0]
840 conn = self.connection or await pool.get_connection()
842 # Start timing for observability
843 start_time = time.monotonic()
844 # Track actual retry attempts for error reporting
845 actual_retry_attempts = 0
847 def failure_callback(error, failure_count):
848 nonlocal actual_retry_attempts
849 actual_retry_attempts = failure_count
850 return self._close_connection(
851 conn, error, failure_count, start_time, command_name
852 )
854 if self.single_connection_client:
855 await self._single_conn_lock.acquire()
856 try:
857 result = await conn.retry.call_with_retry(
858 lambda: self._send_command_parse_response(
859 conn, command_name, *args, **options
860 ),
861 failure_callback,
862 with_failure_count=True,
863 )
865 await record_operation_duration(
866 command_name=command_name,
867 duration_seconds=time.monotonic() - start_time,
868 server_address=getattr(conn, "host", None),
869 server_port=getattr(conn, "port", None),
870 db_namespace=str(conn.db),
871 )
872 return result
873 except Exception as e:
874 await record_error_count(
875 server_address=getattr(conn, "host", None),
876 server_port=getattr(conn, "port", None),
877 network_peer_address=getattr(conn, "host", None),
878 network_peer_port=getattr(conn, "port", None),
879 error_type=e,
880 retry_attempts=actual_retry_attempts,
881 is_internal=False,
882 )
883 raise
884 finally:
885 if self.single_connection_client:
886 self._single_conn_lock.release()
887 if not self.connection:
888 await pool.release(conn)
890 async def parse_response(
891 self, connection: Connection, command_name: Union[str, bytes], **options
892 ):
893 """Parses a response from the Redis server"""
894 try:
895 if NEVER_DECODE in options:
896 response = await connection.read_response(disable_decoding=True)
897 options.pop(NEVER_DECODE)
898 else:
899 response = await connection.read_response()
900 except ResponseError:
901 if EMPTY_RESPONSE in options:
902 return options[EMPTY_RESPONSE]
903 raise
905 if EMPTY_RESPONSE in options:
906 options.pop(EMPTY_RESPONSE)
908 # Remove keys entry, it needs only for cache.
909 options.pop("keys", None)
911 if command_name in self.response_callbacks:
912 # Mypy bug: https://github.com/python/mypy/issues/10977
913 command_name = cast(str, command_name)
914 retval = self.response_callbacks[command_name](response, **options)
915 return await retval if inspect.isawaitable(retval) else retval
916 return response
919StrictRedis = Redis
922class MonitorCommandInfo(TypedDict):
923 time: float
924 db: int
925 client_address: str
926 client_port: str
927 client_type: str
928 command: str
931class Monitor:
932 """
933 Monitor is useful for handling the MONITOR command to the redis server.
934 next_command() method returns one command from monitor
935 listen() method yields commands from monitor.
936 """
938 monitor_re = re.compile(r"\[(\d+) (.*?)\] (.*)")
939 command_re = re.compile(r'"(.*?)(?<!\\)"')
941 def __init__(self, connection_pool: ConnectionPool):
942 self.connection_pool = connection_pool
943 self.connection: Optional[Connection] = None
945 async def connect(self):
946 if self.connection is None:
947 self.connection = await self.connection_pool.get_connection()
949 async def __aenter__(self):
950 await self.connect()
951 await self.connection.send_command("MONITOR")
952 # check that monitor returns 'OK', but don't return it to user
953 response = await self.connection.read_response()
954 if not bool_ok(response):
955 raise RedisError(f"MONITOR failed: {response}")
956 return self
958 async def __aexit__(self, *args):
959 await self.connection.disconnect()
960 await self.connection_pool.release(self.connection)
962 async def next_command(self) -> MonitorCommandInfo:
963 """Parse the response from a monitor command"""
964 await self.connect()
965 response = await self.connection.read_response()
966 if isinstance(response, bytes):
967 response = self.connection.encoder.decode(response, force=True)
968 command_time, command_data = response.split(" ", 1)
969 m = self.monitor_re.match(command_data)
970 db_id, client_info, command = m.groups()
971 command = " ".join(self.command_re.findall(command))
972 # Redis escapes double quotes because each piece of the command
973 # string is surrounded by double quotes. We don't have that
974 # requirement so remove the escaping and leave the quote.
975 command = command.replace('\\"', '"')
977 if client_info == "lua":
978 client_address = "lua"
979 client_port = ""
980 client_type = "lua"
981 elif client_info.startswith("unix"):
982 client_address = "unix"
983 client_port = client_info[5:]
984 client_type = "unix"
985 else:
986 # use rsplit as ipv6 addresses contain colons
987 client_address, client_port = client_info.rsplit(":", 1)
988 client_type = "tcp"
989 return {
990 "time": float(command_time),
991 "db": int(db_id),
992 "client_address": client_address,
993 "client_port": client_port,
994 "client_type": client_type,
995 "command": command,
996 }
998 async def listen(self) -> AsyncIterator[MonitorCommandInfo]:
999 """Listen for commands coming to the server."""
1000 while True:
1001 yield await self.next_command()
1004class PubSub:
1005 """
1006 PubSub provides publish, subscribe and listen support to Redis channels.
1008 After subscribing to one or more channels, the listen() method will block
1009 until a message arrives on one of the subscribed channels. That message
1010 will be returned and it's safe to start listening again.
1011 """
1013 PUBLISH_MESSAGE_TYPES = ("message", "pmessage", "smessage")
1014 UNSUBSCRIBE_MESSAGE_TYPES = ("unsubscribe", "punsubscribe", "sunsubscribe")
1015 HEALTH_CHECK_MESSAGE = "redis-py-health-check"
1017 def __init__(
1018 self,
1019 connection_pool: ConnectionPool,
1020 shard_hint: Optional[str] = None,
1021 ignore_subscribe_messages: bool = False,
1022 encoder=None,
1023 push_handler_func: Optional[Callable] = None,
1024 event_dispatcher: Optional["EventDispatcher"] = None,
1025 ):
1026 if event_dispatcher is None:
1027 self._event_dispatcher = EventDispatcher()
1028 else:
1029 self._event_dispatcher = event_dispatcher
1030 self.connection_pool = connection_pool
1031 self.shard_hint = shard_hint
1032 self.ignore_subscribe_messages = ignore_subscribe_messages
1033 self.connection = None
1034 # we need to know the encoding options for this connection in order
1035 # to lookup channel and pattern names for callback handlers.
1036 self.encoder = encoder
1037 self.push_handler_func = push_handler_func
1038 if self.encoder is None:
1039 self.encoder = self.connection_pool.get_encoder()
1040 if self.encoder.decode_responses:
1041 self.health_check_response = [
1042 ["pong", self.HEALTH_CHECK_MESSAGE],
1043 self.HEALTH_CHECK_MESSAGE,
1044 ]
1045 else:
1046 self.health_check_response = [
1047 [b"pong", self.encoder.encode(self.HEALTH_CHECK_MESSAGE)],
1048 self.encoder.encode(self.HEALTH_CHECK_MESSAGE),
1049 ]
1050 if self.push_handler_func is None:
1051 _set_info_logger()
1052 self.channels = {}
1053 self.pending_unsubscribe_channels = set()
1054 self.patterns = {}
1055 self.pending_unsubscribe_patterns = set()
1056 self.shard_channels = {}
1057 self.pending_unsubscribe_shard_channels = set()
1058 self._lock = asyncio.Lock()
1060 async def __aenter__(self):
1061 return self
1063 async def __aexit__(self, exc_type, exc_value, traceback):
1064 await self.aclose()
1066 def __del__(self):
1067 if self.connection:
1068 self.connection.deregister_connect_callback(self.on_connect)
1070 async def aclose(self):
1071 # In case a connection property does not yet exist
1072 # (due to a crash earlier in the Redis() constructor), return
1073 # immediately as there is nothing to clean-up.
1074 if not hasattr(self, "connection"):
1075 return
1076 async with self._lock:
1077 if self.connection:
1078 # Use nowait=True to avoid awaiting StreamWriter.wait_closed(),
1079 # which can deadlock when a concurrent reader task (e.g. one
1080 # running pubsub.run() or get_message(block=True)) still holds
1081 # the transport. See https://github.com/redis/redis-py/issues/3941
1082 await self.connection.disconnect(nowait=True)
1083 self.connection.deregister_connect_callback(self.on_connect)
1084 await self.connection_pool.release(self.connection)
1085 self.connection = None
1086 self.channels = {}
1087 self.pending_unsubscribe_channels = set()
1088 self.patterns = {}
1089 self.pending_unsubscribe_patterns = set()
1090 self.shard_channels = {}
1091 self.pending_unsubscribe_shard_channels = set()
1093 @deprecated_function(version="5.0.1", reason="Use aclose() instead", name="close")
1094 async def close(self) -> None:
1095 """Alias for aclose(), for backwards compatibility"""
1096 await self.aclose()
1098 @deprecated_function(version="5.0.1", reason="Use aclose() instead", name="reset")
1099 async def reset(self) -> None:
1100 """Alias for aclose(), for backwards compatibility"""
1101 await self.aclose()
1103 async def _resubscribe(self, subscribed, subscribe_fn) -> None:
1104 # Replay handler-backed subscriptions as positional Subscription objects
1105 # so binary names never need to be decoded into keyword argument keys.
1106 subscriptions = pubsub_subscription_args(subscribed)
1107 if subscriptions:
1108 await subscribe_fn(*subscriptions)
1110 async def _resubscribe_shard_channels(self) -> None:
1111 await self._resubscribe(self.shard_channels, self.ssubscribe)
1113 async def on_connect(self, connection: Connection):
1114 """Re-subscribe to any channels and patterns previously subscribed to"""
1115 self.pending_unsubscribe_channels.clear()
1116 self.pending_unsubscribe_patterns.clear()
1117 self.pending_unsubscribe_shard_channels.clear()
1118 if self.channels:
1119 await self._resubscribe(self.channels, self.subscribe)
1120 if self.patterns:
1121 await self._resubscribe(self.patterns, self.psubscribe)
1122 if self.shard_channels:
1123 await self._resubscribe_shard_channels()
1125 @property
1126 def subscribed(self):
1127 """Indicates if there are subscriptions to any channels or patterns"""
1128 return bool(self.channels or self.patterns or self.shard_channels)
1130 async def execute_command(self, *args: EncodableT):
1131 """Execute a publish/subscribe command"""
1133 # NOTE: don't parse the response in this function -- it could pull a
1134 # legitimate message off the stack if the connection is already
1135 # subscribed to one or more channels
1137 await self.connect()
1138 connection = self.connection
1139 kwargs = {"check_health": not self.subscribed}
1140 await self._execute(connection, connection.send_command, *args, **kwargs)
1142 async def connect(self):
1143 """
1144 Ensure that the PubSub is connected
1145 """
1146 if self.connection is None:
1147 self.connection = await self.connection_pool.get_connection()
1148 # register a callback that re-subscribes to any channels we
1149 # were listening to when we were disconnected
1150 self.connection.register_connect_callback(self.on_connect)
1151 else:
1152 await self.connection.connect()
1153 if self.push_handler_func is not None:
1154 self.connection._parser.set_pubsub_push_handler(self.push_handler_func)
1156 self._event_dispatcher.dispatch(
1157 AfterPubSubConnectionInstantiationEvent(
1158 self.connection, self.connection_pool, ClientType.ASYNC, self._lock
1159 )
1160 )
1162 async def _reconnect(
1163 self,
1164 conn,
1165 error: Optional[BaseException] = None,
1166 failure_count: Optional[int] = None,
1167 start_time: Optional[float] = None,
1168 command_name: Optional[str] = None,
1169 ):
1170 """
1171 The supported exceptions are already checked in the
1172 retry object so we don't need to do it here.
1174 In this error handler we are trying to reconnect to the server.
1175 """
1176 if (
1177 error
1178 and failure_count is not None
1179 and failure_count <= conn.retry.get_retries()
1180 ):
1181 if command_name:
1182 await record_operation_duration(
1183 command_name=command_name,
1184 duration_seconds=time.monotonic() - start_time,
1185 server_address=getattr(conn, "host", None),
1186 server_port=getattr(conn, "port", None),
1187 db_namespace=str(conn.db),
1188 error=error,
1189 retry_attempts=failure_count,
1190 )
1191 await conn.disconnect(error=error, failure_count=failure_count)
1192 await conn.connect()
1194 async def _execute(self, conn, command, *args, **kwargs):
1195 """
1196 Connect manually upon disconnection. If the Redis server is down,
1197 this will fail and raise a ConnectionError as desired.
1198 After reconnection, the ``on_connect`` callback should have been
1199 called by the # connection to resubscribe us to any channels and
1200 patterns we were previously listening to
1201 """
1202 if not len(args) == 0:
1203 command_name = args[0]
1204 else:
1205 command_name = None
1207 # Start timing for observability
1208 start_time = time.monotonic()
1209 # Track actual retry attempts for error reporting
1210 actual_retry_attempts = 0
1212 def failure_callback(error, failure_count):
1213 nonlocal actual_retry_attempts
1214 actual_retry_attempts = failure_count
1215 return self._reconnect(conn, error, failure_count, start_time, command_name)
1217 try:
1218 response = await conn.retry.call_with_retry(
1219 lambda: command(*args, **kwargs),
1220 failure_callback,
1221 with_failure_count=True,
1222 )
1224 if command_name:
1225 await record_operation_duration(
1226 command_name=command_name,
1227 duration_seconds=time.monotonic() - start_time,
1228 server_address=getattr(conn, "host", None),
1229 server_port=getattr(conn, "port", None),
1230 db_namespace=str(conn.db),
1231 )
1233 return response
1234 except Exception as e:
1235 await record_error_count(
1236 server_address=getattr(conn, "host", None),
1237 server_port=getattr(conn, "port", None),
1238 network_peer_address=getattr(conn, "host", None),
1239 network_peer_port=getattr(conn, "port", None),
1240 error_type=e,
1241 retry_attempts=actual_retry_attempts,
1242 is_internal=False,
1243 )
1244 raise
1246 async def parse_response(self, block: bool = True, timeout: float = 0):
1247 """
1248 Parse the response from a publish/subscribe command.
1250 Args:
1251 block: If True, block indefinitely until a message is available.
1252 If False, return immediately if no message is available.
1253 Default: True
1254 timeout: The timeout in seconds for reading a response when block=False.
1255 This parameter is ignored when block=True.
1256 Default: 0 (return immediately if no data available)
1258 Returns:
1259 The parsed response from the server, or None if no message is available
1260 within the timeout period (when block=False).
1262 Important:
1263 The block and timeout parameters work together:
1264 - When block=True: timeout is IGNORED, method blocks indefinitely
1265 - When block=False: timeout is USED, method returns after timeout expires
1267 Typically, you should use get_message(timeout=X) instead of calling
1268 parse_response() directly. The get_message() method automatically sets
1269 block=False when a timeout is provided, and block=True when timeout=None.
1271 Example:
1272 # Block indefinitely (timeout is ignored)
1273 response = await pubsub.parse_response(block=True, timeout=0.1)
1275 # Non-blocking with 0.1 second timeout
1276 response = await pubsub.parse_response(block=False, timeout=0.1)
1278 # Non-blocking, return immediately
1279 response = await pubsub.parse_response(block=False, timeout=0)
1281 # Recommended: use get_message() instead
1282 msg = await pubsub.get_message(timeout=0.1) # automatically sets block=False
1283 msg = await pubsub.get_message(timeout=None) # automatically sets block=True
1284 """
1285 conn = self.connection
1286 if conn is None:
1287 raise RuntimeError(
1288 "pubsub connection not set: "
1289 "did you forget to call subscribe() or psubscribe()?"
1290 )
1292 await self.check_health()
1294 if not conn.is_connected:
1295 await conn.connect()
1297 # Block=True: signal "no timeout" to conn.read_response via
1298 # math.inf. The connection treats math.inf as the per-read
1299 # opt-in for blocking indefinitely without falling back to
1300 # self.socket_timeout. Reconnect/AUTH/HELLO/resubscribe
1301 # operations performed by the retry layer continue to honor
1302 # self.socket_timeout because they do not pass math.inf.
1303 #
1304 # TODO(next-major): when the async Connection.read_response
1305 # default for ``timeout`` is changed to SENTINEL, passing
1306 # ``timeout=None`` from this method will become the natural
1307 # "no timeout" signal and the math.inf hand-off can be
1308 # removed. That swap is a breaking change to the
1309 # Connection.read_response signature so it must wait for a
1310 # major release.
1311 read_timeout = math.inf if block else timeout
1312 response = await self._execute(
1313 conn,
1314 conn.read_response,
1315 timeout=read_timeout,
1316 disconnect_on_error=False,
1317 push_request=True,
1318 )
1320 if conn.health_check_interval and response in self.health_check_response:
1321 # ignore the health check message as user might not expect it
1322 return None
1323 return response
1325 async def check_health(self):
1326 conn = self.connection
1327 if conn is None:
1328 raise RuntimeError(
1329 "pubsub connection not set: "
1330 "did you forget to call subscribe() or psubscribe()?"
1331 )
1333 if (
1334 conn.health_check_interval
1335 and asyncio.get_running_loop().time() > conn.next_health_check
1336 ):
1337 await conn.send_command(
1338 "PING", self.HEALTH_CHECK_MESSAGE, check_health=False
1339 )
1341 def _normalize_keys(self, data: _NormalizeKeysT) -> _NormalizeKeysT:
1342 """
1343 normalize channel/pattern names to be either bytes or strings
1344 based on whether responses are automatically decoded. this saves us
1345 from coercing the value for each message coming in.
1346 """
1347 encode = self.encoder.encode
1348 decode = self.encoder.decode
1349 return {decode(encode(k)): v for k, v in data.items()} # type: ignore[return-value] # noqa: E501
1351 async def psubscribe(
1352 self, *args: ChannelT | Subscription, **kwargs: PubSubHandler
1353 ) -> None:
1354 """
1355 Subscribe to channel patterns.
1356 Patterns supplied as keyword arguments expect a pattern name as the
1357 key and a callable as the value.
1358 ``Subscription`` objects can also be supplied positionally with an
1359 optional handler.
1360 A pattern's callable will be invoked automatically
1361 when a message is received on that pattern rather than producing a
1362 message via ``listen()``.
1363 """
1364 new_patterns = parse_pubsub_subscriptions(args, kwargs)
1365 ret_val = await self.execute_command("PSUBSCRIBE", *new_patterns.keys())
1366 # update the patterns dict AFTER we send the command. we don't want to
1367 # subscribe twice to these patterns, once for the command and again
1368 # for the reconnection.
1369 new_patterns = self._normalize_keys(new_patterns)
1370 self.patterns.update(new_patterns)
1371 self.pending_unsubscribe_patterns.difference_update(new_patterns)
1372 return ret_val
1374 def punsubscribe(self, *args: ChannelT) -> Awaitable:
1375 """
1376 Unsubscribe from the supplied patterns. If empty, unsubscribe from
1377 all patterns.
1378 """
1379 patterns: Iterable[ChannelT]
1380 if args:
1381 parsed_args = list_or_args((args[0],), args[1:])
1382 patterns = self._normalize_keys(dict.fromkeys(parsed_args)).keys()
1383 else:
1384 parsed_args = []
1385 patterns = self.patterns
1386 self.pending_unsubscribe_patterns.update(patterns)
1387 return self.execute_command("PUNSUBSCRIBE", *parsed_args)
1389 async def subscribe(
1390 self, *args: ChannelT | Subscription, **kwargs: PubSubHandler
1391 ) -> None:
1392 """
1393 Subscribe to channels.
1394 Channels supplied as keyword arguments expect
1395 a channel name as the key and a callable as the value.
1396 ``Subscription`` objects can also be supplied positionally with an
1397 optional handler.
1398 A channel's callable will be invoked automatically
1399 when a message is received on that channel rather than producing a
1400 message via ``listen()`` or ``get_message()``.
1401 """
1402 new_channels = parse_pubsub_subscriptions(args, kwargs)
1403 ret_val = await self.execute_command("SUBSCRIBE", *new_channels.keys())
1404 # update the channels dict AFTER we send the command. we don't want to
1405 # subscribe twice to these channels, once for the command and again
1406 # for the reconnection.
1407 new_channels = self._normalize_keys(new_channels)
1408 self.channels.update(new_channels)
1409 self.pending_unsubscribe_channels.difference_update(new_channels)
1410 return ret_val
1412 def unsubscribe(self, *args) -> Awaitable:
1413 """
1414 Unsubscribe from the supplied channels. If empty, unsubscribe from
1415 all channels
1416 """
1417 if args:
1418 parsed_args = list_or_args(args[0], args[1:])
1419 channels = self._normalize_keys(dict.fromkeys(parsed_args))
1420 else:
1421 parsed_args = []
1422 channels = self.channels
1423 self.pending_unsubscribe_channels.update(channels)
1424 return self.execute_command("UNSUBSCRIBE", *parsed_args)
1426 async def ssubscribe(
1427 self,
1428 *args: ChannelT | Subscription,
1429 target_node: Any = None,
1430 **kwargs: PubSubHandler,
1431 ) -> None:
1432 """
1433 Subscribes the client to the specified shard channels.
1434 Channels supplied as keyword arguments expect a channel name as the key
1435 and a callable as the value.
1436 ``Subscription`` objects can also be supplied positionally
1437 with an optional handler.
1438 A channel's callable will be invoked automatically when a message
1439 is received on that channel rather than producing a message
1440 via ``listen()`` or ``get_sharded_message()``.
1441 """
1442 new_s_channels = parse_pubsub_subscriptions(args, kwargs)
1443 ret_val = await self.execute_command("SSUBSCRIBE", *new_s_channels.keys())
1444 # update the s_channels dict AFTER we send the command. we don't want to
1445 # subscribe twice to these channels, once for the command and again
1446 # for the reconnection.
1447 new_s_channels = self._normalize_keys(new_s_channels)
1448 self.shard_channels.update(new_s_channels)
1449 self.pending_unsubscribe_shard_channels.difference_update(new_s_channels)
1450 return ret_val
1452 def sunsubscribe(self, *args, target_node=None) -> Awaitable:
1453 """
1454 Unsubscribe from the supplied shard_channels. If empty, unsubscribe from
1455 all shard_channels
1456 """
1457 if args:
1458 args = list_or_args(args[0], args[1:])
1459 s_channels = self._normalize_keys(dict.fromkeys(args))
1460 else:
1461 s_channels = self.shard_channels
1462 self.pending_unsubscribe_shard_channels.update(s_channels)
1463 return self.execute_command("SUNSUBSCRIBE", *args)
1465 async def listen(self) -> AsyncIterator:
1466 """Listen for messages on channels this client has been subscribed to"""
1467 while self.subscribed:
1468 response = await self.handle_message(await self.parse_response(block=True))
1469 if response is not None:
1470 yield response
1472 async def get_message(
1473 self, ignore_subscribe_messages: bool = False, timeout: Optional[float] = 0.0
1474 ):
1475 """
1476 Get the next message if one is available, otherwise None.
1478 If timeout is specified, the system will wait for `timeout` seconds
1479 before returning. Timeout should be specified as a floating point
1480 number or None to wait indefinitely.
1481 """
1482 response = await self.parse_response(block=(timeout is None), timeout=timeout)
1483 if response:
1484 return await self.handle_message(response, ignore_subscribe_messages)
1485 return None
1487 def ping(self, message=None) -> Awaitable[bool]:
1488 """
1489 Ping the Redis server to test connectivity.
1491 Sends a PING command to the Redis server and returns True if the server
1492 responds with "PONG".
1493 """
1494 args = ["PING", message] if message is not None else ["PING"]
1495 return self.execute_command(*args)
1497 async def handle_message(self, response, ignore_subscribe_messages=False):
1498 """
1499 Parses a pub/sub message. If the channel or pattern was subscribed to
1500 with a message handler, the handler is invoked instead of a parsed
1501 message being returned.
1502 """
1503 if response is None:
1504 return None
1505 if isinstance(response, bytes):
1506 response = [b"pong", response] if response != b"PONG" else [b"pong", b""]
1507 message_type = str_if_bytes(response[0])
1508 if message_type == "pmessage":
1509 message = {
1510 "type": message_type,
1511 "pattern": response[1],
1512 "channel": response[2],
1513 "data": response[3],
1514 }
1515 elif message_type == "pong":
1516 message = {
1517 "type": message_type,
1518 "pattern": None,
1519 "channel": None,
1520 "data": response[1],
1521 }
1522 else:
1523 message = {
1524 "type": message_type,
1525 "pattern": None,
1526 "channel": response[1],
1527 "data": response[2],
1528 }
1530 if message_type in ["message", "pmessage"]:
1531 channel = str_if_bytes(message["channel"])
1532 await record_pubsub_message(
1533 direction=PubSubDirection.RECEIVE,
1534 channel=channel,
1535 )
1536 elif message_type == "smessage":
1537 channel = str_if_bytes(message["channel"])
1538 await record_pubsub_message(
1539 direction=PubSubDirection.RECEIVE,
1540 channel=channel,
1541 sharded=True,
1542 )
1544 # if this is an unsubscribe message, remove it from memory
1545 if message_type in self.UNSUBSCRIBE_MESSAGE_TYPES:
1546 if message_type == "punsubscribe":
1547 pattern = response[1]
1548 if pattern in self.pending_unsubscribe_patterns:
1549 self.pending_unsubscribe_patterns.remove(pattern)
1550 self.patterns.pop(pattern, None)
1551 elif message_type == "sunsubscribe":
1552 s_channel = response[1]
1553 if s_channel in self.pending_unsubscribe_shard_channels:
1554 self.pending_unsubscribe_shard_channels.remove(s_channel)
1555 self.shard_channels.pop(s_channel, None)
1556 else:
1557 channel = response[1]
1558 if channel in self.pending_unsubscribe_channels:
1559 self.pending_unsubscribe_channels.remove(channel)
1560 self.channels.pop(channel, None)
1562 if message_type in self.PUBLISH_MESSAGE_TYPES:
1563 # if there's a message handler, invoke it
1564 if message_type == "pmessage":
1565 handler = self.patterns.get(message["pattern"], None)
1566 elif message_type == "smessage":
1567 handler = self.shard_channels.get(message["channel"], None)
1568 else:
1569 handler = self.channels.get(message["channel"], None)
1570 if handler:
1571 if inspect.iscoroutinefunction(handler):
1572 await handler(message)
1573 else:
1574 handler(message)
1575 return None
1576 elif message_type != "pong":
1577 # this is a subscribe/unsubscribe message. ignore if we don't
1578 # want them
1579 if ignore_subscribe_messages or self.ignore_subscribe_messages:
1580 return None
1582 return message
1584 async def run(
1585 self,
1586 *,
1587 exception_handler: Optional["PSWorkerThreadExcHandlerT"] = None,
1588 poll_timeout: float = 1.0,
1589 pubsub=None,
1590 ) -> None:
1591 """Process pub/sub messages using registered callbacks.
1593 This is the equivalent of :py:meth:`redis.PubSub.run_in_thread` in
1594 redis-py, but it is a coroutine. To launch it as a separate task, use
1595 ``asyncio.create_task``:
1597 >>> task = asyncio.create_task(pubsub.run())
1599 To shut it down, use asyncio cancellation:
1601 >>> task.cancel()
1602 >>> await task
1603 """
1604 for channel, handler in self.channels.items():
1605 if handler is None:
1606 raise PubSubError(f"Channel: '{channel}' has no handler registered")
1607 for pattern, handler in self.patterns.items():
1608 if handler is None:
1609 raise PubSubError(f"Pattern: '{pattern}' has no handler registered")
1611 await self.connect()
1612 while True:
1613 try:
1614 if pubsub is None:
1615 await self.get_message(
1616 ignore_subscribe_messages=True, timeout=poll_timeout
1617 )
1618 else:
1619 await pubsub.get_message(
1620 ignore_subscribe_messages=True, timeout=poll_timeout
1621 )
1622 except asyncio.CancelledError:
1623 raise
1624 except BaseException as e:
1625 if exception_handler is None:
1626 raise
1627 res = exception_handler(e, self)
1628 if inspect.isawaitable(res):
1629 await res
1630 # Ensure that other tasks on the event loop get a chance to run
1631 # if we didn't have to block for I/O anywhere.
1632 await asyncio.sleep(0)
1635class PubsubWorkerExceptionHandler(Protocol):
1636 def __call__(self, e: BaseException, pubsub: PubSub): ...
1639class AsyncPubsubWorkerExceptionHandler(Protocol):
1640 async def __call__(self, e: BaseException, pubsub: PubSub): ...
1643PSWorkerThreadExcHandlerT = Union[
1644 PubsubWorkerExceptionHandler, AsyncPubsubWorkerExceptionHandler
1645]
1648CommandT = Tuple[Tuple[Union[str, bytes], ...], Mapping[str, Any]]
1649CommandStackT = List[CommandT]
1652class Pipeline(Redis): # lgtm [py/init-calls-subclass]
1653 """
1654 Pipelines provide a way to transmit multiple commands to the Redis server
1655 in one transmission. This is convenient for batch processing, such as
1656 saving all the values in a list to Redis.
1658 All commands executed within a pipeline(when running in transactional mode,
1659 which is the default behavior) are wrapped with MULTI and EXEC
1660 calls. This guarantees all commands executed in the pipeline will be
1661 executed atomically.
1663 Any command raising an exception does *not* halt the execution of
1664 subsequent commands in the pipeline. Instead, the exception is caught
1665 and its instance is placed into the response list returned by execute().
1666 Code iterating over the response list should be able to deal with an
1667 instance of an exception as a potential value. In general, these will be
1668 ResponseError exceptions, such as those raised when issuing a command
1669 on a key of a different datatype.
1670 """
1672 UNWATCH_COMMANDS = {"DISCARD", "EXEC", "UNWATCH"}
1674 def __init__(
1675 self,
1676 connection_pool: ConnectionPool,
1677 response_callbacks: MutableMapping[Union[str, bytes], ResponseCallbackT],
1678 transaction: bool,
1679 shard_hint: Optional[str],
1680 ):
1681 self.connection_pool = connection_pool
1682 self.connection = None
1683 self.response_callbacks = response_callbacks
1684 self.is_transaction = transaction
1685 self.shard_hint = shard_hint
1686 self.watching = False
1687 self.command_stack: CommandStackT = []
1688 self.scripts: Set[Script] = set()
1689 self.explicit_transaction = False
1691 async def __aenter__(self: _RedisT) -> _RedisT:
1692 return self
1694 async def __aexit__(self, exc_type, exc_value, traceback):
1695 await self.reset()
1697 def __await__(self):
1698 return self._async_self().__await__()
1700 _DEL_MESSAGE = "Unclosed Pipeline client"
1702 def __len__(self):
1703 return len(self.command_stack)
1705 def __bool__(self):
1706 """Pipeline instances should always evaluate to True"""
1707 return True
1709 async def _async_self(self):
1710 return self
1712 async def reset(self):
1713 self.command_stack = []
1714 self.scripts = set()
1715 try:
1716 # make sure to reset the connection state in the event that we were
1717 # watching something
1718 if self.watching and self.connection:
1719 try:
1720 # call this manually since our unwatch or
1721 # immediate_execute_command methods can call reset()
1722 await self.connection.send_command("UNWATCH")
1723 await self.connection.read_response()
1724 except ConnectionError:
1725 # disconnect will also remove any previous WATCHes
1726 if self.connection:
1727 await self.connection.disconnect()
1728 except asyncio.CancelledError:
1729 # Disconnect so any unread UNWATCH reply does not get
1730 # served to the next caller that takes the connection.
1731 if self.connection:
1732 await self.connection.disconnect()
1733 raise
1734 finally:
1735 self.watching = False
1736 self.explicit_transaction = False
1737 # We can safely return the connection to the pool here since we're
1738 # sure we're no longer WATCHing anything. Detach self.connection
1739 # before awaiting release: if a second cancel aborts the await,
1740 # the pipeline must not be left holding a reference to a
1741 # connection that is being returned to the pool. Shield the
1742 # release itself so a second cancel cannot split the pool's
1743 # internal in-use/available bookkeeping mid-update.
1744 if self.connection:
1745 connection, self.connection = self.connection, None
1746 await asyncio.shield(self.connection_pool.release(connection))
1748 async def aclose(self) -> None:
1749 """Alias for reset(), a standard method name for cleanup"""
1750 await self.reset()
1752 def multi(self):
1753 """
1754 Start a transactional block of the pipeline after WATCH commands
1755 are issued. End the transactional block with `execute`.
1756 """
1757 if self.explicit_transaction:
1758 raise RedisError("Cannot issue nested calls to MULTI")
1759 if self.command_stack:
1760 raise RedisError(
1761 "Commands without an initial WATCH have already been issued"
1762 )
1763 self.explicit_transaction = True
1765 def execute_command(
1766 self, *args, **kwargs
1767 ) -> Union["Pipeline", Awaitable["Pipeline"]]:
1768 if (self.watching or args[0] == "WATCH") and not self.explicit_transaction:
1769 return self.immediate_execute_command(*args, **kwargs)
1770 return self.pipeline_execute_command(*args, **kwargs)
1772 async def _disconnect_reset_raise_on_watching(
1773 self,
1774 conn: Connection,
1775 error: Exception,
1776 failure_count: Optional[int] = None,
1777 start_time: Optional[float] = None,
1778 command_name: Optional[str] = None,
1779 ) -> None:
1780 """
1781 Close the connection reset watching state and
1782 raise an exception if we were watching.
1784 The supported exceptions are already checked in the
1785 retry object so we don't need to do it here.
1787 After we disconnect the connection, it will try to reconnect and
1788 do a health check as part of the send_command logic(on connection level).
1789 """
1790 if (
1791 error
1792 and failure_count is not None
1793 and failure_count <= conn.retry.get_retries()
1794 ):
1795 await record_operation_duration(
1796 command_name=command_name,
1797 duration_seconds=time.monotonic() - start_time,
1798 server_address=getattr(conn, "host", None),
1799 server_port=getattr(conn, "port", None),
1800 db_namespace=str(conn.db),
1801 error=error,
1802 retry_attempts=failure_count,
1803 )
1804 await conn.disconnect(error=error, failure_count=failure_count)
1805 # if we were already watching a variable, the watch is no longer
1806 # valid since this connection has died. raise a WatchError, which
1807 # indicates the user should retry this transaction.
1808 if self.watching:
1809 await self.reset()
1810 raise WatchError(
1811 f"A {type(error).__name__} occurred while watching one or more keys"
1812 )
1814 async def immediate_execute_command(self, *args, **options):
1815 """
1816 Execute a command immediately, but don't auto-retry on the supported
1817 errors for retry if we're already WATCHing a variable.
1818 Used when issuing WATCH or subsequent commands retrieving their values but before
1819 MULTI is called.
1820 """
1821 command_name = args[0]
1822 conn = self.connection
1823 # if this is the first call, we need a connection
1824 if not conn:
1825 conn = await self.connection_pool.get_connection()
1826 self.connection = conn
1828 # Start timing for observability
1829 start_time = time.monotonic()
1830 # Track actual retry attempts for error reporting
1831 actual_retry_attempts = 0
1833 def failure_callback(error, failure_count):
1834 nonlocal actual_retry_attempts
1835 actual_retry_attempts = failure_count
1836 return self._disconnect_reset_raise_on_watching(
1837 conn, error, failure_count, start_time, command_name
1838 )
1840 try:
1841 response = await conn.retry.call_with_retry(
1842 lambda: self._send_command_parse_response(
1843 conn, command_name, *args, **options
1844 ),
1845 failure_callback,
1846 with_failure_count=True,
1847 )
1849 await record_operation_duration(
1850 command_name=command_name,
1851 duration_seconds=time.monotonic() - start_time,
1852 server_address=getattr(conn, "host", None),
1853 server_port=getattr(conn, "port", None),
1854 db_namespace=str(conn.db),
1855 )
1857 return response
1858 except Exception as e:
1859 await record_error_count(
1860 server_address=getattr(conn, "host", None),
1861 server_port=getattr(conn, "port", None),
1862 network_peer_address=getattr(conn, "host", None),
1863 network_peer_port=getattr(conn, "port", None),
1864 error_type=e,
1865 retry_attempts=actual_retry_attempts,
1866 is_internal=False,
1867 )
1868 raise
1870 def pipeline_execute_command(self, *args, **options):
1871 """
1872 Stage a command to be executed when execute() is next called
1874 Returns the current Pipeline object back so commands can be
1875 chained together, such as:
1877 pipe = pipe.set('foo', 'bar').incr('baz').decr('bang')
1879 At some other point, you can then run: pipe.execute(),
1880 which will execute all commands queued in the pipe.
1881 """
1882 self.command_stack.append((args, options))
1883 return self
1885 async def _execute_transaction( # noqa: C901
1886 self, connection: Connection, commands: CommandStackT, raise_on_error
1887 ):
1888 pre: CommandT = (("MULTI",), {})
1889 post: CommandT = (("EXEC",), {})
1890 cmds = (pre, *commands, post)
1891 all_cmds = connection.pack_commands(
1892 args for args, options in cmds if EMPTY_RESPONSE not in options
1893 )
1894 await connection.send_packed_command(all_cmds)
1895 errors = []
1897 # parse off the response for MULTI
1898 # NOTE: we need to handle ResponseErrors here and continue
1899 # so that we read all the additional command messages from
1900 # the socket
1901 try:
1902 await self.parse_response(connection, "_")
1903 except ResponseError as err:
1904 errors.append((0, err))
1906 # and all the other commands
1907 for i, command in enumerate(commands):
1908 if EMPTY_RESPONSE in command[1]:
1909 errors.append((i, command[1][EMPTY_RESPONSE]))
1910 else:
1911 try:
1912 await self.parse_response(connection, "_")
1913 except ResponseError as err:
1914 self.annotate_exception(err, i + 1, command[0])
1915 errors.append((i, err))
1917 # parse the EXEC.
1918 try:
1919 response = await self.parse_response(connection, "_")
1920 except ExecAbortError as err:
1921 if errors:
1922 raise errors[0][1] from err
1923 raise
1925 # EXEC clears any watched keys
1926 self.watching = False
1928 if response is None:
1929 raise WatchError("Watched variable changed.") from None
1931 # put any parse errors into the response
1932 for i, e in errors:
1933 response.insert(i, e)
1935 if len(response) != len(commands):
1936 if self.connection:
1937 await self.connection.disconnect()
1938 raise ResponseError(
1939 "Wrong number of response items from pipeline execution"
1940 ) from None
1942 # find any errors in the response and raise if necessary
1943 if raise_on_error:
1944 self.raise_first_error(commands, response)
1946 # We have to run response callbacks manually
1947 data = []
1948 for r, cmd in zip(response, commands):
1949 if not isinstance(r, Exception):
1950 args, options = cmd
1951 command_name = args[0]
1953 # Remove keys entry, it needs only for cache.
1954 options.pop("keys", None)
1956 if command_name in self.response_callbacks:
1957 r = self.response_callbacks[command_name](r, **options)
1958 if inspect.isawaitable(r):
1959 r = await r
1960 data.append(r)
1961 return data
1963 async def _execute_pipeline(
1964 self, connection: Connection, commands: CommandStackT, raise_on_error: bool
1965 ):
1966 # build up all commands into a single request to increase network perf
1967 all_cmds = connection.pack_commands([args for args, _ in commands])
1968 await connection.send_packed_command(all_cmds)
1970 response = []
1971 for args, options in commands:
1972 try:
1973 response.append(
1974 await self.parse_response(connection, args[0], **options)
1975 )
1976 except ResponseError as e:
1977 response.append(e)
1979 if raise_on_error:
1980 self.raise_first_error(commands, response)
1981 return response
1983 def raise_first_error(self, commands: CommandStackT, response: Iterable[Any]):
1984 for i, r in enumerate(response):
1985 if isinstance(r, ResponseError):
1986 self.annotate_exception(r, i + 1, commands[i][0])
1987 raise r
1989 def annotate_exception(
1990 self, exception: Exception, number: int, command: Iterable[object]
1991 ) -> None:
1992 cmd = " ".join(map(safe_str, command))
1993 msg = (
1994 f"Command # {number} ({truncate_text(cmd)}) "
1995 f"of pipeline caused error: {exception.args}"
1996 )
1997 exception.args = (msg,) + exception.args[1:]
1999 async def parse_response(
2000 self, connection: Connection, command_name: Union[str, bytes], **options
2001 ):
2002 result = await super().parse_response(connection, command_name, **options)
2003 if command_name in self.UNWATCH_COMMANDS:
2004 self.watching = False
2005 elif command_name == "WATCH":
2006 self.watching = True
2007 return result
2009 async def load_scripts(self):
2010 # make sure all scripts that are about to be run on this pipeline exist
2011 scripts = list(self.scripts)
2012 immediate = self.immediate_execute_command
2013 shas = [s.sha for s in scripts]
2014 # we can't use the normal script_* methods because they would just
2015 # get buffered in the pipeline.
2016 exists = await immediate("SCRIPT EXISTS", *shas)
2017 if not all(exists):
2018 for s, exist in zip(scripts, exists):
2019 if not exist:
2020 s.sha = await immediate("SCRIPT LOAD", s.script)
2022 async def _disconnect_raise_on_watching(
2023 self,
2024 conn: Connection,
2025 error: Exception,
2026 failure_count: Optional[int] = None,
2027 start_time: Optional[float] = None,
2028 command_name: Optional[str] = None,
2029 ):
2030 """
2031 Close the connection, raise an exception if we were watching.
2033 The supported exceptions are already checked in the
2034 retry object so we don't need to do it here.
2036 After we disconnect the connection, it will try to reconnect and
2037 do a health check as part of the send_command logic(on connection level).
2038 """
2039 if (
2040 error
2041 and failure_count is not None
2042 and failure_count <= conn.retry.get_retries()
2043 ):
2044 await record_operation_duration(
2045 command_name=command_name,
2046 duration_seconds=time.monotonic() - start_time,
2047 server_address=getattr(conn, "host", None),
2048 server_port=getattr(conn, "port", None),
2049 db_namespace=str(conn.db),
2050 error=error,
2051 retry_attempts=failure_count,
2052 )
2053 await conn.disconnect(error=error, failure_count=failure_count)
2054 # if we were watching a variable, the watch is no longer valid
2055 # since this connection has died. raise a WatchError, which
2056 # indicates the user should retry this transaction.
2057 if self.watching:
2058 raise WatchError(
2059 f"A {type(error).__name__} occurred while watching one or more keys"
2060 )
2062 async def execute(self, raise_on_error: bool = True) -> List[Any]:
2063 """Execute all the commands in the current pipeline"""
2064 stack = self.command_stack
2065 if not stack and not self.watching:
2066 return []
2067 if self.scripts:
2068 await self.load_scripts()
2069 if self.is_transaction or self.explicit_transaction:
2070 execute = self._execute_transaction
2071 operation_name = "MULTI"
2072 else:
2073 execute = self._execute_pipeline
2074 operation_name = "PIPELINE"
2076 conn = self.connection
2077 if not conn:
2078 conn = await self.connection_pool.get_connection()
2079 # assign to self.connection so reset() releases the connection
2080 # back to the pool after we're done
2081 self.connection = conn
2082 conn = cast(Connection, conn)
2084 # Start timing for observability
2085 start_time = time.monotonic()
2086 # Track actual retry attempts for error reporting
2087 actual_retry_attempts = 0
2089 def failure_callback(error, failure_count):
2090 nonlocal actual_retry_attempts
2091 actual_retry_attempts = failure_count
2092 return self._disconnect_raise_on_watching(
2093 conn, error, failure_count, start_time, operation_name
2094 )
2096 try:
2097 response = await conn.retry.call_with_retry(
2098 lambda: execute(conn, stack, raise_on_error),
2099 failure_callback,
2100 with_failure_count=True,
2101 )
2103 await record_operation_duration(
2104 command_name=operation_name,
2105 duration_seconds=time.monotonic() - start_time,
2106 server_address=getattr(conn, "host", None),
2107 server_port=getattr(conn, "port", None),
2108 db_namespace=str(conn.db),
2109 )
2110 return response
2111 except Exception as e:
2112 await record_error_count(
2113 server_address=getattr(conn, "host", None),
2114 server_port=getattr(conn, "port", None),
2115 network_peer_address=getattr(conn, "host", None),
2116 network_peer_port=getattr(conn, "port", None),
2117 error_type=e,
2118 retry_attempts=actual_retry_attempts,
2119 is_internal=False,
2120 )
2121 raise
2122 finally:
2123 await self.reset()
2125 async def discard(self):
2126 """Flushes all previously queued commands
2127 See: https://redis.io/commands/DISCARD
2128 """
2129 await self.execute_command("DISCARD")
2131 async def watch(self, *names: KeyT):
2132 """Watches the values at keys ``names``"""
2133 if self.explicit_transaction:
2134 raise RedisError("Cannot issue a WATCH after a MULTI")
2135 return await self.execute_command("WATCH", *names)
2137 async def unwatch(self):
2138 """Unwatches all previously specified keys"""
2139 return self.watching and await self.execute_command("UNWATCH") or True