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