1import asyncio
2import contextlib
3import copy
4import inspect
5import math
6import socket
7import sys
8import time
9import warnings
10import weakref
11from abc import ABC, abstractmethod
12from itertools import chain
13from types import MappingProxyType
14from typing import (
15 Any,
16 AsyncIterator,
17 Callable,
18 Iterable,
19 List,
20 Literal,
21 Mapping,
22 Optional,
23 Protocol,
24 Set,
25 Tuple,
26 Type,
27 TypedDict,
28 TypeVar,
29 Union,
30)
31from urllib.parse import ParseResult, parse_qs, unquote, urlparse
32
33from ..observability.attributes import (
34 DB_CLIENT_CONNECTION_POOL_NAME,
35 DB_CLIENT_CONNECTION_STATE,
36 AttributeBuilder,
37 ConnectionState,
38 get_pool_name,
39)
40from ..utils import SSL_AVAILABLE, deprecated_function
41
42if SSL_AVAILABLE:
43 import ssl
44 from ssl import SSLContext, TLSVersion, VerifyFlags
45else:
46 ssl = None
47 TLSVersion = None
48 SSLContext = None
49 VerifyFlags = None
50
51from ..auth.token import TokenInterface
52from ..driver_info import DriverInfo, resolve_driver_info
53from ..event import AsyncAfterConnectionReleasedEvent, EventDispatcher
54from ..utils import deprecated_args, format_error_message
55
56# the functionality is available in 3.11.x but has a major issue before
57# 3.11.3. See https://github.com/redis/redis-py/issues/2633
58if sys.version_info >= (3, 11, 3):
59 from asyncio import timeout as async_timeout
60else:
61 from async_timeout import timeout as async_timeout
62
63from redis.asyncio.maint_notifications import (
64 AsyncMaintNotificationsConnectionHandler,
65 AsyncMaintNotificationsPoolHandler,
66 AsyncOSSMaintNotificationsHandler,
67)
68from redis.asyncio.observability.recorder import (
69 record_connection_closed,
70 record_connection_count,
71 record_connection_create_time,
72 record_connection_wait_time,
73 record_error_count,
74)
75from redis.asyncio.retry import Retry
76from redis.backoff import NoBackoff
77from redis.credentials import CredentialProvider, UsernamePasswordCredentialProvider
78from redis.exceptions import (
79 AuthenticationError,
80 AuthenticationWrongNumberOfArgsError,
81 ConnectionError,
82 DataError,
83 MaxConnectionsError,
84 RedisError,
85 ResponseError,
86 TimeoutError,
87)
88from redis.maint_notifications import (
89 MaintenanceState,
90 MaintNotificationsConfig,
91 NodeMovingNotification,
92 _build_moving_cleanup_connection_kwargs,
93 _build_moving_connection_kwargs,
94)
95from redis.observability.metrics import CloseReason
96from redis.typing import EncodableT
97from redis.utils import (
98 DEFAULT_RESP_VERSION,
99 HIREDIS_AVAILABLE,
100 SENTINEL,
101 check_protocol_version,
102 str_if_bytes,
103)
104
105from .._defaults import (
106 DEFAULT_SOCKET_CONNECT_TIMEOUT,
107 DEFAULT_SOCKET_READ_SIZE,
108 DEFAULT_SOCKET_TIMEOUT,
109 get_default_socket_keepalive_options,
110)
111from .._parsers import (
112 AsyncPushNotificationsParser,
113 BaseParser,
114 Encoder,
115 _AsyncHiredisParser,
116 _AsyncRESP2Parser,
117 _AsyncRESP3Parser,
118)
119
120SYM_STAR = b"*"
121SYM_DOLLAR = b"$"
122SYM_CRLF = b"\r\n"
123SYM_LF = b"\n"
124SYM_EMPTY = b""
125
126DefaultParser: Type[Union[_AsyncRESP2Parser, _AsyncRESP3Parser, _AsyncHiredisParser]]
127if HIREDIS_AVAILABLE:
128 DefaultParser = _AsyncHiredisParser
129else:
130 DefaultParser = _AsyncRESP3Parser
131
132
133class ConnectCallbackProtocol(Protocol):
134 def __call__(self, connection: "AbstractConnection"): ...
135
136
137class AsyncConnectCallbackProtocol(Protocol):
138 async def __call__(self, connection: "AbstractConnection"): ...
139
140
141ConnectCallbackT = Union[ConnectCallbackProtocol, AsyncConnectCallbackProtocol]
142
143
144class AsyncMaintNotificationsAbstractConnection:
145 """
146 Internal mixin for async maintenance notification state and parser handlers.
147
148 The sync implementation uses the same mixin-style structure. The async
149 version keeps the notification state and parser handler installation close
150 to the connection without sending the server-side handshake; that is wired
151 in a later step.
152 """
153
154 __slots__ = ()
155
156 def __init__(
157 self,
158 maint_notifications_config: MaintNotificationsConfig | None,
159 maint_notifications_pool_handler: (
160 AsyncMaintNotificationsPoolHandler | None
161 ) = None,
162 maintenance_state: MaintenanceState = MaintenanceState.NONE,
163 maintenance_notification_hash: int | None = None,
164 orig_host_address: str | None = None,
165 orig_socket_timeout: float | None = None,
166 orig_socket_connect_timeout: float | None = None,
167 oss_cluster_maint_notifications_handler: (
168 AsyncOSSMaintNotificationsHandler | None
169 ) = None,
170 parser: BaseParser | None = None,
171 ) -> None:
172 self.maint_notifications_config = maint_notifications_config
173 self.maintenance_state = maintenance_state
174 self.maintenance_notification_hash = maintenance_notification_hash
175 self._processed_start_maint_notifications: set[int] = set()
176 self._skipped_end_maint_notifications: set[int] = set()
177 self._configure_maintenance_notifications(
178 maint_notifications_pool_handler,
179 orig_host_address,
180 orig_socket_timeout,
181 orig_socket_connect_timeout,
182 oss_cluster_maint_notifications_handler,
183 parser,
184 )
185
186 @abstractmethod
187 def _get_parser(self) -> BaseParser:
188 pass
189
190 def _get_push_notifications_parser(self) -> AsyncPushNotificationsParser:
191 parser = self._get_parser()
192 if not isinstance(parser, (_AsyncHiredisParser, _AsyncRESP3Parser)):
193 raise RedisError(
194 "Maintenance notifications are only supported with hiredis and RESP3 parsers!"
195 )
196 return parser
197
198 @abstractmethod
199 def get_protocol(self):
200 pass
201
202 @abstractmethod
203 async def send_command(self, *args: Any, **kwargs: Any) -> None:
204 pass
205
206 @abstractmethod
207 async def read_response(
208 self,
209 disable_decoding: bool = False,
210 timeout: float | None = None,
211 *,
212 disconnect_on_error: bool = True,
213 push_request: bool | None = False,
214 ) -> Any:
215 pass
216
217 @abstractmethod
218 def getpeername(self) -> str | None:
219 pass
220
221 def _configure_maintenance_notifications(
222 self,
223 maint_notifications_pool_handler: (
224 AsyncMaintNotificationsPoolHandler | None
225 ) = None,
226 orig_host_address: str | None = None,
227 orig_socket_timeout: float | None = None,
228 orig_socket_connect_timeout: float | None = None,
229 oss_cluster_maint_notifications_handler: (
230 AsyncOSSMaintNotificationsHandler | None
231 ) = None,
232 parser: BaseParser | None = None,
233 ) -> None:
234 if (
235 not self.maint_notifications_config
236 or not self.maint_notifications_config.enabled
237 ):
238 self._maint_notifications_pool_handler = None
239 self._maint_notifications_connection_handler = None
240 self._oss_cluster_maint_notifications_handler = None
241 return
242
243 if not parser:
244 raise RedisError(
245 "To configure maintenance notifications, a parser must be provided!"
246 )
247
248 if not isinstance(parser, _AsyncHiredisParser) and not isinstance(
249 parser, _AsyncRESP3Parser
250 ):
251 raise RedisError(
252 "Maintenance notifications are only supported with hiredis and RESP3 parsers!"
253 )
254
255 if maint_notifications_pool_handler:
256 # Extract a reference to a new pool handler that copies all properties
257 # of the original one and has a different connection reference
258 # This is needed because when we attach the handler to the parser
259 # we need to make sure that the handler has a reference to the
260 # connection that the parser is attached to.
261 self._maint_notifications_pool_handler = (
262 maint_notifications_pool_handler.get_handler_for_connection()
263 )
264 self._maint_notifications_pool_handler.set_connection(self)
265 else:
266 self._maint_notifications_pool_handler = None
267
268 self._maint_notifications_connection_handler = (
269 AsyncMaintNotificationsConnectionHandler(
270 self, self.maint_notifications_config
271 )
272 )
273
274 if oss_cluster_maint_notifications_handler:
275 self._oss_cluster_maint_notifications_handler = (
276 oss_cluster_maint_notifications_handler
277 )
278 parser.set_oss_cluster_maint_push_handler(
279 oss_cluster_maint_notifications_handler.handle_notification
280 )
281 else:
282 self._oss_cluster_maint_notifications_handler = None
283
284 # Set up pool handler to parser if available
285 if self._maint_notifications_pool_handler:
286 parser.set_node_moving_push_handler(
287 self._maint_notifications_pool_handler.handle_notification
288 )
289
290 # Set up connection handler
291 parser.set_maintenance_push_handler(
292 self._maint_notifications_connection_handler.handle_notification
293 )
294
295 self.orig_host_address = orig_host_address if orig_host_address else self.host
296 self.orig_socket_timeout = (
297 orig_socket_timeout if orig_socket_timeout else self.socket_timeout
298 )
299 self.orig_socket_connect_timeout = (
300 orig_socket_connect_timeout
301 if orig_socket_connect_timeout
302 else self.socket_connect_timeout
303 )
304
305 def set_maint_notifications_pool_handler_for_connection(
306 self, maint_notifications_pool_handler: AsyncMaintNotificationsPoolHandler
307 ) -> None:
308 # Deep copy the pool handler to avoid sharing the same pool handler
309 # between multiple connections, because otherwise each connection will override
310 # the connection reference and the pool handler will only hold a reference
311 # to the last connection that was set.
312 maint_notifications_pool_handler_copy = (
313 maint_notifications_pool_handler.get_handler_for_connection()
314 )
315 maint_notifications_pool_handler_copy.set_connection(self)
316 parser = self._get_push_notifications_parser()
317 parser.set_node_moving_push_handler(
318 maint_notifications_pool_handler_copy.handle_notification
319 )
320 self._maint_notifications_pool_handler = maint_notifications_pool_handler_copy
321
322 # Update maintenance notification connection handler if it doesn't exist
323 if not self._maint_notifications_connection_handler:
324 self._maint_notifications_connection_handler = (
325 AsyncMaintNotificationsConnectionHandler(
326 self, maint_notifications_pool_handler.config
327 )
328 )
329 parser.set_maintenance_push_handler(
330 self._maint_notifications_connection_handler.handle_notification
331 )
332 else:
333 self._maint_notifications_connection_handler.config = (
334 maint_notifications_pool_handler.config
335 )
336
337 def set_maint_notifications_cluster_handler_for_connection(
338 self,
339 oss_cluster_maint_notifications_handler: AsyncOSSMaintNotificationsHandler,
340 ) -> None:
341 parser = self._get_push_notifications_parser()
342 parser.set_oss_cluster_maint_push_handler(
343 oss_cluster_maint_notifications_handler.handle_notification
344 )
345 # OSS cluster mode and pool-handler mode are mutually exclusive. Clear
346 # any node-moving/pool handler a default (RESP3 "auto") pool wired in
347 # __init__ so this existing connection is not configured with both.
348 parser.set_node_moving_push_handler(None)
349 self._maint_notifications_pool_handler = None
350
351 self._oss_cluster_maint_notifications_handler = (
352 oss_cluster_maint_notifications_handler
353 )
354
355 # Update maintenance notification connection handler if it doesn't exist
356 if not self._maint_notifications_connection_handler:
357 self._maint_notifications_connection_handler = (
358 AsyncMaintNotificationsConnectionHandler(
359 self, oss_cluster_maint_notifications_handler.config
360 )
361 )
362 parser.set_maintenance_push_handler(
363 self._maint_notifications_connection_handler.handle_notification
364 )
365 else:
366 self._maint_notifications_connection_handler.config = (
367 oss_cluster_maint_notifications_handler.config
368 )
369
370 async def activate_maint_notifications_handling_if_enabled(
371 self, check_health: bool = True
372 ) -> None:
373 # Send maintenance notifications handshake if RESP3 is active
374 # and maintenance notifications are enabled
375 # and we have a host to determine the endpoint type from
376 # When the maint_notifications_config enabled mode is "auto",
377 # we just log a warning if the handshake fails
378 # When the mode is enabled=True, we raise an exception in case of failure
379 host = getattr(self, "host", None)
380 if (
381 check_protocol_version(self.get_protocol(), 3)
382 and self.maint_notifications_config
383 and self.maint_notifications_config.enabled
384 and self._maint_notifications_connection_handler
385 and host is not None
386 ):
387 await self._enable_maintenance_notifications(
388 maint_notifications_config=self.maint_notifications_config,
389 check_health=check_health,
390 )
391
392 async def _enable_maintenance_notifications(
393 self,
394 maint_notifications_config: MaintNotificationsConfig,
395 check_health: bool = True,
396 ) -> None:
397 try:
398 host = getattr(self, "host", None)
399 if host is None:
400 raise ValueError(
401 "Cannot enable maintenance notifications for connection"
402 " object that doesn't have a host attribute."
403 )
404
405 endpoint_type = maint_notifications_config.get_endpoint_type(host, self)
406 await self.send_command(
407 "CLIENT",
408 "MAINT_NOTIFICATIONS",
409 "ON",
410 "moving-endpoint-type",
411 endpoint_type.value,
412 check_health=check_health,
413 )
414 response = await self.read_response()
415 if not response or str_if_bytes(response) != "OK":
416 raise ResponseError(
417 "The server doesn't support maintenance notifications"
418 )
419 except Exception as e:
420 if (
421 isinstance(e, ResponseError)
422 and maint_notifications_config.enabled == "auto"
423 ):
424 # Log warning but don't fail the connection
425 import logging
426
427 logger = logging.getLogger(__name__)
428 logger.debug(f"Failed to enable maintenance notifications: {e}")
429 else:
430 raise
431
432 def get_resolved_ip(self) -> str | None:
433 """
434 Extract the resolved IP address from an established connection or host.
435
436 First tries to get the actual peer IP from the async stream writer, then
437 falls back to DNS resolution if needed.
438
439 Returns:
440 The resolved IP address, or None if it cannot be determined.
441 """
442
443 # Method 1: Try to get the actual IP from the established stream.
444 # This is most accurate as it shows the exact IP being used.
445 try:
446 peer_addr = self.getpeername()
447 if peer_addr:
448 return peer_addr
449 except (AttributeError, OSError):
450 # Stream might not be connected or peer address lookup might fail.
451 pass
452
453 # Method 2: Fall back to the configured host (which may be an IP or an
454 # FQDN). Unlike the sync client we intentionally do NOT call
455 # socket.getaddrinfo() here: this method runs on the event loop, so a
456 # blocking DNS resolution would stall it. On the endpoint-type handshake
457 # path (get_endpoint_type) getpeername() above always succeeds because the
458 # writer was just connected, so this fallback is only reached by the
459 # debug-log call sites during reconnects — where returning the host is
460 # fine. A blocking getaddrinfo on an FQDN host there can freeze the loop
461 # for seconds and trip unrelated connect timeouts.
462 return getattr(self, "host", None)
463
464 @property
465 def maintenance_state(self) -> MaintenanceState:
466 return self._maintenance_state
467
468 @maintenance_state.setter
469 def maintenance_state(self, state: MaintenanceState) -> None:
470 self._maintenance_state = state
471
472 def add_maint_start_notification(self, id: int) -> None:
473 self._processed_start_maint_notifications.add(id)
474
475 def get_processed_start_notifications(self) -> set[int]:
476 return self._processed_start_maint_notifications
477
478 def add_skipped_end_notification(self, id: int) -> None:
479 self._skipped_end_maint_notifications.add(id)
480
481 def get_skipped_end_notifications(self) -> set[int]:
482 return self._skipped_end_maint_notifications
483
484 def reset_received_notifications(self) -> None:
485 self._processed_start_maint_notifications.clear()
486 self._skipped_end_maint_notifications.clear()
487
488 def update_current_socket_timeout(
489 self, relaxed_timeout: float | None = None
490 ) -> None:
491 timeout = relaxed_timeout if relaxed_timeout != -1 else self.socket_timeout
492 self._reschedule_active_read_timeout(timeout)
493
494 def _reschedule_active_read_timeout(self, timeout: float | None) -> None:
495 timeout_context = getattr(self, "_active_read_timeout", None)
496 if timeout_context is None:
497 # No read_response call is currently inside its socket timeout
498 # context, so there is no in-flight deadline to relax or restore.
499 return
500
501 if timeout is None:
502 # A None socket timeout means the active read should become blocking.
503 # Python 3.11's timeout context supports clearing the deadline.
504 if hasattr(timeout_context, "reschedule"):
505 timeout_context.reschedule(None)
506 # Older async-timeout contexts cannot clear a deadline, so reject the
507 # current timeout instead of leaving a stale relaxed deadline active.
508 elif hasattr(timeout_context, "reject"):
509 timeout_context.reject()
510 return
511
512 # Active read timeouts are stored as loop-time deadlines, not durations.
513 deadline = asyncio.get_running_loop().time() + timeout
514 if hasattr(timeout_context, "reschedule"):
515 # Python 3.11 asyncio.timeout exposes reschedule().
516 timeout_context.reschedule(deadline)
517 elif hasattr(timeout_context, "update"):
518 # async-timeout exposes update() for the same deadline adjustment.
519 timeout_context.update(deadline)
520
521 def set_tmp_settings(
522 self,
523 tmp_host_address: str | object | None = SENTINEL,
524 tmp_relaxed_timeout: float | None = -1,
525 ) -> None:
526 """
527 SENTINEL keeps the host unchanged. -1 keeps the relaxed timeout unchanged.
528 """
529 if tmp_host_address and tmp_host_address != SENTINEL:
530 self.host = str(tmp_host_address)
531 if tmp_relaxed_timeout != -1:
532 self.socket_timeout = tmp_relaxed_timeout
533 self.socket_connect_timeout = tmp_relaxed_timeout
534
535 def reset_tmp_settings(
536 self,
537 reset_host_address: bool = False,
538 reset_relaxed_timeout: bool = False,
539 ) -> None:
540 if reset_host_address:
541 self.host = self.orig_host_address
542 if reset_relaxed_timeout:
543 self.socket_timeout = self.orig_socket_timeout
544 self.socket_connect_timeout = self.orig_socket_connect_timeout
545
546
547class AbstractConnection(AsyncMaintNotificationsAbstractConnection):
548 """Manages communication to and from a Redis server"""
549
550 __slots__ = (
551 "db",
552 "username",
553 "client_name",
554 "lib_name",
555 "lib_version",
556 "credential_provider",
557 "password",
558 "socket_timeout",
559 "socket_connect_timeout",
560 "redis_connect_func",
561 "retry_on_timeout",
562 "retry_on_error",
563 "health_check_interval",
564 "next_health_check",
565 "last_active_at",
566 "encoder",
567 "ssl_context",
568 "protocol",
569 "_reader",
570 "_writer",
571 "_parser",
572 "_active_read_timeout",
573 "_connect_callbacks",
574 "_buffer_cutoff",
575 "_lock",
576 "_socket_read_size",
577 "__dict__",
578 )
579
580 @deprecated_args(
581 args_to_warn=["lib_name", "lib_version"],
582 reason="Use 'driver_info' parameter instead. "
583 "lib_name and lib_version will be removed in a future version.",
584 )
585 def __init__(
586 self,
587 *,
588 db: str | int = 0,
589 password: str | None = None,
590 socket_timeout: float | None = DEFAULT_SOCKET_TIMEOUT,
591 socket_connect_timeout: float | None = DEFAULT_SOCKET_CONNECT_TIMEOUT,
592 retry_on_timeout: bool = False,
593 retry_on_error: list | object = SENTINEL,
594 encoding: str = "utf-8",
595 encoding_errors: str = "strict",
596 decode_responses: bool = False,
597 parser_class: Type[BaseParser] = DefaultParser,
598 socket_read_size: int = DEFAULT_SOCKET_READ_SIZE,
599 health_check_interval: float = 0,
600 client_name: str | None = None,
601 lib_name: str | object | None = SENTINEL,
602 lib_version: str | object | None = SENTINEL,
603 driver_info: DriverInfo | object | None = SENTINEL,
604 username: str | None = None,
605 retry: Retry | None = None,
606 redis_connect_func: ConnectCallbackT | None = None,
607 encoder_class: Type[Encoder] = Encoder,
608 credential_provider: CredentialProvider | None = None,
609 protocol: int | None = None,
610 legacy_responses: bool = True,
611 event_dispatcher: EventDispatcher | None = None,
612 maint_notifications_config: MaintNotificationsConfig | None = None,
613 maint_notifications_pool_handler: (
614 AsyncMaintNotificationsPoolHandler | None
615 ) = None,
616 maintenance_state: MaintenanceState = MaintenanceState.NONE,
617 maintenance_notification_hash: int | None = None,
618 orig_host_address: str | None = None,
619 orig_socket_timeout: float | None = None,
620 orig_socket_connect_timeout: float | None = None,
621 oss_cluster_maint_notifications_handler: (
622 AsyncOSSMaintNotificationsHandler | None
623 ) = None,
624 ):
625 """
626 Initialize a new async Connection.
627
628 Parameters
629 ----------
630 driver_info : DriverInfo, optional
631 Driver metadata for CLIENT SETINFO. If provided, lib_name and lib_version
632 are ignored. If not provided, a DriverInfo will be created from lib_name
633 and lib_version. Explicit None disables CLIENT SETINFO.
634 lib_name : str, optional
635 **Deprecated.** Use driver_info instead. Library name for CLIENT SETINFO.
636 lib_version : str, optional
637 **Deprecated.** Use driver_info instead. Library version for CLIENT SETINFO.
638 """
639 if (username or password) and credential_provider is not None:
640 raise DataError(
641 "'username' and 'password' cannot be passed along with 'credential_"
642 "provider'. Please provide only one of the following arguments: \n"
643 "1. 'password' and (optional) 'username'\n"
644 "2. 'credential_provider'"
645 )
646 if event_dispatcher is None:
647 self._event_dispatcher = EventDispatcher()
648 else:
649 self._event_dispatcher = event_dispatcher
650 self.db = db
651 self.client_name = client_name
652
653 # Handle driver_info: if provided, use it; otherwise create from lib_name/lib_version.
654 self.driver_info = resolve_driver_info(driver_info, lib_name, lib_version)
655
656 self.credential_provider = credential_provider
657 self.password = password
658 self.username = username
659 self.socket_timeout = socket_timeout
660 if socket_connect_timeout is None:
661 socket_connect_timeout = socket_timeout
662 self.socket_connect_timeout = socket_connect_timeout
663 self.retry_on_timeout = retry_on_timeout
664 if retry_on_error is SENTINEL:
665 retry_on_error = []
666 if retry_on_timeout:
667 retry_on_error.append(TimeoutError)
668 retry_on_error.append(socket.timeout)
669 retry_on_error.append(asyncio.TimeoutError)
670 self.retry_on_error = retry_on_error
671 if retry or retry_on_error:
672 if not retry:
673 self.retry = Retry(NoBackoff(), 1)
674 else:
675 # deep-copy the Retry object as it is mutable
676 self.retry = copy.deepcopy(retry)
677 # Update the retry's supported errors with the specified errors
678 self.retry.update_supported_errors(retry_on_error)
679 else:
680 self.retry = Retry(NoBackoff(), 0)
681 self.health_check_interval = health_check_interval
682 self.next_health_check: float = -1
683 self.encoder = encoder_class(encoding, encoding_errors, decode_responses)
684 self.redis_connect_func = redis_connect_func
685 self._reader: Optional[asyncio.StreamReader] = None
686 self._writer: Optional[asyncio.StreamWriter] = None
687 self._socket_read_size = socket_read_size
688 self._active_read_timeout = None
689 self._connect_callbacks: List[weakref.WeakMethod[ConnectCallbackT]] = []
690 self._buffer_cutoff = 6000
691 self._re_auth_token: Optional[TokenInterface] = None
692 self._should_reconnect = False
693
694 try:
695 p = int(protocol)
696 except TypeError:
697 p = DEFAULT_RESP_VERSION
698 except ValueError:
699 raise ConnectionError("protocol must be an integer")
700 else:
701 if p < 2 or p > 3:
702 raise ConnectionError("protocol must be either 2 or 3")
703 self.protocol = p
704 self.legacy_responses = legacy_responses
705 if parser_class != _AsyncHiredisParser:
706 # The Python parsers are protocol-specific; hiredis supports both.
707 if self.protocol == 3 and parser_class == _AsyncRESP2Parser:
708 parser_class = _AsyncRESP3Parser
709 elif self.protocol == 2 and parser_class == _AsyncRESP3Parser:
710 parser_class = _AsyncRESP2Parser
711 self.set_parser(parser_class)
712 AsyncMaintNotificationsAbstractConnection.__init__(
713 self,
714 maint_notifications_config,
715 maint_notifications_pool_handler,
716 maintenance_state,
717 maintenance_notification_hash,
718 orig_host_address,
719 orig_socket_timeout,
720 orig_socket_connect_timeout,
721 oss_cluster_maint_notifications_handler,
722 self._parser,
723 )
724
725 def __del__(self, _warnings: Any = warnings):
726 # For some reason, the individual streams don't get properly garbage
727 # collected and therefore produce no resource warnings. We add one
728 # here, in the same style as those from the stdlib.
729 if getattr(self, "_writer", None):
730 _warnings.warn(
731 f"unclosed Connection {self!r}", ResourceWarning, source=self
732 )
733
734 try:
735 asyncio.get_running_loop()
736 self._close()
737 except RuntimeError:
738 # No actions been taken if pool already closed.
739 pass
740
741 def _close(self):
742 """
743 Internal method to silently close the connection without waiting
744 """
745 if self._writer:
746 self._writer.close()
747 self._writer = self._reader = None
748
749 def __repr__(self):
750 repr_args = ",".join((f"{k}={v}" for k, v in self.repr_pieces()))
751 return f"<{self.__class__.__module__}.{self.__class__.__name__}({repr_args})>"
752
753 @abstractmethod
754 def repr_pieces(self):
755 pass
756
757 @property
758 def is_connected(self):
759 return self._reader is not None and self._writer is not None
760
761 def register_connect_callback(self, callback):
762 """
763 Register a callback to be called when the connection is established either
764 initially or reconnected. This allows listeners to issue commands that
765 are ephemeral to the connection, for example pub/sub subscription or
766 key tracking. The callback must be a _method_ and will be kept as
767 a weak reference.
768 """
769 wm = weakref.WeakMethod(callback)
770 if wm not in self._connect_callbacks:
771 self._connect_callbacks.append(wm)
772
773 def deregister_connect_callback(self, callback):
774 """
775 De-register a previously registered callback. It will no-longer receive
776 notifications on connection events. Calling this is not required when the
777 listener goes away, since the callbacks are kept as weak methods.
778 """
779 try:
780 self._connect_callbacks.remove(weakref.WeakMethod(callback))
781 except ValueError:
782 pass
783
784 def set_parser(self, parser_class: Type[BaseParser]) -> None:
785 """
786 Creates a new instance of parser_class with socket size:
787 _socket_read_size and assigns it to the parser for the connection
788 :param parser_class: The required parser class
789 """
790 self._parser = parser_class(socket_read_size=self._socket_read_size)
791
792 def _get_parser(self) -> BaseParser:
793 return self._parser
794
795 def getpeername(self) -> str | None:
796 """
797 Returns the peer name of the connection.
798 """
799 writer = self._writer
800 if writer is None:
801 return None
802 peername = writer.get_extra_info("peername")
803 if isinstance(peername, tuple) and peername:
804 return str(peername[0])
805 return None
806
807 async def connect(self):
808 """Connects to the Redis server if not already connected"""
809 # try once the socket connect with the handshake, retry the whole
810 # connect/handshake flow based on retry policy
811 await self.retry.call_with_retry(
812 lambda: self.connect_check_health(
813 check_health=True, retry_socket_connect=False
814 ),
815 lambda error, failure_count: self.disconnect(
816 error=error, failure_count=failure_count
817 ),
818 with_failure_count=True,
819 )
820
821 async def connect_check_health(
822 self, check_health: bool = True, retry_socket_connect: bool = True
823 ):
824 if self.is_connected:
825 return
826 # Track actual retry attempts for error reporting
827 actual_retry_attempts = 0
828
829 def failure_callback(error, failure_count):
830 nonlocal actual_retry_attempts
831 actual_retry_attempts = failure_count
832 return self.disconnect(error=error, failure_count=failure_count)
833
834 try:
835 if retry_socket_connect:
836 await self.retry.call_with_retry(
837 lambda: self._connect(),
838 failure_callback,
839 with_failure_count=True,
840 )
841 else:
842 await self._connect()
843 except asyncio.CancelledError:
844 raise # in 3.7 and earlier, this is an Exception, not BaseException
845 except (socket.timeout, asyncio.TimeoutError):
846 e = TimeoutError("Timeout connecting to server")
847 await record_error_count(
848 server_address=getattr(self, "host", None),
849 server_port=getattr(self, "port", None),
850 network_peer_address=getattr(self, "host", None),
851 network_peer_port=getattr(self, "port", None),
852 error_type=e,
853 retry_attempts=actual_retry_attempts,
854 is_internal=False,
855 )
856 raise e
857 except OSError as e:
858 e = ConnectionError(self._error_message(e))
859 await record_error_count(
860 server_address=getattr(self, "host", None),
861 server_port=getattr(self, "port", None),
862 network_peer_address=getattr(self, "host", None),
863 network_peer_port=getattr(self, "port", None),
864 error_type=e,
865 retry_attempts=actual_retry_attempts,
866 is_internal=False,
867 )
868 raise e
869 except Exception as exc:
870 raise ConnectionError(exc) from exc
871
872 try:
873 if not self.redis_connect_func:
874 # Use the default on_connect function
875 await self.on_connect_check_health(check_health=check_health)
876 else:
877 # Use the passed function redis_connect_func
878 (
879 await self.redis_connect_func(self)
880 if asyncio.iscoroutinefunction(self.redis_connect_func)
881 else self.redis_connect_func(self)
882 )
883 except RedisError:
884 # clean up after any error in on_connect
885 await self.disconnect()
886 raise
887
888 # run any user callbacks. right now the only internal callback
889 # is for pubsub channel/pattern resubscription
890 # first, remove any dead weakrefs
891 self._connect_callbacks = [ref for ref in self._connect_callbacks if ref()]
892 for ref in self._connect_callbacks:
893 callback = ref()
894 task = callback(self)
895 if task and inspect.isawaitable(task):
896 await task
897
898 def mark_for_reconnect(self):
899 self._should_reconnect = True
900
901 def should_reconnect(self):
902 return self._should_reconnect
903
904 def reset_should_reconnect(self):
905 self._should_reconnect = False
906
907 @abstractmethod
908 async def _connect(self):
909 pass
910
911 @abstractmethod
912 def _host_error(self) -> str:
913 pass
914
915 def _error_message(self, exception: BaseException) -> str:
916 return format_error_message(self._host_error(), exception)
917
918 def get_protocol(self):
919 return self.protocol
920
921 async def on_connect(self) -> None:
922 """Initialize the connection, authenticate and select a database"""
923 await self.on_connect_check_health(check_health=True)
924
925 async def on_connect_check_health(self, check_health: bool = True) -> None:
926 self._parser.on_connect(self)
927 parser = self._parser
928
929 auth_args = None
930 # if credential provider or username and/or password are set, authenticate
931 if self.credential_provider or (self.username or self.password):
932 cred_provider = (
933 self.credential_provider
934 or UsernamePasswordCredentialProvider(self.username, self.password)
935 )
936 auth_args = await cred_provider.get_credentials_async()
937
938 # if resp version is specified and we have auth args,
939 # we need to send them via HELLO
940 if auth_args and check_protocol_version(self.protocol, 3):
941 if isinstance(self._parser, _AsyncRESP2Parser):
942 self.set_parser(_AsyncRESP3Parser)
943 # update cluster exception classes
944 self._parser.EXCEPTION_CLASSES = parser.EXCEPTION_CLASSES
945 self._parser.on_connect(self)
946 if len(auth_args) == 1:
947 auth_args = ["default", auth_args[0]]
948 # avoid checking health here -- PING will fail if we try
949 # to check the health prior to the AUTH
950 await self.send_command(
951 "HELLO", self.protocol, "AUTH", *auth_args, check_health=False
952 )
953 response = await self.read_response()
954 if response.get(b"proto") != int(self.protocol) and response.get(
955 "proto"
956 ) != int(self.protocol):
957 raise ConnectionError("Invalid RESP version")
958 # avoid checking health here -- PING will fail if we try
959 # to check the health prior to the AUTH
960 elif auth_args:
961 await self.send_command("AUTH", *auth_args, check_health=False)
962
963 try:
964 auth_response = await self.read_response()
965 except AuthenticationWrongNumberOfArgsError:
966 # a username and password were specified but the Redis
967 # server seems to be < 6.0.0 which expects a single password
968 # arg. retry auth with just the password.
969 # https://github.com/andymccurdy/redis-py/issues/1274
970 await self.send_command("AUTH", auth_args[-1], check_health=False)
971 auth_response = await self.read_response()
972
973 if str_if_bytes(auth_response) != "OK":
974 raise AuthenticationError("Invalid Username or Password")
975
976 # if resp version is specified, switch to it
977 elif check_protocol_version(self.protocol, 3):
978 if isinstance(self._parser, _AsyncRESP2Parser):
979 self.set_parser(_AsyncRESP3Parser)
980 # update cluster exception classes
981 self._parser.EXCEPTION_CLASSES = parser.EXCEPTION_CLASSES
982 self._parser.on_connect(self)
983 await self.send_command("HELLO", self.protocol, check_health=check_health)
984 response = await self.read_response()
985 # if response.get(b"proto") != self.protocol and response.get(
986 # "proto"
987 # ) != self.protocol:
988 # raise ConnectionError("Invalid RESP version")
989
990 # Activate maintenance notifications for this connection
991 # if enabled in the configuration
992 # This is a no-op if maintenance notifications are not enabled
993 await self.activate_maint_notifications_handling_if_enabled(
994 check_health=check_health
995 )
996
997 # if a client_name is given, set it
998 if self.client_name:
999 await self.send_command(
1000 "CLIENT",
1001 "SETNAME",
1002 self.client_name,
1003 check_health=check_health,
1004 )
1005 if str_if_bytes(await self.read_response()) != "OK":
1006 raise ConnectionError("Error setting client name")
1007
1008 # Set the library name and version from driver_info, pipeline for lower startup latency
1009 lib_name_sent = False
1010 lib_version_sent = False
1011
1012 if self.driver_info and self.driver_info.formatted_name:
1013 await self.send_command(
1014 "CLIENT",
1015 "SETINFO",
1016 "LIB-NAME",
1017 self.driver_info.formatted_name,
1018 check_health=check_health,
1019 )
1020 lib_name_sent = True
1021
1022 if self.driver_info and self.driver_info.lib_version:
1023 await self.send_command(
1024 "CLIENT",
1025 "SETINFO",
1026 "LIB-VER",
1027 self.driver_info.lib_version,
1028 check_health=check_health,
1029 )
1030 lib_version_sent = True
1031
1032 # if a database is specified, switch to it. Also pipeline this
1033 if self.db:
1034 await self.send_command("SELECT", self.db, check_health=check_health)
1035
1036 # read responses from pipeline
1037 for _ in range(sum([lib_name_sent, lib_version_sent])):
1038 try:
1039 await self.read_response()
1040 except ResponseError:
1041 pass
1042
1043 if self.db:
1044 if str_if_bytes(await self.read_response()) != "OK":
1045 raise ConnectionError("Invalid Database")
1046
1047 async def disconnect(
1048 self,
1049 nowait: bool = False,
1050 error: Optional[Exception] = None,
1051 failure_count: Optional[int] = None,
1052 health_check_failed: bool = False,
1053 ) -> None:
1054 """Disconnects from the Redis server"""
1055 # On Python 3.13+, asyncio.timeout() raises RuntimeError when called
1056 # outside a running Task (e.g. during GC finalization or event-loop
1057 # callbacks). In that context we fall back to a synchronous close.
1058 # See https://github.com/redis/redis-py/issues/3856
1059 if asyncio.current_task() is None:
1060 self._parser.on_disconnect()
1061 self.reset_should_reconnect()
1062 self._close()
1063 return
1064
1065 try:
1066 async with async_timeout(self.socket_connect_timeout):
1067 self._parser.on_disconnect()
1068 # Reset the reconnect flag
1069 self.reset_should_reconnect()
1070 if not self.is_connected:
1071 return
1072 try:
1073 self._writer.close() # type: ignore[union-attr]
1074 # wait for close to finish, except when handling errors and
1075 # forcefully disconnecting.
1076 if not nowait:
1077 await self._writer.wait_closed() # type: ignore[union-attr]
1078 except OSError:
1079 pass
1080 finally:
1081 self._reader = None
1082 self._writer = None
1083 except asyncio.TimeoutError:
1084 raise TimeoutError(
1085 f"Timed out closing connection after {self.socket_connect_timeout}"
1086 ) from None
1087
1088 if error:
1089 if health_check_failed:
1090 close_reason = CloseReason.HEALTHCHECK_FAILED
1091 else:
1092 close_reason = CloseReason.ERROR
1093
1094 if failure_count is not None and failure_count > self.retry.get_retries():
1095 await record_error_count(
1096 server_address=getattr(self, "host", None),
1097 server_port=getattr(self, "port", None),
1098 network_peer_address=getattr(self, "host", None),
1099 network_peer_port=getattr(self, "port", None),
1100 error_type=error,
1101 retry_attempts=failure_count,
1102 )
1103
1104 await record_connection_closed(
1105 close_reason=close_reason,
1106 error_type=error,
1107 )
1108 else:
1109 await record_connection_closed(
1110 close_reason=CloseReason.APPLICATION_CLOSE,
1111 )
1112
1113 if self.maintenance_state == MaintenanceState.MAINTENANCE:
1114 # MOVING state is owned by the pool-level TTL cleanup. Regular
1115 # maintenance timeout relaxation can be restored when this
1116 # connection closes, matching the sync lifecycle.
1117 self.reset_tmp_settings(reset_relaxed_timeout=True)
1118 self.maintenance_state = MaintenanceState.NONE
1119 # reset the sets that keep track of received start maint
1120 # notifications and skipped end maint notifications
1121 self.reset_received_notifications()
1122
1123 async def _send_ping(self):
1124 """Send PING, expect PONG in return"""
1125 await self.send_command("PING", check_health=False)
1126 if str_if_bytes(await self.read_response()) != "PONG":
1127 raise ConnectionError("Bad response from PING health check")
1128
1129 async def _ping_failed(self, error, failure_count):
1130 """Function to call when PING fails"""
1131 await self.disconnect(
1132 error=error, failure_count=failure_count, health_check_failed=True
1133 )
1134
1135 async def check_health(self):
1136 """Check the health of the connection with a PING/PONG"""
1137 if (
1138 self.health_check_interval
1139 and asyncio.get_running_loop().time() > self.next_health_check
1140 ):
1141 await self.retry.call_with_retry(
1142 self._send_ping, self._ping_failed, with_failure_count=True
1143 )
1144
1145 async def _send_packed_command(self, command: Iterable[bytes]) -> None:
1146 self._writer.writelines(command)
1147 await self._writer.drain()
1148
1149 async def send_packed_command(
1150 self, command: Union[bytes, str, Iterable[bytes]], check_health: bool = True
1151 ) -> None:
1152 if not self.is_connected:
1153 await self.connect_check_health(check_health=False)
1154 if check_health:
1155 await self.check_health()
1156
1157 try:
1158 if isinstance(command, str):
1159 command = command.encode()
1160 if isinstance(command, bytes):
1161 command = [command]
1162 if self.socket_timeout:
1163 await asyncio.wait_for(
1164 self._send_packed_command(command), self.socket_timeout
1165 )
1166 else:
1167 self._writer.writelines(command)
1168 await self._writer.drain()
1169 except asyncio.TimeoutError:
1170 await self.disconnect(nowait=True)
1171 raise TimeoutError("Timeout writing to socket") from None
1172 except OSError as e:
1173 await self.disconnect(nowait=True)
1174 if len(e.args) == 1:
1175 err_no, errmsg = "UNKNOWN", e.args[0]
1176 else:
1177 err_no = e.args[0]
1178 errmsg = e.args[1]
1179 raise ConnectionError(
1180 f"Error {err_no} while writing to socket. {errmsg}."
1181 ) from e
1182 except BaseException:
1183 # BaseExceptions can be raised when a socket send operation is not
1184 # finished, e.g. due to a timeout. Ideally, a caller could then re-try
1185 # to send un-sent data. However, the send_packed_command() API
1186 # does not support it so there is no point in keeping the connection open.
1187 await self.disconnect(nowait=True)
1188 raise
1189
1190 async def send_command(self, *args: Any, **kwargs: Any) -> None:
1191 """Pack and send a command to the Redis server"""
1192 await self.send_packed_command(
1193 self.pack_command(*args), check_health=kwargs.get("check_health", True)
1194 )
1195
1196 @deprecated_function(
1197 version="8.0.0", reason="Use can_read() instead", name="can_read_destructive"
1198 )
1199 async def can_read_destructive(self) -> bool:
1200 """Check the socket to see if there's data loaded in the buffer."""
1201 try:
1202 return await self._parser.can_read()
1203 except OSError as e:
1204 await self.disconnect(nowait=True)
1205 host_error = self._host_error()
1206 raise ConnectionError(f"Error while reading from {host_error}: {e.args}")
1207
1208 async def can_read(self) -> bool:
1209 """Check the socket to see if there's data loaded in the buffer."""
1210 # TODO: Rename this API; it detects pending data or dirty/closed
1211 # connection state, not only whether application data can be read.
1212 try:
1213 return await self._parser.can_read()
1214 except OSError as e:
1215 await self.disconnect(nowait=True)
1216 host_error = self._host_error()
1217 raise ConnectionError(f"Error while reading from {host_error}: {e.args}")
1218
1219 async def read_response(
1220 self,
1221 disable_decoding: bool = False,
1222 timeout: float | None = None,
1223 *,
1224 disconnect_on_error: bool = True,
1225 push_request: bool | None = False,
1226 ):
1227 """Read the response from a previously sent command.
1228
1229 ``timeout`` semantics:
1230 - ``None`` (default): fall back to ``self.socket_timeout``.
1231 - ``math.inf``: block indefinitely with no timeout. Used by PubSub
1232 blocking reads (``listen()`` / ``get_message(timeout=None)`` /
1233 ``parse_response(block=True)``) where the configured
1234 ``socket_timeout`` must not abort the read.
1235 - ``float``: apply that timeout in seconds for this single read.
1236
1237 TODO(next-major): replace the ``math.inf`` opt-in with a SENTINEL
1238 default for ``timeout``. After that change, ``timeout=None`` will
1239 mean "no timeout, block until a response arrives" (matching the
1240 long-standing PubSub docstring contract) and the SENTINEL default
1241 will be the value that falls back to ``self.socket_timeout``.
1242 That swap is a breaking change, so it must wait for a major
1243 release. Until then, callers that need an indefinitely blocking
1244 read pass ``math.inf`` explicitly.
1245 """
1246 # TODO(next-major): drop the math.inf branch. Use SENTINEL as the
1247 # default for ``timeout`` and treat ``timeout is None`` as the
1248 # "no timeout" signal (matching the PubSub docstring contract).
1249 # Match only positive infinity here. ``-math.inf`` is not a valid
1250 # "block forever" signal and historically behaved as an already-
1251 # expired timeout; preserve that.
1252 if timeout == math.inf:
1253 read_timeout = None
1254 else:
1255 read_timeout = timeout if timeout is not None else self.socket_timeout
1256 host_error = self._host_error()
1257 try:
1258 if read_timeout is not None:
1259 timeout_context = async_timeout(read_timeout)
1260 if timeout is None:
1261 async with timeout_context as active_timeout:
1262 self._active_read_timeout = active_timeout
1263 try:
1264 response = await self._read_response_from_parser(
1265 disable_decoding=disable_decoding,
1266 push_request=push_request,
1267 )
1268 finally:
1269 self._active_read_timeout = None
1270 else:
1271 async with timeout_context:
1272 response = await self._read_response_from_parser(
1273 disable_decoding=disable_decoding,
1274 push_request=push_request,
1275 )
1276 else:
1277 response = await self._read_response_from_parser(
1278 disable_decoding=disable_decoding,
1279 push_request=push_request,
1280 )
1281 except asyncio.TimeoutError:
1282 if timeout is not None:
1283 # user requested timeout, return None. Operation can be retried
1284 return None
1285 # it was a self.socket_timeout error.
1286 if disconnect_on_error:
1287 await self.disconnect(nowait=True)
1288 raise TimeoutError(f"Timeout reading from {host_error}")
1289 except OSError as e:
1290 if disconnect_on_error:
1291 await self.disconnect(nowait=True)
1292 raise ConnectionError(f"Error while reading from {host_error} : {e.args}")
1293 except BaseException:
1294 # Also by default close in case of BaseException. A lot of code
1295 # relies on this behaviour when doing Command/Response pairs.
1296 # See #1128.
1297 if disconnect_on_error:
1298 await self.disconnect(nowait=True)
1299 raise
1300
1301 if self.health_check_interval:
1302 next_time = asyncio.get_running_loop().time() + self.health_check_interval
1303 self.next_health_check = next_time
1304
1305 if isinstance(response, ResponseError):
1306 raise response from None
1307 return response
1308
1309 async def _read_response_from_parser(
1310 self, disable_decoding: bool = False, push_request: bool | None = False
1311 ):
1312 if check_protocol_version(self.protocol, 3):
1313 return await self._parser.read_response(
1314 disable_decoding=disable_decoding, push_request=push_request
1315 )
1316 return await self._parser.read_response(disable_decoding=disable_decoding)
1317
1318 def pack_command(self, *args: EncodableT) -> List[bytes]:
1319 """Pack a series of arguments into the Redis protocol"""
1320 output = []
1321 # the client might have included 1 or more literal arguments in
1322 # the command name, e.g., 'CONFIG GET'. The Redis server expects these
1323 # arguments to be sent separately, so split the first argument
1324 # manually. These arguments should be bytestrings so that they are
1325 # not encoded.
1326 assert not isinstance(args[0], float)
1327 if isinstance(args[0], str):
1328 args = tuple(args[0].encode().split()) + args[1:]
1329 elif b" " in args[0]:
1330 args = tuple(args[0].split()) + args[1:]
1331
1332 buff = SYM_EMPTY.join((SYM_STAR, str(len(args)).encode(), SYM_CRLF))
1333
1334 buffer_cutoff = self._buffer_cutoff
1335 for arg in map(self.encoder.encode, args):
1336 # to avoid large string mallocs, chunk the command into the
1337 # output list if we're sending large values or memoryviews
1338 arg_length = len(arg)
1339 if (
1340 len(buff) > buffer_cutoff
1341 or arg_length > buffer_cutoff
1342 or isinstance(arg, memoryview)
1343 ):
1344 buff = SYM_EMPTY.join(
1345 (buff, SYM_DOLLAR, str(arg_length).encode(), SYM_CRLF)
1346 )
1347 output.append(buff)
1348 output.append(arg)
1349 buff = SYM_CRLF
1350 else:
1351 buff = SYM_EMPTY.join(
1352 (
1353 buff,
1354 SYM_DOLLAR,
1355 str(arg_length).encode(),
1356 SYM_CRLF,
1357 arg,
1358 SYM_CRLF,
1359 )
1360 )
1361 output.append(buff)
1362 return output
1363
1364 def pack_commands(self, commands: Iterable[Iterable[EncodableT]]) -> List[bytes]:
1365 """Pack multiple commands into the Redis protocol"""
1366 output: List[bytes] = []
1367 pieces: List[bytes] = []
1368 buffer_length = 0
1369 buffer_cutoff = self._buffer_cutoff
1370
1371 for cmd in commands:
1372 for chunk in self.pack_command(*cmd):
1373 chunklen = len(chunk)
1374 if (
1375 buffer_length > buffer_cutoff
1376 or chunklen > buffer_cutoff
1377 or isinstance(chunk, memoryview)
1378 ):
1379 if pieces:
1380 output.append(SYM_EMPTY.join(pieces))
1381 buffer_length = 0
1382 pieces = []
1383
1384 if chunklen > buffer_cutoff or isinstance(chunk, memoryview):
1385 output.append(chunk)
1386 else:
1387 pieces.append(chunk)
1388 buffer_length += chunklen
1389
1390 if pieces:
1391 output.append(SYM_EMPTY.join(pieces))
1392 return output
1393
1394 def _socket_is_empty(self):
1395 """Check if the socket is empty"""
1396 return len(self._reader._buffer) == 0
1397
1398 async def process_invalidation_messages(self):
1399 while not self._socket_is_empty():
1400 await self.read_response(push_request=True)
1401
1402 def set_re_auth_token(self, token: TokenInterface):
1403 self._re_auth_token = token
1404
1405 async def re_auth(self):
1406 if self._re_auth_token is not None:
1407 await self.send_command(
1408 "AUTH",
1409 self._re_auth_token.try_get("oid"),
1410 self._re_auth_token.get_value(),
1411 )
1412 await self.read_response()
1413 self._re_auth_token = None
1414
1415
1416class Connection(AbstractConnection):
1417 "Manages TCP communication to and from a Redis server"
1418
1419 def __init__(
1420 self,
1421 *,
1422 host: str = "localhost",
1423 port: str | int = 6379,
1424 socket_keepalive: bool = True,
1425 socket_keepalive_options: Mapping[int, int | bytes] | object | None = SENTINEL,
1426 socket_type: int = 0,
1427 **kwargs,
1428 ):
1429 """
1430 Initialize a TCP connection.
1431
1432 Parameters
1433 ----------
1434 socket_keepalive : bool
1435 If `True`, TCP keepalive is enabled for TCP socket connections.
1436 socket_keepalive_options : Mapping[int, int | bytes] | object | None
1437 Mapping of TCP keepalive socket option constants to values, for
1438 example `{socket.TCP_KEEPIDLE: 30}`. If left unspecified, redis-py
1439 uses TCP keepalive defaults when `socket_keepalive` is enabled:
1440 idle 30 seconds, interval 5 seconds, and 3 probes. Platform-specific
1441 options that are not available are skipped. Pass `None` or `{}` to
1442 avoid setting additional TCP keepalive options.
1443 """
1444 self.host = host
1445 self.port = int(port)
1446 self.socket_keepalive = socket_keepalive
1447 if socket_keepalive_options is SENTINEL:
1448 socket_keepalive_options = get_default_socket_keepalive_options()
1449 self.socket_keepalive_options = socket_keepalive_options or {}
1450 self.socket_type = socket_type
1451 super().__init__(**kwargs)
1452
1453 def repr_pieces(self):
1454 pieces = [("host", self.host), ("port", self.port), ("db", self.db)]
1455 if self.client_name:
1456 pieces.append(("client_name", self.client_name))
1457 return pieces
1458
1459 def _connection_arguments(self) -> Mapping:
1460 return {"host": self.host, "port": self.port}
1461
1462 async def _connect(self):
1463 """Create a TCP socket connection"""
1464 async with async_timeout(self.socket_connect_timeout):
1465 reader, writer = await asyncio.open_connection(
1466 **self._connection_arguments()
1467 )
1468 self._reader = reader
1469 self._writer = writer
1470 sock = writer.transport.get_extra_info("socket")
1471 if sock:
1472 sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
1473 try:
1474 # TCP_KEEPALIVE
1475 if self.socket_keepalive:
1476 sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
1477 for k, v in self.socket_keepalive_options.items():
1478 sock.setsockopt(socket.SOL_TCP, k, v)
1479
1480 except (OSError, TypeError):
1481 # `socket_keepalive_options` might contain invalid options
1482 # causing an error. Do not leave the connection open.
1483 writer.close()
1484 raise
1485
1486 def _host_error(self) -> str:
1487 return f"{self.host}:{self.port}"
1488
1489
1490class SSLConnection(Connection):
1491 """Manages SSL connections to and from the Redis server(s).
1492 This class extends the Connection class, adding SSL functionality, and making
1493 use of ssl.SSLContext (https://docs.python.org/3/library/ssl.html#ssl.SSLContext)
1494 """
1495
1496 def __init__(
1497 self,
1498 ssl_keyfile: Optional[str] = None,
1499 ssl_certfile: Optional[str] = None,
1500 ssl_cert_reqs: Union[str, ssl.VerifyMode] = "required",
1501 ssl_include_verify_flags: Optional[List["ssl.VerifyFlags"]] = None,
1502 ssl_exclude_verify_flags: Optional[List["ssl.VerifyFlags"]] = None,
1503 ssl_ca_certs: Optional[str] = None,
1504 ssl_ca_data: Optional[str] = None,
1505 ssl_ca_path: Optional[str] = None,
1506 ssl_check_hostname: bool = True,
1507 ssl_min_version: Optional[TLSVersion] = None,
1508 ssl_ciphers: Optional[str] = None,
1509 ssl_password: Optional[str] = None,
1510 **kwargs,
1511 ):
1512 if not SSL_AVAILABLE:
1513 raise RedisError("Python wasn't built with SSL support")
1514
1515 self.ssl_context: RedisSSLContext = RedisSSLContext(
1516 keyfile=ssl_keyfile,
1517 certfile=ssl_certfile,
1518 cert_reqs=ssl_cert_reqs,
1519 include_verify_flags=ssl_include_verify_flags,
1520 exclude_verify_flags=ssl_exclude_verify_flags,
1521 ca_certs=ssl_ca_certs,
1522 ca_data=ssl_ca_data,
1523 ca_path=ssl_ca_path,
1524 check_hostname=ssl_check_hostname,
1525 min_version=ssl_min_version,
1526 ciphers=ssl_ciphers,
1527 password=ssl_password,
1528 )
1529 super().__init__(**kwargs)
1530
1531 def _connection_arguments(self) -> Mapping:
1532 kwargs = super()._connection_arguments()
1533 kwargs["ssl"] = self.ssl_context.get()
1534 return kwargs
1535
1536 @property
1537 def keyfile(self):
1538 return self.ssl_context.keyfile
1539
1540 @property
1541 def certfile(self):
1542 return self.ssl_context.certfile
1543
1544 @property
1545 def cert_reqs(self):
1546 return self.ssl_context.cert_reqs
1547
1548 @property
1549 def include_verify_flags(self):
1550 return self.ssl_context.include_verify_flags
1551
1552 @property
1553 def exclude_verify_flags(self):
1554 return self.ssl_context.exclude_verify_flags
1555
1556 @property
1557 def ca_certs(self):
1558 return self.ssl_context.ca_certs
1559
1560 @property
1561 def ca_data(self):
1562 return self.ssl_context.ca_data
1563
1564 @property
1565 def check_hostname(self):
1566 return self.ssl_context.check_hostname
1567
1568 @property
1569 def min_version(self):
1570 return self.ssl_context.min_version
1571
1572
1573class RedisSSLContext:
1574 __slots__ = (
1575 "keyfile",
1576 "certfile",
1577 "cert_reqs",
1578 "include_verify_flags",
1579 "exclude_verify_flags",
1580 "ca_certs",
1581 "ca_data",
1582 "ca_path",
1583 "context",
1584 "check_hostname",
1585 "min_version",
1586 "ciphers",
1587 "password",
1588 )
1589
1590 def __init__(
1591 self,
1592 keyfile: Optional[str] = None,
1593 certfile: Optional[str] = None,
1594 cert_reqs: Optional[Union[str, ssl.VerifyMode]] = None,
1595 include_verify_flags: Optional[List["ssl.VerifyFlags"]] = None,
1596 exclude_verify_flags: Optional[List["ssl.VerifyFlags"]] = None,
1597 ca_certs: Optional[str] = None,
1598 ca_data: Optional[str] = None,
1599 ca_path: Optional[str] = None,
1600 check_hostname: bool = False,
1601 min_version: Optional[TLSVersion] = None,
1602 ciphers: Optional[str] = None,
1603 password: Optional[str] = None,
1604 ):
1605 if not SSL_AVAILABLE:
1606 raise RedisError("Python wasn't built with SSL support")
1607
1608 self.keyfile = keyfile
1609 self.certfile = certfile
1610 if cert_reqs is None:
1611 cert_reqs = ssl.CERT_NONE
1612 elif isinstance(cert_reqs, str):
1613 CERT_REQS = { # noqa: N806
1614 "none": ssl.CERT_NONE,
1615 "optional": ssl.CERT_OPTIONAL,
1616 "required": ssl.CERT_REQUIRED,
1617 }
1618 if cert_reqs not in CERT_REQS:
1619 raise RedisError(
1620 f"Invalid SSL Certificate Requirements Flag: {cert_reqs}"
1621 )
1622 cert_reqs = CERT_REQS[cert_reqs]
1623 self.cert_reqs = cert_reqs
1624 self.include_verify_flags = include_verify_flags
1625 self.exclude_verify_flags = exclude_verify_flags
1626 self.ca_certs = ca_certs
1627 self.ca_data = ca_data
1628 self.ca_path = ca_path
1629 self.check_hostname = (
1630 check_hostname if self.cert_reqs != ssl.CERT_NONE else False
1631 )
1632 self.min_version = min_version
1633 self.ciphers = ciphers
1634 self.password = password
1635 self.context: Optional[SSLContext] = None
1636
1637 def get(self) -> SSLContext:
1638 if not self.context:
1639 context = ssl.create_default_context()
1640 context.check_hostname = self.check_hostname
1641 context.verify_mode = self.cert_reqs
1642 if self.include_verify_flags:
1643 for flag in self.include_verify_flags:
1644 context.verify_flags |= flag
1645 if self.exclude_verify_flags:
1646 for flag in self.exclude_verify_flags:
1647 context.verify_flags &= ~flag
1648 if self.certfile or self.keyfile:
1649 context.load_cert_chain(
1650 certfile=self.certfile,
1651 keyfile=self.keyfile,
1652 password=self.password,
1653 )
1654 if self.ca_certs or self.ca_data or self.ca_path:
1655 context.load_verify_locations(
1656 cafile=self.ca_certs, capath=self.ca_path, cadata=self.ca_data
1657 )
1658 if self.min_version is not None:
1659 context.minimum_version = self.min_version
1660 if self.ciphers is not None:
1661 context.set_ciphers(self.ciphers)
1662 self.context = context
1663 return self.context
1664
1665
1666class UnixDomainSocketConnection(AbstractConnection):
1667 "Manages UDS communication to and from a Redis server"
1668
1669 def __init__(self, *, path: str = "", **kwargs):
1670 self.path = path
1671 super().__init__(**kwargs)
1672
1673 def repr_pieces(self) -> Iterable[Tuple[str, Union[str, int]]]:
1674 pieces = [("path", self.path), ("db", self.db)]
1675 if self.client_name:
1676 pieces.append(("client_name", self.client_name))
1677 return pieces
1678
1679 async def _connect(self):
1680 async with async_timeout(self.socket_connect_timeout):
1681 reader, writer = await asyncio.open_unix_connection(path=self.path)
1682 self._reader = reader
1683 self._writer = writer
1684 await self.on_connect()
1685
1686 def _host_error(self) -> str:
1687 return self.path
1688
1689
1690FALSE_STRINGS = ("0", "F", "FALSE", "N", "NO")
1691
1692
1693def to_bool(value) -> Optional[bool]:
1694 if value is None or value == "":
1695 return None
1696 if isinstance(value, str) and value.upper() in FALSE_STRINGS:
1697 return False
1698 return bool(value)
1699
1700
1701def parse_ssl_verify_flags(value):
1702 # flags are passed in as a string representation of a list,
1703 # e.g. VERIFY_X509_STRICT, VERIFY_X509_PARTIAL_CHAIN
1704 verify_flags_str = value.replace("[", "").replace("]", "")
1705
1706 verify_flags = []
1707 for flag in verify_flags_str.split(","):
1708 flag = flag.strip()
1709 if not hasattr(VerifyFlags, flag):
1710 raise ValueError(f"Invalid ssl verify flag: {flag}")
1711 verify_flags.append(getattr(VerifyFlags, flag))
1712 return verify_flags
1713
1714
1715URL_QUERY_ARGUMENT_PARSERS: Mapping[str, Callable[..., object]] = MappingProxyType(
1716 {
1717 "db": int,
1718 "socket_timeout": float,
1719 "socket_connect_timeout": float,
1720 "socket_read_size": int,
1721 "socket_keepalive": to_bool,
1722 "retry_on_timeout": to_bool,
1723 "max_connections": int,
1724 "health_check_interval": int,
1725 "ssl_check_hostname": to_bool,
1726 "ssl_include_verify_flags": parse_ssl_verify_flags,
1727 "ssl_exclude_verify_flags": parse_ssl_verify_flags,
1728 "ssl_min_version": int,
1729 "timeout": float,
1730 "protocol": int,
1731 "legacy_responses": to_bool,
1732 }
1733)
1734
1735
1736class ConnectKwargs(TypedDict, total=False):
1737 username: str
1738 password: str
1739 connection_class: Type[AbstractConnection]
1740 host: str
1741 port: int
1742 db: int
1743 path: str
1744
1745
1746def parse_url(url: str) -> ConnectKwargs:
1747 parsed: ParseResult = urlparse(url)
1748 kwargs: ConnectKwargs = {}
1749
1750 for name, value_list in parse_qs(parsed.query).items():
1751 if value_list and len(value_list) > 0:
1752 value = unquote(value_list[0])
1753 parser = URL_QUERY_ARGUMENT_PARSERS.get(name)
1754 if parser:
1755 try:
1756 kwargs[name] = parser(value)
1757 except (TypeError, ValueError):
1758 raise ValueError(f"Invalid value for '{name}' in connection URL.")
1759 else:
1760 kwargs[name] = value
1761
1762 if parsed.username:
1763 kwargs["username"] = unquote(parsed.username)
1764 if parsed.password:
1765 kwargs["password"] = unquote(parsed.password)
1766
1767 # We only support redis://, rediss:// and unix:// schemes.
1768 if parsed.scheme == "unix":
1769 if parsed.path:
1770 kwargs["path"] = unquote(parsed.path)
1771 kwargs["connection_class"] = UnixDomainSocketConnection
1772
1773 elif parsed.scheme in ("redis", "rediss"):
1774 if parsed.hostname:
1775 kwargs["host"] = unquote(parsed.hostname)
1776 if parsed.port:
1777 kwargs["port"] = int(parsed.port)
1778
1779 # If there's a path argument, use it as the db argument if a
1780 # querystring value wasn't specified
1781 if parsed.path and "db" not in kwargs:
1782 try:
1783 kwargs["db"] = int(unquote(parsed.path).replace("/", ""))
1784 except (AttributeError, ValueError):
1785 pass
1786
1787 if parsed.scheme == "rediss":
1788 kwargs["connection_class"] = SSLConnection
1789
1790 else:
1791 valid_schemes = "redis://, rediss://, unix://"
1792 raise ValueError(
1793 f"Redis URL must specify one of the following schemes ({valid_schemes})"
1794 )
1795
1796 return kwargs
1797
1798
1799_CP = TypeVar("_CP", bound="ConnectionPool")
1800
1801
1802class ConnectionPoolInterface(ABC):
1803 @abstractmethod
1804 def get_protocol(self):
1805 pass
1806
1807 @abstractmethod
1808 def reset(self) -> None:
1809 pass
1810
1811 @abstractmethod
1812 @deprecated_args(
1813 args_to_warn=["*"],
1814 reason="Use get_connection() without args instead",
1815 version="5.3.0",
1816 )
1817 async def get_connection(
1818 self, command_name: Optional[str] = None, *keys: Any, **options: Any
1819 ) -> "AbstractConnection":
1820 pass
1821
1822 @abstractmethod
1823 def get_encoder(self) -> "Encoder":
1824 pass
1825
1826 @abstractmethod
1827 async def release(self, connection: "AbstractConnection") -> None:
1828 pass
1829
1830 @abstractmethod
1831 async def disconnect(self, inuse_connections: bool = True) -> None:
1832 pass
1833
1834 @abstractmethod
1835 async def aclose(self) -> None:
1836 pass
1837
1838 @abstractmethod
1839 def set_retry(self, retry: "Retry") -> None:
1840 pass
1841
1842 @abstractmethod
1843 async def re_auth_callback(self, token: TokenInterface) -> None:
1844 pass
1845
1846 @abstractmethod
1847 def get_connection_count(self) -> List[Tuple[int, dict]]:
1848 """
1849 Returns a connection count (both idle and in use).
1850 """
1851 pass
1852
1853
1854class AsyncMaintNotificationsAbstractConnectionPool:
1855 """
1856 Internal mixin for async maintenance notification pool wiring.
1857
1858 The handler owns notification policy.
1859 This mixin owns pool state mutation because `_available_connections`,
1860 `_in_use_connections`, `connection_kwargs`, and the non-reentrant `asyncio.Lock`
1861 all live on the pool.
1862 """
1863
1864 def __init__(
1865 self,
1866 maint_notifications_config: MaintNotificationsConfig | None = None,
1867 oss_cluster_maint_notifications_handler: (
1868 "AsyncOSSMaintNotificationsHandler | None"
1869 ) = None,
1870 **kwargs: Any,
1871 ) -> None:
1872 protocol = kwargs.get("protocol")
1873 is_protocol_supported = check_protocol_version(protocol, 3)
1874 is_connection_supported = self._maintenance_notifications_supported()
1875
1876 if (
1877 maint_notifications_config is None
1878 and is_protocol_supported
1879 and is_connection_supported
1880 ):
1881 maint_notifications_config = MaintNotificationsConfig()
1882
1883 if maint_notifications_config and maint_notifications_config.enabled:
1884 if not is_connection_supported:
1885 if maint_notifications_config.enabled is True:
1886 # Unix sockets do not have a host endpoint for CLIENT
1887 # MAINT_NOTIFICATIONS to describe.
1888 if "path" in self.connection_kwargs:
1889 raise RedisError(
1890 "Maintenance notifications are not supported for "
1891 "Unix domain socket connections"
1892 )
1893
1894 # Custom connection classes must inherit the async maintenance
1895 # mixin so handlers can update connection state safely.
1896 if not self._maintenance_notifications_connection_class_supported():
1897 connection_class = getattr(self, "connection_class", None)
1898 connection_class_name = getattr(
1899 connection_class, "__name__", connection_class
1900 )
1901 raise RedisError(
1902 "Maintenance notifications are not supported for "
1903 f"connection class {connection_class_name}"
1904 )
1905
1906 # TCP-like connections still need a host to identify the
1907 # endpoint that can move during maintenance.
1908 raise RedisError(
1909 "Maintenance notifications are not supported for connections "
1910 "without a host"
1911 )
1912 self._maint_notifications_pool_handler = None
1913 self._oss_cluster_maint_notifications_handler = None
1914 return
1915
1916 if not is_protocol_supported:
1917 raise RedisError(
1918 "Maintenance notifications handlers on connection are only supported with RESP version 3"
1919 )
1920
1921 if oss_cluster_maint_notifications_handler:
1922 self._oss_cluster_maint_notifications_handler = (
1923 oss_cluster_maint_notifications_handler
1924 )
1925 self._update_connection_kwargs_for_maint_notifications(
1926 oss_cluster_maint_notifications_handler=self._oss_cluster_maint_notifications_handler
1927 )
1928 self._maint_notifications_pool_handler = None
1929 else:
1930 self._oss_cluster_maint_notifications_handler = None
1931 self._maint_notifications_pool_handler = (
1932 AsyncMaintNotificationsPoolHandler(self, maint_notifications_config)
1933 )
1934 self._update_connection_kwargs_for_maint_notifications(
1935 maint_notifications_pool_handler=self._maint_notifications_pool_handler
1936 )
1937 else:
1938 self._maint_notifications_pool_handler = None
1939 self._oss_cluster_maint_notifications_handler = None
1940
1941 async def _on_close(self) -> None:
1942 """Hook invoked from the pool's ``aclose()`` before the pool is shut down."""
1943 if self._maint_notifications_pool_handler is not None:
1944 await self._maint_notifications_pool_handler.cancel_scheduled_tasks()
1945
1946 @property
1947 @abstractmethod
1948 def connection_kwargs(self) -> dict[str, Any]:
1949 pass
1950
1951 @connection_kwargs.setter
1952 @abstractmethod
1953 def connection_kwargs(self, value: dict[str, Any]) -> None:
1954 pass
1955
1956 @abstractmethod
1957 def _get_pool_lock(self) -> asyncio.Lock:
1958 pass
1959
1960 @abstractmethod
1961 def _get_free_connections(self) -> Iterable["AbstractConnection"]:
1962 pass
1963
1964 @abstractmethod
1965 def _get_in_use_connections(self) -> Iterable["AbstractConnection"]:
1966 pass
1967
1968 def _maintenance_notifications_supported(self) -> bool:
1969 if "path" in self.connection_kwargs:
1970 return False
1971 if not self._maintenance_notifications_connection_class_supported():
1972 return False
1973 return bool(self.connection_kwargs.get("host"))
1974
1975 def _maintenance_notifications_connection_class_supported(self) -> bool:
1976 connection_class = getattr(self, "connection_class", None)
1977 if connection_class is None:
1978 return False
1979 try:
1980 return issubclass(
1981 connection_class, AsyncMaintNotificationsAbstractConnection
1982 )
1983 except TypeError:
1984 return False
1985
1986 def maint_notifications_enabled(self):
1987 """
1988 Returns:
1989 True if the maintenance notifications are enabled, False otherwise.
1990 The maintenance notifications config is stored in the pool handler.
1991 If the pool handler is not set, the maintenance notifications are not enabled.
1992 """
1993 if self._oss_cluster_maint_notifications_handler:
1994 maint_notifications_config = (
1995 self._oss_cluster_maint_notifications_handler.config
1996 )
1997 else:
1998 maint_notifications_config = (
1999 self._maint_notifications_pool_handler.config
2000 if self._maint_notifications_pool_handler
2001 else None
2002 )
2003 return maint_notifications_config and maint_notifications_config.enabled
2004
2005 async def update_maint_notifications_config(
2006 self,
2007 maint_notifications_config: MaintNotificationsConfig,
2008 oss_cluster_maint_notifications_handler: (
2009 AsyncOSSMaintNotificationsHandler | None
2010 ) = None,
2011 ) -> None:
2012 """
2013 Updates the maintenance notifications configuration.
2014 This method should be called only if the pool was created
2015 without enabling the maintenance notifications and
2016 in a later point in time maintenance notifications
2017 are requested to be enabled.
2018 """
2019 if (
2020 self.maint_notifications_enabled()
2021 and not maint_notifications_config.enabled
2022 ):
2023 raise ValueError(
2024 "Cannot disable maintenance notifications after enabling them"
2025 )
2026
2027 if oss_cluster_maint_notifications_handler:
2028 self._oss_cluster_maint_notifications_handler = (
2029 oss_cluster_maint_notifications_handler
2030 )
2031 # OSS cluster mode and pool-handler mode are mutually exclusive
2032 # (see __init__). A pool created with the default RESP3 "auto"
2033 # config wires a pool handler before this method runs; clear it so
2034 # new and existing connections are not configured with both handlers.
2035 self._maint_notifications_pool_handler = None
2036 else:
2037 if (
2038 maint_notifications_config.enabled
2039 and not self._maintenance_notifications_supported()
2040 ):
2041 if maint_notifications_config.enabled is True:
2042 # Unix sockets do not have a host endpoint for CLIENT
2043 # MAINT_NOTIFICATIONS to describe.
2044 if "path" in self.connection_kwargs:
2045 raise RedisError(
2046 "Maintenance notifications are not supported for "
2047 "Unix domain socket connections"
2048 )
2049
2050 # Custom connection classes must inherit the async maintenance
2051 # mixin so handlers can update connection state safely.
2052 if not self._maintenance_notifications_connection_class_supported():
2053 connection_class = getattr(self, "connection_class", None)
2054 connection_class_name = getattr(
2055 connection_class, "__name__", connection_class
2056 )
2057 raise RedisError(
2058 "Maintenance notifications are not supported for "
2059 f"connection class {connection_class_name}"
2060 )
2061
2062 # TCP-like connections still need a host to identify the
2063 # endpoint that can move during maintenance.
2064 raise RedisError(
2065 "Maintenance notifications are not supported for connections "
2066 "without a host"
2067 )
2068 self._maint_notifications_pool_handler = None
2069 return
2070
2071 if self._oss_cluster_maint_notifications_handler:
2072 # Pool already in OSS cluster mode; update the OSS handler config
2073 # instead of creating a mutually-exclusive pool handler (which
2074 # would be silently ignored because the OSS handler wins priority
2075 # in both update helpers below).
2076 self._oss_cluster_maint_notifications_handler.config = (
2077 maint_notifications_config
2078 )
2079 elif not self._maint_notifications_pool_handler:
2080 self._maint_notifications_pool_handler = (
2081 AsyncMaintNotificationsPoolHandler(self, maint_notifications_config)
2082 )
2083 else:
2084 self._maint_notifications_pool_handler.config = (
2085 maint_notifications_config
2086 )
2087
2088 self._update_connection_kwargs_for_maint_notifications(
2089 maint_notifications_pool_handler=self._maint_notifications_pool_handler,
2090 oss_cluster_maint_notifications_handler=self._oss_cluster_maint_notifications_handler,
2091 )
2092 await self._update_maint_notifications_configs_for_connections(
2093 maint_notifications_pool_handler=self._maint_notifications_pool_handler,
2094 oss_cluster_maint_notifications_handler=self._oss_cluster_maint_notifications_handler,
2095 )
2096
2097 def _update_connection_kwargs_for_maint_notifications(
2098 self,
2099 maint_notifications_pool_handler: (
2100 AsyncMaintNotificationsPoolHandler | None
2101 ) = None,
2102 oss_cluster_maint_notifications_handler: (
2103 AsyncOSSMaintNotificationsHandler | None
2104 ) = None,
2105 ) -> None:
2106 """
2107 Update the connection kwargs for all future connections.
2108 """
2109 if not self.maint_notifications_enabled():
2110 return
2111
2112 if maint_notifications_pool_handler:
2113 self.connection_kwargs.update(
2114 {
2115 "maint_notifications_pool_handler": maint_notifications_pool_handler,
2116 "maint_notifications_config": maint_notifications_pool_handler.config,
2117 }
2118 )
2119 if oss_cluster_maint_notifications_handler:
2120 self.connection_kwargs.update(
2121 {
2122 "oss_cluster_maint_notifications_handler": oss_cluster_maint_notifications_handler,
2123 "maint_notifications_config": oss_cluster_maint_notifications_handler.config,
2124 }
2125 )
2126 # OSS cluster mode and pool-handler mode are mutually exclusive.
2127 # Drop any pool handler a default (RESP3 "auto") pool creation may
2128 # have wired so future connections are not configured with both.
2129 self.connection_kwargs.pop("maint_notifications_pool_handler", None)
2130
2131 # Store original connection parameters for maintenance notifications.
2132 if self.connection_kwargs.get("orig_host_address", None) is None:
2133 # If orig_host_address is None it means we haven't
2134 # configured the original values yet
2135 self.connection_kwargs.update(
2136 {
2137 "orig_host_address": self.connection_kwargs.get("host"),
2138 "orig_socket_timeout": self.connection_kwargs.get(
2139 "socket_timeout", DEFAULT_SOCKET_TIMEOUT
2140 ),
2141 "orig_socket_connect_timeout": self.connection_kwargs.get(
2142 "socket_connect_timeout", DEFAULT_SOCKET_CONNECT_TIMEOUT
2143 ),
2144 }
2145 )
2146
2147 async def _update_maint_notifications_configs_for_connections(
2148 self,
2149 maint_notifications_pool_handler: (
2150 AsyncMaintNotificationsPoolHandler | None
2151 ) = None,
2152 oss_cluster_maint_notifications_handler: (
2153 AsyncOSSMaintNotificationsHandler | None
2154 ) = None,
2155 ) -> None:
2156 """Update the maintenance notifications config for all connections in the pool."""
2157 async with self._get_pool_lock():
2158 for conn in list(self._get_free_connections()):
2159 if oss_cluster_maint_notifications_handler:
2160 conn.set_maint_notifications_cluster_handler_for_connection(
2161 oss_cluster_maint_notifications_handler
2162 )
2163 conn.maint_notifications_config = (
2164 oss_cluster_maint_notifications_handler.config
2165 )
2166 elif maint_notifications_pool_handler:
2167 conn.set_maint_notifications_pool_handler_for_connection(
2168 maint_notifications_pool_handler
2169 )
2170 conn.maint_notifications_config = (
2171 maint_notifications_pool_handler.config
2172 )
2173 else:
2174 raise ValueError(
2175 "Either maint_notifications_pool_handler or "
2176 "oss_cluster_maint_notifications_handler must be set"
2177 )
2178 await conn.disconnect()
2179
2180 for conn in list(self._get_in_use_connections()):
2181 if oss_cluster_maint_notifications_handler:
2182 # Use set_maint_notifications_cluster_handler_for_connection
2183 # (not _configure_maintenance_notifications) so the parser is
2184 # obtained from the connection itself. _configure_* requires a
2185 # parser argument and would raise here; it would also reset the
2186 # connection's orig_* settings, which is wrong for an in-use
2187 # (active) connection. This mirrors the idle-connection branch
2188 # above and the pool-handler branches.
2189 conn.set_maint_notifications_cluster_handler_for_connection(
2190 oss_cluster_maint_notifications_handler
2191 )
2192 conn.maint_notifications_config = (
2193 oss_cluster_maint_notifications_handler.config
2194 )
2195 elif maint_notifications_pool_handler:
2196 conn.set_maint_notifications_pool_handler_for_connection(
2197 maint_notifications_pool_handler
2198 )
2199 conn.maint_notifications_config = (
2200 maint_notifications_pool_handler.config
2201 )
2202 else:
2203 raise ValueError(
2204 "Either maint_notifications_pool_handler or "
2205 "oss_cluster_maint_notifications_handler must be set"
2206 )
2207 conn.mark_for_reconnect()
2208
2209 def _should_update_connection(
2210 self,
2211 conn: "AbstractConnection",
2212 matching_pattern: str = "connected_address",
2213 matching_address: str | None = None,
2214 matching_notification_hash: int | None = None,
2215 ) -> bool:
2216 """
2217 Check if the connection should be updated based on the matching criteria.
2218 """
2219 if matching_pattern == "connected_address":
2220 if matching_address and conn.getpeername() != matching_address:
2221 return False
2222 elif matching_pattern == "configured_address":
2223 if matching_address and conn.host != matching_address:
2224 return False
2225 elif matching_pattern == "notification_hash":
2226 if (
2227 matching_notification_hash is not None
2228 and conn.maintenance_notification_hash != matching_notification_hash
2229 ):
2230 return False
2231 return True
2232
2233 def update_connection_settings(
2234 self,
2235 conn: "AsyncMaintNotificationsAbstractConnection",
2236 state: MaintenanceState | None = None,
2237 maintenance_notification_hash: int | None = None,
2238 host_address: str | None = None,
2239 relaxed_timeout: float | None = None,
2240 update_notification_hash: bool = False,
2241 reset_host_address: bool = False,
2242 reset_relaxed_timeout: bool = False,
2243 ) -> None:
2244 """
2245 Update the settings for a single connection.
2246 """
2247 if state:
2248 conn.maintenance_state = state
2249
2250 if update_notification_hash:
2251 # update the notification hash only if requested
2252 conn.maintenance_notification_hash = maintenance_notification_hash
2253
2254 if host_address is not None:
2255 conn.set_tmp_settings(tmp_host_address=host_address)
2256
2257 if relaxed_timeout is not None:
2258 conn.set_tmp_settings(tmp_relaxed_timeout=relaxed_timeout)
2259
2260 if reset_relaxed_timeout or reset_host_address:
2261 conn.reset_tmp_settings(
2262 reset_host_address=reset_host_address,
2263 reset_relaxed_timeout=reset_relaxed_timeout,
2264 )
2265
2266 conn.update_current_socket_timeout(relaxed_timeout)
2267
2268 async def update_connections_settings(
2269 self,
2270 state: MaintenanceState | None = None,
2271 maintenance_notification_hash: int | None = None,
2272 host_address: str | None = None,
2273 relaxed_timeout: float | None = None,
2274 matching_address: str | None = None,
2275 matching_notification_hash: int | None = None,
2276 matching_pattern: Literal[
2277 "connected_address", "configured_address", "notification_hash"
2278 ] = "connected_address",
2279 update_notification_hash: bool = False,
2280 reset_host_address: bool = False,
2281 reset_relaxed_timeout: bool = False,
2282 include_free_connections: bool = True,
2283 ) -> None:
2284 """
2285 Update the settings for all matching connections in the pool.
2286
2287 This method does not create new connections.
2288 This method does not affect the connection kwargs.
2289
2290 :param state: The maintenance state to set for the connection.
2291 :param maintenance_notification_hash: The hash of the maintenance notification
2292 to set for the connection.
2293 :param host_address: The host address to set for the connection.
2294 :param relaxed_timeout: The relaxed timeout to set for the connection.
2295 :param matching_address: The address to match for the connection.
2296 :param matching_notification_hash: The notification hash to match for the connection.
2297 :param matching_pattern: The pattern to match for the connection.
2298 :param update_notification_hash: Whether to update the notification hash for the connection.
2299 :param reset_host_address: Whether to reset the host address to the original address.
2300 :param reset_relaxed_timeout: Whether to reset the relaxed timeout to the original timeout.
2301 :param include_free_connections: Whether to include free/available connections.
2302 """
2303 async with self._get_pool_lock():
2304 self._update_connections_settings_without_locking(
2305 state=state,
2306 maintenance_notification_hash=maintenance_notification_hash,
2307 host_address=host_address,
2308 relaxed_timeout=relaxed_timeout,
2309 matching_address=matching_address,
2310 matching_notification_hash=matching_notification_hash,
2311 matching_pattern=matching_pattern,
2312 update_notification_hash=update_notification_hash,
2313 reset_host_address=reset_host_address,
2314 reset_relaxed_timeout=reset_relaxed_timeout,
2315 include_free_connections=include_free_connections,
2316 )
2317
2318 def _update_connections_settings_without_locking(
2319 self,
2320 state: MaintenanceState | None = None,
2321 maintenance_notification_hash: int | None = None,
2322 host_address: str | None = None,
2323 relaxed_timeout: float | None = None,
2324 matching_address: str | None = None,
2325 matching_notification_hash: int | None = None,
2326 matching_pattern: Literal[
2327 "connected_address", "configured_address", "notification_hash"
2328 ] = "connected_address",
2329 update_notification_hash: bool = False,
2330 reset_host_address: bool = False,
2331 reset_relaxed_timeout: bool = False,
2332 include_free_connections: bool = True,
2333 ) -> None:
2334 """
2335 Update matching connections while the caller already holds the pool lock.
2336
2337 This helper intentionally does not acquire the pool lock so callers can
2338 compose several pool mutations inside one critical section without
2339 deadlocking the non-reentrant `asyncio.Lock`.
2340 """
2341 for conn in self._get_in_use_connections():
2342 if self._should_update_connection(
2343 conn,
2344 matching_pattern,
2345 matching_address,
2346 matching_notification_hash,
2347 ):
2348 self.update_connection_settings(
2349 conn,
2350 state=state,
2351 maintenance_notification_hash=maintenance_notification_hash,
2352 host_address=host_address,
2353 relaxed_timeout=relaxed_timeout,
2354 update_notification_hash=update_notification_hash,
2355 reset_host_address=reset_host_address,
2356 reset_relaxed_timeout=reset_relaxed_timeout,
2357 )
2358
2359 if include_free_connections:
2360 for conn in self._get_free_connections():
2361 if self._should_update_connection(
2362 conn,
2363 matching_pattern,
2364 matching_address,
2365 matching_notification_hash,
2366 ):
2367 self.update_connection_settings(
2368 conn,
2369 state=state,
2370 maintenance_notification_hash=maintenance_notification_hash,
2371 host_address=host_address,
2372 relaxed_timeout=relaxed_timeout,
2373 update_notification_hash=update_notification_hash,
2374 reset_host_address=reset_host_address,
2375 reset_relaxed_timeout=reset_relaxed_timeout,
2376 )
2377
2378 def update_connection_kwargs(self, **kwargs: Any) -> None:
2379 """
2380 Update the connection kwargs for all future connections.
2381
2382 This method updates the connection kwargs for all future connections created by the pool.
2383 Existing connections are not affected.
2384 """
2385 self.connection_kwargs.update(kwargs)
2386
2387 async def apply_moving_notification(
2388 self,
2389 notification: NodeMovingNotification,
2390 config: MaintNotificationsConfig,
2391 moving_address_src: str | None,
2392 run_proactive_reconnect: bool = False,
2393 ) -> None:
2394 """
2395 Apply the pool state transition for a MOVING notification atomically.
2396
2397 Async pools use a non-reentrant `asyncio.Lock`, so the handler cannot
2398 safely compose several separately locked calls. Existing connection
2399 updates, optional proactive reconnect, and future `connection_kwargs`
2400 changes must happen under one pool-owned lock; otherwise a connection
2401 can move between active/free lists and escape handling.
2402 """
2403 async with self._get_pool_lock():
2404 # Opt BlockingConnectionPool into serializing its get/release
2405 # with this critical section. Other pools do not define
2406 # set_in_maintenance and this is a no-op for them.
2407 self._set_in_maintenance(True)
2408 try:
2409 self._update_connections_settings_without_locking(
2410 state=MaintenanceState.MOVING,
2411 maintenance_notification_hash=hash(notification),
2412 relaxed_timeout=config.relaxed_timeout,
2413 host_address=notification.new_node_host,
2414 matching_address=moving_address_src,
2415 matching_pattern="connected_address",
2416 update_notification_hash=True,
2417 include_free_connections=True,
2418 )
2419
2420 if run_proactive_reconnect:
2421 await self._run_proactive_reconnect_without_locking(
2422 moving_address_src
2423 )
2424
2425 self.update_connection_kwargs(
2426 **_build_moving_connection_kwargs(notification, config)
2427 )
2428 finally:
2429 self._set_in_maintenance(False)
2430
2431 async def run_proactive_reconnect(
2432 self,
2433 moving_address_src: str | None = None,
2434 ) -> None:
2435 """
2436 Mark active connections and disconnect free connections atomically.
2437
2438 This operation is pool-owned because the active/free lists can change
2439 while tasks acquire or release connections. Keeping the mark/disconnect
2440 pass under one lock avoids a connection moving between lists between
2441 separately locked calls.
2442 """
2443 async with self._get_pool_lock():
2444 await self._run_proactive_reconnect_without_locking(moving_address_src)
2445
2446 async def _run_proactive_reconnect_without_locking(
2447 self,
2448 moving_address_src: str | None = None,
2449 ) -> None:
2450 """
2451 Mark and disconnect matching connections while the caller holds the pool lock.
2452
2453 This helper intentionally does not acquire the pool lock so it can be
2454 reused by larger atomic operations that already hold the non-reentrant
2455 `asyncio.Lock`.
2456 """
2457 for conn in self._get_in_use_connections():
2458 if self._should_update_connection(
2459 conn, "connected_address", moving_address_src
2460 ):
2461 conn.mark_for_reconnect()
2462
2463 free_connections = [
2464 conn
2465 for conn in self._get_free_connections()
2466 if self._should_update_connection(
2467 conn, "connected_address", moving_address_src
2468 )
2469 ]
2470 await self._disconnect_connections(free_connections)
2471
2472 async def cleanup_moving_notification(
2473 self,
2474 notification_hash: int,
2475 reset_relaxed_timeout: bool,
2476 reset_host_address: bool,
2477 ) -> None:
2478 """
2479 Revert MOVING pool state atomically after the notification TTL.
2480
2481 Future connection kwargs and existing connection state must be cleaned
2482 up in the same critical section. Splitting the cleanup lets an
2483 acquire/release interleave, which can leave stale MOVING state or undo a
2484 newer overlapping MOVING notification.
2485 """
2486 async with self._get_pool_lock():
2487 kwargs = _build_moving_cleanup_connection_kwargs(
2488 self.connection_kwargs, notification_hash
2489 )
2490 if kwargs is not None:
2491 self.update_connection_kwargs(**kwargs)
2492
2493 self._update_connections_settings_without_locking(
2494 relaxed_timeout=-1,
2495 state=MaintenanceState.NONE,
2496 maintenance_notification_hash=None,
2497 matching_notification_hash=notification_hash,
2498 matching_pattern="notification_hash",
2499 update_notification_hash=True,
2500 reset_relaxed_timeout=reset_relaxed_timeout,
2501 reset_host_address=reset_host_address,
2502 include_free_connections=True,
2503 )
2504
2505 async def _disconnect_connections(
2506 self, connections: Iterable["AbstractConnection"]
2507 ) -> None:
2508 connections = tuple(connections)
2509 if not connections:
2510 return
2511 results = await asyncio.gather(
2512 *(connection.disconnect() for connection in connections),
2513 return_exceptions=True,
2514 )
2515 exc = next(
2516 (result for result in results if isinstance(result, BaseException)), None
2517 )
2518 if exc:
2519 raise exc
2520
2521 def _set_in_maintenance(self, in_maintenance: bool) -> None:
2522 """Flip the pool's maintenance flag if it exposes one (BlockingConnectionPool)."""
2523 set_in_maintenance = getattr(self, "set_in_maintenance", None)
2524 if callable(set_in_maintenance):
2525 set_in_maintenance(in_maintenance)
2526
2527
2528class ConnectionPool(
2529 AsyncMaintNotificationsAbstractConnectionPool, ConnectionPoolInterface
2530):
2531 """
2532 Create a connection pool. ``If max_connections`` is set, then this
2533 object raises :py:class:`~redis.ConnectionError` when the pool's
2534 limit is reached.
2535
2536 By default, TCP connections are created unless ``connection_class``
2537 is specified. Use :py:class:`~redis.UnixDomainSocketConnection` for
2538 unix sockets.
2539 :py:class:`~redis.SSLConnection` can be used for SSL enabled connections.
2540
2541 Any additional keyword arguments are passed to the constructor of
2542 ``connection_class``.
2543 """
2544
2545 @classmethod
2546 def from_url(cls: Type[_CP], url: str, **kwargs) -> _CP:
2547 """
2548 Return a connection pool configured from the given URL.
2549
2550 For example::
2551
2552 redis://[[username]:[password]]@localhost:6379/0
2553 rediss://[[username]:[password]]@localhost:6379/0
2554 unix://[username@]/path/to/socket.sock?db=0[&password=password]
2555
2556 Three URL schemes are supported:
2557
2558 - `redis://` creates a TCP socket connection. See more at:
2559 <https://www.iana.org/assignments/uri-schemes/prov/redis>
2560 - `rediss://` creates a SSL wrapped TCP socket connection. See more at:
2561 <https://www.iana.org/assignments/uri-schemes/prov/rediss>
2562 - ``unix://``: creates a Unix Domain Socket connection.
2563
2564 The username, password, hostname, path and all querystring values
2565 are passed through urllib.parse.unquote in order to replace any
2566 percent-encoded values with their corresponding characters.
2567
2568 There are several ways to specify a database number. The first value
2569 found will be used:
2570
2571 1. A ``db`` querystring option, e.g. redis://localhost?db=0
2572
2573 2. If using the redis:// or rediss:// schemes, the path argument
2574 of the url, e.g. redis://localhost/0
2575
2576 3. A ``db`` keyword argument to this function.
2577
2578 If none of these options are specified, the default db=0 is used.
2579
2580 All querystring options are cast to their appropriate Python types.
2581 Boolean arguments can be specified with string values "True"/"False"
2582 or "Yes"/"No". Values that cannot be properly cast cause a
2583 ``ValueError`` to be raised. Once parsed, the querystring arguments
2584 and keyword arguments are passed to the ``ConnectionPool``'s
2585 class initializer. In the case of conflicting arguments, querystring
2586 arguments always win.
2587 """
2588 url_options = parse_url(url)
2589 kwargs.update(url_options)
2590 return cls(**kwargs)
2591
2592 def __init__(
2593 self,
2594 connection_class: Type[AbstractConnection] = Connection,
2595 max_connections: Optional[int] = None,
2596 maint_notifications_config: MaintNotificationsConfig | None = None,
2597 **connection_kwargs,
2598 ):
2599 max_connections = max_connections or 100
2600 if not isinstance(max_connections, int) or max_connections < 0:
2601 raise ValueError('"max_connections" must be a positive integer')
2602
2603 self.connection_class = connection_class
2604 self._connection_kwargs = connection_kwargs
2605 self.max_connections = max_connections
2606
2607 self._available_connections: List[AbstractConnection] = []
2608 self._in_use_connections: Set[AbstractConnection] = set()
2609 self.encoder_class = self.connection_kwargs.get("encoder_class", Encoder)
2610 self._lock = asyncio.Lock()
2611 self._event_dispatcher = self.connection_kwargs.get("event_dispatcher", None)
2612 if self._event_dispatcher is None:
2613 self._event_dispatcher = EventDispatcher()
2614
2615 AsyncMaintNotificationsAbstractConnectionPool.__init__(
2616 self,
2617 maint_notifications_config=maint_notifications_config,
2618 **connection_kwargs,
2619 )
2620
2621 # Keys that should be redacted in __repr__ to avoid exposing sensitive information
2622 SENSITIVE_REPR_KEYS = frozenset(
2623 {
2624 "password",
2625 "username",
2626 "ssl_password",
2627 "credential_provider",
2628 }
2629 )
2630
2631 def __repr__(self):
2632 conn_kwargs = ",".join(
2633 [
2634 f"{k}={'<REDACTED>' if k in self.SENSITIVE_REPR_KEYS else v}"
2635 for k, v in self.connection_kwargs.items()
2636 ]
2637 )
2638 return (
2639 f"<{self.__class__.__module__}.{self.__class__.__name__}"
2640 f"(<{self.connection_class.__module__}.{self.connection_class.__name__}"
2641 f"({conn_kwargs})>)>"
2642 )
2643
2644 @property
2645 def connection_kwargs(self) -> dict[str, Any]:
2646 return self._connection_kwargs
2647
2648 @connection_kwargs.setter
2649 def connection_kwargs(self, value: dict[str, Any]) -> None:
2650 self._connection_kwargs = value
2651
2652 def _get_pool_lock(self) -> asyncio.Lock:
2653 return self._lock
2654
2655 def _get_free_connections(self) -> Iterable[AbstractConnection]:
2656 return self._available_connections
2657
2658 def _get_in_use_connections(self) -> Iterable[AbstractConnection]:
2659 return self._in_use_connections
2660
2661 def get_protocol(self):
2662 """
2663 Returns:
2664 The RESP protocol version, or ``None`` if the protocol is not specified,
2665 in which case the server default will be used.
2666 """
2667 return self.connection_kwargs.get("protocol", None)
2668
2669 def reset(self):
2670 # Record metrics for connections being removed before clearing
2671 # (only if attributes exist - they won't during __init__)
2672 if hasattr(self, "_available_connections") and hasattr(
2673 self, "_in_use_connections"
2674 ):
2675 idle_count = len(self._available_connections)
2676 in_use_count = len(self._in_use_connections)
2677 if idle_count > 0 or in_use_count > 0:
2678 pool_name = get_pool_name(self)
2679 # Note: Using sync version since reset() is sync
2680 from redis.observability.recorder import (
2681 record_connection_count as sync_record_connection_count,
2682 )
2683
2684 if idle_count > 0:
2685 sync_record_connection_count(
2686 pool_name=pool_name,
2687 connection_state=ConnectionState.IDLE,
2688 counter=-idle_count,
2689 )
2690 if in_use_count > 0:
2691 sync_record_connection_count(
2692 pool_name=pool_name,
2693 connection_state=ConnectionState.USED,
2694 counter=-in_use_count,
2695 )
2696
2697 self._available_connections = []
2698 self._in_use_connections = weakref.WeakSet()
2699
2700 def __del__(self) -> None:
2701 """Clean up connection pool and record metrics when garbage collected."""
2702 try:
2703 if not hasattr(self, "_available_connections") or not hasattr(
2704 self, "_in_use_connections"
2705 ):
2706 return
2707 idle_count = len(self._available_connections)
2708 in_use_count = len(self._in_use_connections)
2709 if idle_count > 0 or in_use_count > 0:
2710 pool_name = get_pool_name(self)
2711 # Note: Using sync version since __del__ is sync
2712 from redis.observability.recorder import (
2713 record_connection_count as sync_record_connection_count,
2714 )
2715
2716 if idle_count > 0:
2717 sync_record_connection_count(
2718 pool_name=pool_name,
2719 connection_state=ConnectionState.IDLE,
2720 counter=-idle_count,
2721 )
2722 if in_use_count > 0:
2723 sync_record_connection_count(
2724 pool_name=pool_name,
2725 connection_state=ConnectionState.USED,
2726 counter=-in_use_count,
2727 )
2728 except Exception:
2729 pass
2730
2731 def can_get_connection(self) -> bool:
2732 """Return True if a connection can be retrieved from the pool."""
2733 return (
2734 self._available_connections
2735 or len(self._in_use_connections) < self.max_connections
2736 )
2737
2738 @deprecated_args(
2739 args_to_warn=["*"],
2740 reason="Use get_connection() without args instead",
2741 version="5.3.0",
2742 )
2743 async def get_connection(self, command_name=None, *keys, **options):
2744 """Get a connected connection from the pool"""
2745 # Track connection count before to detect if a new connection is created
2746 async with self._lock:
2747 connections_before = len(self._available_connections) + len(
2748 self._in_use_connections
2749 )
2750 start_time_created = time.monotonic()
2751 connection = self.get_available_connection()
2752 connections_after = len(self._available_connections) + len(
2753 self._in_use_connections
2754 )
2755 is_created = connections_after > connections_before
2756
2757 # Record state transition for observability
2758 # This ensures counters stay balanced if ensure_connection() fails and release() is called
2759 pool_name = get_pool_name(self)
2760 if is_created:
2761 # New connection created and acquired: just USED +1
2762 await record_connection_count(
2763 pool_name=pool_name,
2764 connection_state=ConnectionState.USED,
2765 counter=1,
2766 )
2767 else:
2768 # Existing connection acquired from pool: IDLE -> USED
2769 await record_connection_count(
2770 pool_name=pool_name,
2771 connection_state=ConnectionState.IDLE,
2772 counter=-1,
2773 )
2774 await record_connection_count(
2775 pool_name=pool_name,
2776 connection_state=ConnectionState.USED,
2777 counter=1,
2778 )
2779
2780 # We now perform the connection check outside of the lock.
2781 try:
2782 await self.ensure_connection(connection)
2783
2784 if is_created:
2785 await record_connection_create_time(
2786 connection_pool=self,
2787 duration_seconds=time.monotonic() - start_time_created,
2788 )
2789
2790 return connection
2791 except BaseException:
2792 await self.release(connection)
2793 raise
2794
2795 def get_available_connection(self):
2796 """Get a connection from the pool, without making sure it is connected"""
2797 try:
2798 connection = self._available_connections.pop()
2799 except IndexError:
2800 if len(self._in_use_connections) >= self.max_connections:
2801 raise MaxConnectionsError("Too many connections") from None
2802 connection = self.make_connection()
2803 self._in_use_connections.add(connection)
2804 return connection
2805
2806 def get_encoder(self):
2807 """Return an encoder based on encoding settings"""
2808 kwargs = self.connection_kwargs
2809 return self.encoder_class(
2810 encoding=kwargs.get("encoding", "utf-8"),
2811 encoding_errors=kwargs.get("encoding_errors", "strict"),
2812 decode_responses=kwargs.get("decode_responses", False),
2813 )
2814
2815 def make_connection(self):
2816 """Create a new connection. Can be overridden by child classes."""
2817 # Note: We don't record IDLE here because async uses a sync make_connection
2818 # but async record_connection_count. The recording is handled in get_connection.
2819 return self.connection_class(**self.connection_kwargs)
2820
2821 async def ensure_connection(self, connection: AbstractConnection):
2822 """Ensure that the connection object is connected and valid"""
2823 await connection.connect()
2824 # connections that the pool provides should be ready to send
2825 # a command. if not, the connection was either returned to the
2826 # pool before all data has been read or the socket has been
2827 # closed. either way, reconnect and verify everything is good.
2828 try:
2829 if await connection.can_read() and not self.maint_notifications_enabled():
2830 raise ConnectionError("Connection has data") from None
2831 except (ConnectionError, TimeoutError, OSError):
2832 await connection.disconnect()
2833 await connection.connect()
2834 if await connection.can_read() and not self.maint_notifications_enabled():
2835 raise ConnectionError("Connection not ready") from None
2836
2837 async def release(self, connection: AbstractConnection):
2838 """Releases the connection back to the pool"""
2839 # Connections should always be returned to the correct pool,
2840 # not doing so is an error that will cause an exception here.
2841 async with self._lock:
2842 self._in_use_connections.remove(connection)
2843
2844 if connection.should_reconnect():
2845 await connection.disconnect()
2846
2847 self._available_connections.append(connection)
2848
2849 await self._event_dispatcher.dispatch_async(
2850 AsyncAfterConnectionReleasedEvent(connection)
2851 )
2852
2853 # Record state transition: USED -> IDLE
2854 pool_name = get_pool_name(self)
2855 await record_connection_count(
2856 pool_name=pool_name,
2857 connection_state=ConnectionState.USED,
2858 counter=-1,
2859 )
2860 await record_connection_count(
2861 pool_name=pool_name,
2862 connection_state=ConnectionState.IDLE,
2863 counter=1,
2864 )
2865
2866 async def disconnect(self, inuse_connections: bool = True):
2867 """
2868 Disconnects connections in the pool
2869
2870 If ``inuse_connections`` is True, disconnect connections that are
2871 current in use, potentially by other tasks. Otherwise only disconnect
2872 connections that are idle in the pool.
2873 """
2874 if inuse_connections:
2875 connections: Iterable[AbstractConnection] = chain(
2876 self._available_connections, self._in_use_connections
2877 )
2878 else:
2879 connections = self._available_connections
2880 resp = await asyncio.gather(
2881 *(connection.disconnect() for connection in connections),
2882 return_exceptions=True,
2883 )
2884
2885 exc = next((r for r in resp if isinstance(r, BaseException)), None)
2886 if exc:
2887 raise exc
2888
2889 async def update_active_connections_for_reconnect(self):
2890 """
2891 Mark all active connections for reconnect.
2892 """
2893 async with self._lock:
2894 for conn in self._in_use_connections:
2895 conn.mark_for_reconnect()
2896
2897 async def aclose(self) -> None:
2898 """Close the pool, disconnecting all connections"""
2899 await self._on_close()
2900 await self.disconnect()
2901
2902 async def __aenter__(self: _CP) -> _CP:
2903 return self
2904
2905 async def __aexit__(self, exc_type, exc_value, traceback) -> None:
2906 await self.aclose()
2907
2908 def set_retry(self, retry: "Retry") -> None:
2909 for conn in self._available_connections:
2910 conn.retry = retry
2911 for conn in self._in_use_connections:
2912 conn.retry = retry
2913
2914 async def re_auth_callback(self, token: TokenInterface):
2915 async with self._lock:
2916 for conn in self._available_connections:
2917 await conn.retry.call_with_retry(
2918 lambda: conn.send_command(
2919 "AUTH", token.try_get("oid"), token.get_value()
2920 ),
2921 lambda error: self._mock(error),
2922 )
2923 await conn.retry.call_with_retry(
2924 lambda: conn.read_response(), lambda error: self._mock(error)
2925 )
2926 for conn in self._in_use_connections:
2927 conn.set_re_auth_token(token)
2928
2929 async def _mock(self, error: RedisError):
2930 """
2931 Dummy functions, needs to be passed as error callback to retry object.
2932 :param error:
2933 :return:
2934 """
2935 pass
2936
2937 def get_connection_count(self) -> List[tuple[int, dict]]:
2938 """
2939 Returns a connection count (both idle and in use).
2940 """
2941 attributes = AttributeBuilder.build_base_attributes()
2942 attributes[DB_CLIENT_CONNECTION_POOL_NAME] = get_pool_name(self)
2943 free_connections_attributes = attributes.copy()
2944 in_use_connections_attributes = attributes.copy()
2945
2946 free_connections_attributes[DB_CLIENT_CONNECTION_STATE] = (
2947 ConnectionState.IDLE.value
2948 )
2949 in_use_connections_attributes[DB_CLIENT_CONNECTION_STATE] = (
2950 ConnectionState.USED.value
2951 )
2952
2953 return [
2954 (len(self._available_connections), free_connections_attributes),
2955 (len(self._in_use_connections), in_use_connections_attributes),
2956 ]
2957
2958
2959class BlockingConnectionPool(ConnectionPool):
2960 """
2961 A blocking connection pool::
2962
2963 >>> from redis.asyncio import Redis, BlockingConnectionPool
2964 >>> client = Redis.from_pool(BlockingConnectionPool())
2965
2966 It performs the same function as the default
2967 :py:class:`~redis.asyncio.ConnectionPool` implementation, in that,
2968 it maintains a pool of reusable connections that can be shared by
2969 multiple async redis clients.
2970
2971 The difference is that, in the event that a client tries to get a
2972 connection from the pool when all of connections are in use, rather than
2973 raising a :py:class:`~redis.ConnectionError` (as the default
2974 :py:class:`~redis.asyncio.ConnectionPool` implementation does), it
2975 blocks the current `Task` for a specified number of seconds until
2976 a connection becomes available.
2977
2978 Use ``max_connections`` to increase / decrease the pool size::
2979
2980 >>> pool = BlockingConnectionPool(max_connections=10)
2981
2982 Use ``timeout`` to tell it either how many seconds to wait for a connection
2983 to become available, or to block forever:
2984
2985 >>> # Block forever.
2986 >>> pool = BlockingConnectionPool(timeout=None)
2987
2988 >>> # Raise a ``ConnectionError`` after five seconds if a connection is
2989 >>> # not available.
2990 >>> pool = BlockingConnectionPool(timeout=5)
2991 """
2992
2993 def __init__(
2994 self,
2995 max_connections: int = 50,
2996 timeout: Optional[float] = 20,
2997 connection_class: Type[AbstractConnection] = Connection,
2998 queue_class: Type[asyncio.Queue] = asyncio.LifoQueue, # deprecated
2999 **connection_kwargs,
3000 ):
3001 super().__init__(
3002 connection_class=connection_class,
3003 max_connections=max_connections,
3004 **connection_kwargs,
3005 )
3006 self._condition = asyncio.Condition()
3007 self.timeout = timeout
3008 self._in_maintenance = False
3009
3010 def set_in_maintenance(self, in_maintenance: bool) -> None:
3011 """
3012 Toggle the pool's maintenance mode.
3013
3014 While maintenance mode is on, ``get_connection`` and ``release``
3015 serialize their pool mutations through ``self._lock`` so they cannot
3016 interleave with a MOVING notification handler that is currently
3017 rewriting pool state under the same lock. Outside of maintenance the
3018 mutations skip the lock, since their critical sections are pure-Python
3019 and already atomic under asyncio's single-threaded scheduling.
3020 """
3021 self._in_maintenance = in_maintenance
3022
3023 @contextlib.asynccontextmanager
3024 async def _maybe_pool_lock(self) -> AsyncIterator[None]:
3025 if self._in_maintenance:
3026 async with self._lock:
3027 yield
3028 else:
3029 yield
3030
3031 @deprecated_args(
3032 args_to_warn=["*"],
3033 reason="Use get_connection() without args instead",
3034 version="5.3.0",
3035 )
3036 async def get_connection(self, command_name=None, *keys, **options):
3037 """Gets a connection from the pool, blocking until one is available"""
3038 # Start timing for wait time observability
3039 start_time_acquired = time.monotonic()
3040
3041 try:
3042 async with self._condition:
3043 async with async_timeout(self.timeout):
3044 await self._condition.wait_for(self.can_get_connection)
3045 async with self._maybe_pool_lock():
3046 # Track connection count before to detect if a new connection is created
3047 connections_before = len(self._available_connections) + len(
3048 self._in_use_connections
3049 )
3050 start_time_created = time.monotonic()
3051 connection = super().get_available_connection()
3052 connections_after = len(self._available_connections) + len(
3053 self._in_use_connections
3054 )
3055 is_created = connections_after > connections_before
3056 except asyncio.TimeoutError as err:
3057 raise ConnectionError("No connection available.") from err
3058
3059 # We now perform the connection check outside of the lock.
3060 try:
3061 await self.ensure_connection(connection)
3062
3063 if is_created:
3064 await record_connection_create_time(
3065 connection_pool=self,
3066 duration_seconds=time.monotonic() - start_time_created,
3067 )
3068
3069 await record_connection_wait_time(
3070 pool_name=get_pool_name(self),
3071 duration_seconds=time.monotonic() - start_time_acquired,
3072 )
3073
3074 return connection
3075 except BaseException:
3076 await self.release(connection)
3077 raise
3078
3079 async def release(self, connection: AbstractConnection):
3080 """Releases the connection back to the pool."""
3081 async with self._condition:
3082 await super().release(connection)
3083 self._condition.notify()