Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/redis/connection.py: 23%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1import copy
2import os
3import socket
4import sys
5import threading
6import time
7import weakref
8from abc import ABC, abstractmethod
9from itertools import chain
10from queue import Empty, Full, LifoQueue
11from typing import (
12 Any,
13 Callable,
14 Dict,
15 Iterable,
16 List,
17 Literal,
18 Optional,
19 Type,
20 TypeVar,
21 Union,
22)
23from urllib.parse import parse_qs, unquote, urlparse
25from redis.cache import (
26 CacheEntry,
27 CacheEntryStatus,
28 CacheFactory,
29 CacheFactoryInterface,
30 CacheInterface,
31 CacheKey,
32 CacheProxy,
33)
35from ._defaults import (
36 DEFAULT_SOCKET_CONNECT_TIMEOUT,
37 DEFAULT_SOCKET_READ_SIZE,
38 DEFAULT_SOCKET_TIMEOUT,
39 get_default_socket_keepalive_options,
40)
41from ._parsers import Encoder, _HiredisParser, _RESP2Parser, _RESP3Parser
42from .auth.token import TokenInterface
43from .backoff import NoBackoff
44from .credentials import CredentialProvider, UsernamePasswordCredentialProvider
45from .driver_info import DriverInfo, resolve_driver_info
46from .event import AfterConnectionReleasedEvent, EventDispatcher
47from .exceptions import (
48 AuthenticationError,
49 AuthenticationWrongNumberOfArgsError,
50 ChildDeadlockedError,
51 ConnectionError,
52 DataError,
53 MaxConnectionsError,
54 RedisError,
55 ResponseError,
56 TimeoutError,
57)
58from .maint_notifications import (
59 MaintenanceState,
60 MaintNotificationsConfig,
61 MaintNotificationsConnectionHandler,
62 MaintNotificationsPoolHandler,
63 OSSMaintNotificationsHandler,
64)
65from .observability.attributes import (
66 DB_CLIENT_CONNECTION_POOL_NAME,
67 DB_CLIENT_CONNECTION_STATE,
68 AttributeBuilder,
69 ConnectionState,
70 CSCReason,
71 CSCResult,
72 get_pool_name,
73)
74from .observability.metrics import CloseReason
75from .observability.recorder import (
76 init_csc_items,
77 record_connection_closed,
78 record_connection_count,
79 record_connection_create_time,
80 record_connection_wait_time,
81 record_csc_eviction,
82 record_csc_network_saved,
83 record_csc_request,
84 record_error_count,
85 register_csc_items_callback,
86)
87from .retry import Retry
88from .utils import (
89 CRYPTOGRAPHY_AVAILABLE,
90 DEFAULT_RESP_VERSION,
91 HIREDIS_AVAILABLE,
92 SENTINEL,
93 SSL_AVAILABLE,
94 check_protocol_version,
95 compare_versions,
96 deprecated_args,
97 ensure_string,
98 format_error_message,
99 str_if_bytes,
100)
102if SSL_AVAILABLE:
103 import ssl
104 from ssl import VerifyFlags
105else:
106 ssl = None
107 VerifyFlags = None
109if HIREDIS_AVAILABLE:
110 import hiredis
112SYM_STAR = b"*"
113SYM_DOLLAR = b"$"
114SYM_CRLF = b"\r\n"
115SYM_EMPTY = b""
117DefaultParser: Type[Union[_RESP2Parser, _RESP3Parser, _HiredisParser]]
118if HIREDIS_AVAILABLE:
119 DefaultParser = _HiredisParser
120else:
121 DefaultParser = _RESP2Parser
124class HiredisRespSerializer:
125 def pack(self, *args: List):
126 """Pack a series of arguments into the Redis protocol"""
127 output = []
129 if isinstance(args[0], str):
130 args = tuple(args[0].encode().split()) + args[1:]
131 elif b" " in args[0]:
132 args = tuple(args[0].split()) + args[1:]
133 args = tuple(
134 bytes(arg) if isinstance(arg, (bytearray, memoryview)) else arg
135 for arg in args
136 )
137 try:
138 output.append(hiredis.pack_command(args))
139 except TypeError:
140 _, value, traceback = sys.exc_info()
141 raise DataError(value).with_traceback(traceback)
143 return output
146class PythonRespSerializer:
147 def __init__(self, buffer_cutoff, encode) -> None:
148 self._buffer_cutoff = buffer_cutoff
149 self.encode = encode
151 def pack(self, *args):
152 """Pack a series of arguments into the Redis protocol"""
153 output = []
154 # the client might have included 1 or more literal arguments in
155 # the command name, e.g., 'CONFIG GET'. The Redis server expects these
156 # arguments to be sent separately, so split the first argument
157 # manually. These arguments should be bytestrings so that they are
158 # not encoded.
159 if isinstance(args[0], str):
160 args = tuple(args[0].encode().split()) + args[1:]
161 elif b" " in args[0]:
162 args = tuple(args[0].split()) + args[1:]
164 buff = SYM_EMPTY.join((SYM_STAR, str(len(args)).encode(), SYM_CRLF))
166 buffer_cutoff = self._buffer_cutoff
167 for arg in map(self.encode, args):
168 # to avoid large string mallocs, chunk the command into the
169 # output list if we're sending large values or memoryviews
170 arg_length = len(arg)
171 if (
172 len(buff) > buffer_cutoff
173 or arg_length > buffer_cutoff
174 or isinstance(arg, memoryview)
175 ):
176 buff = SYM_EMPTY.join(
177 (buff, SYM_DOLLAR, str(arg_length).encode(), SYM_CRLF)
178 )
179 output.append(buff)
180 output.append(arg)
181 buff = SYM_CRLF
182 else:
183 buff = SYM_EMPTY.join(
184 (
185 buff,
186 SYM_DOLLAR,
187 str(arg_length).encode(),
188 SYM_CRLF,
189 arg,
190 SYM_CRLF,
191 )
192 )
193 output.append(buff)
194 return output
197class ConnectionInterface:
198 @abstractmethod
199 def repr_pieces(self):
200 pass
202 @abstractmethod
203 def register_connect_callback(self, callback):
204 pass
206 @abstractmethod
207 def deregister_connect_callback(self, callback):
208 pass
210 @abstractmethod
211 def set_parser(self, parser_class):
212 pass
214 @abstractmethod
215 def get_protocol(self):
216 pass
218 @abstractmethod
219 def connect(self):
220 pass
222 @abstractmethod
223 def on_connect(self):
224 pass
226 @abstractmethod
227 def disconnect(self, *args, **kwargs):
228 pass
230 @abstractmethod
231 def check_health(self):
232 pass
234 @abstractmethod
235 def send_packed_command(self, command, check_health=True):
236 pass
238 @abstractmethod
239 def send_command(self, *args, **kwargs):
240 pass
242 @abstractmethod
243 def can_read(self, timeout: float = 0) -> bool:
244 # TODO: Rename this API; it detects pending data or dirty/closed
245 # connection state, not only whether application data can be read.
246 pass
248 @abstractmethod
249 def read_response(
250 self,
251 disable_decoding=False,
252 *,
253 timeout: Union[float, object] = SENTINEL,
254 disconnect_on_error=True,
255 push_request=False,
256 ):
257 pass
259 @abstractmethod
260 def pack_command(self, *args):
261 pass
263 @abstractmethod
264 def pack_commands(self, commands):
265 pass
267 @property
268 @abstractmethod
269 def handshake_metadata(self) -> Union[Dict[bytes, bytes], Dict[str, str]]:
270 pass
272 @abstractmethod
273 def set_re_auth_token(self, token: TokenInterface):
274 pass
276 @abstractmethod
277 def re_auth(self):
278 pass
280 @abstractmethod
281 def mark_for_reconnect(self):
282 """
283 Mark the connection to be reconnected on the next command.
284 This is useful when a connection is moved to a different node.
285 """
286 pass
288 @abstractmethod
289 def should_reconnect(self):
290 """
291 Returns True if the connection should be reconnected.
292 """
293 pass
295 @abstractmethod
296 def reset_should_reconnect(self):
297 """
298 Reset the internal flag to False.
299 """
300 pass
302 @abstractmethod
303 def extract_connection_details(self) -> str:
304 pass
306 @property
307 @abstractmethod
308 def is_connected(self) -> bool:
309 """
310 Return ``True`` if the connection to the server is active.
311 """
312 pass
315class MaintNotificationsAbstractConnection:
316 """
317 Abstract class for handling maintenance notifications logic.
318 This class is expected to be used as base class together with ConnectionInterface.
320 This class is intended to be used with multiple inheritance!
322 All logic related to maintenance notifications is encapsulated in this class.
323 """
325 def __init__(
326 self,
327 maint_notifications_config: Optional[MaintNotificationsConfig],
328 maint_notifications_pool_handler: Optional[
329 MaintNotificationsPoolHandler
330 ] = None,
331 maintenance_state: "MaintenanceState" = MaintenanceState.NONE,
332 maintenance_notification_hash: Optional[int] = None,
333 orig_host_address: Optional[str] = None,
334 orig_socket_timeout: Optional[float] = None,
335 orig_socket_connect_timeout: Optional[float] = None,
336 oss_cluster_maint_notifications_handler: Optional[
337 OSSMaintNotificationsHandler
338 ] = None,
339 parser: Optional[Union[_HiredisParser, _RESP3Parser]] = None,
340 event_dispatcher: Optional[EventDispatcher] = None,
341 ):
342 """
343 Initialize the maintenance notifications for the connection.
345 Args:
346 maint_notifications_config (MaintNotificationsConfig): The configuration for maintenance notifications.
347 maint_notifications_pool_handler (Optional[MaintNotificationsPoolHandler]): The pool handler for maintenance notifications.
348 maintenance_state (MaintenanceState): The current maintenance state of the connection.
349 maintenance_notification_hash (Optional[int]): The current maintenance notification hash of the connection.
350 orig_host_address (Optional[str]): The original host address of the connection.
351 orig_socket_timeout (Optional[float]): The original socket timeout of the connection.
352 orig_socket_connect_timeout (Optional[float]): The original socket connect timeout of the connection.
353 oss_cluster_maint_notifications_handler (Optional[OSSMaintNotificationsHandler]): The OSS cluster handler for maintenance notifications.
354 parser (Optional[Union[_HiredisParser, _RESP3Parser]]): The parser to use for maintenance notifications.
355 If not provided, the parser from the connection is used.
356 This is useful when the parser is created after this object.
357 """
358 self.maint_notifications_config = maint_notifications_config
359 self.maintenance_state = maintenance_state
360 self.maintenance_notification_hash = maintenance_notification_hash
362 if event_dispatcher is not None:
363 self.event_dispatcher = event_dispatcher
364 else:
365 self.event_dispatcher = EventDispatcher()
367 self._configure_maintenance_notifications(
368 maint_notifications_pool_handler,
369 orig_host_address,
370 orig_socket_timeout,
371 orig_socket_connect_timeout,
372 oss_cluster_maint_notifications_handler,
373 parser,
374 )
375 self._processed_start_maint_notifications = set()
376 self._skipped_end_maint_notifications = set()
378 @abstractmethod
379 def _get_parser(self) -> Union[_HiredisParser, _RESP3Parser]:
380 pass
382 @abstractmethod
383 def _get_socket(self) -> Optional[socket.socket]:
384 pass
386 @abstractmethod
387 def get_protocol(self) -> Union[int, str]:
388 """
389 Returns:
390 The RESP protocol version, or ``None`` if the protocol is not specified,
391 in which case the server default will be used.
392 """
393 pass
395 @property
396 @abstractmethod
397 def host(self) -> str:
398 pass
400 @host.setter
401 @abstractmethod
402 def host(self, value: str):
403 pass
405 @property
406 @abstractmethod
407 def socket_timeout(self) -> Optional[Union[float, int]]:
408 pass
410 @socket_timeout.setter
411 @abstractmethod
412 def socket_timeout(self, value: Optional[Union[float, int]]):
413 pass
415 @property
416 @abstractmethod
417 def socket_connect_timeout(self) -> Optional[Union[float, int]]:
418 pass
420 @socket_connect_timeout.setter
421 @abstractmethod
422 def socket_connect_timeout(self, value: Optional[Union[float, int]]):
423 pass
425 @abstractmethod
426 def send_command(self, *args, **kwargs):
427 pass
429 @abstractmethod
430 def read_response(
431 self,
432 disable_decoding=False,
433 *,
434 timeout: Union[float, object] = SENTINEL,
435 disconnect_on_error=True,
436 push_request=False,
437 ):
438 pass
440 @abstractmethod
441 def disconnect(self, *args, **kwargs):
442 pass
444 def _configure_maintenance_notifications(
445 self,
446 maint_notifications_pool_handler: Optional[
447 MaintNotificationsPoolHandler
448 ] = None,
449 orig_host_address=None,
450 orig_socket_timeout=None,
451 orig_socket_connect_timeout=None,
452 oss_cluster_maint_notifications_handler: Optional[
453 OSSMaintNotificationsHandler
454 ] = None,
455 parser: Optional[Union[_HiredisParser, _RESP3Parser]] = None,
456 ):
457 """
458 Enable maintenance notifications by setting up
459 handlers and storing original connection parameters.
461 Should be used ONLY with parsers that support push notifications.
462 """
463 if (
464 not self.maint_notifications_config
465 or not self.maint_notifications_config.enabled
466 ):
467 self._maint_notifications_pool_handler = None
468 self._maint_notifications_connection_handler = None
469 self._oss_cluster_maint_notifications_handler = None
470 return
472 if not parser:
473 raise RedisError(
474 "To configure maintenance notifications, a parser must be provided!"
475 )
477 if not isinstance(parser, _HiredisParser) and not isinstance(
478 parser, _RESP3Parser
479 ):
480 raise RedisError(
481 "Maintenance notifications are only supported with hiredis and RESP3 parsers!"
482 )
484 if maint_notifications_pool_handler:
485 # Extract a reference to a new pool handler that copies all properties
486 # of the original one and has a different connection reference
487 # This is needed because when we attach the handler to the parser
488 # we need to make sure that the handler has a reference to the
489 # connection that the parser is attached to.
490 self._maint_notifications_pool_handler = (
491 maint_notifications_pool_handler.get_handler_for_connection()
492 )
493 self._maint_notifications_pool_handler.set_connection(self)
494 else:
495 self._maint_notifications_pool_handler = None
497 self._maint_notifications_connection_handler = (
498 MaintNotificationsConnectionHandler(self, self.maint_notifications_config)
499 )
501 if oss_cluster_maint_notifications_handler:
502 self._oss_cluster_maint_notifications_handler = (
503 oss_cluster_maint_notifications_handler
504 )
505 else:
506 self._oss_cluster_maint_notifications_handler = None
508 # Set up OSS cluster handler to parser if available
509 if self._oss_cluster_maint_notifications_handler:
510 parser.set_oss_cluster_maint_push_handler(
511 self._oss_cluster_maint_notifications_handler.handle_notification
512 )
514 # Set up pool handler to parser if available
515 if self._maint_notifications_pool_handler:
516 parser.set_node_moving_push_handler(
517 self._maint_notifications_pool_handler.handle_notification
518 )
520 # Set up connection handler
521 parser.set_maintenance_push_handler(
522 self._maint_notifications_connection_handler.handle_notification
523 )
525 # Store original connection parameters
526 self.orig_host_address = orig_host_address if orig_host_address else self.host
527 self.orig_socket_timeout = (
528 orig_socket_timeout if orig_socket_timeout else self.socket_timeout
529 )
530 self.orig_socket_connect_timeout = (
531 orig_socket_connect_timeout
532 if orig_socket_connect_timeout
533 else self.socket_connect_timeout
534 )
536 def set_maint_notifications_pool_handler_for_connection(
537 self, maint_notifications_pool_handler: MaintNotificationsPoolHandler
538 ):
539 # Deep copy the pool handler to avoid sharing the same pool handler
540 # between multiple connections, because otherwise each connection will override
541 # the connection reference and the pool handler will only hold a reference
542 # to the last connection that was set.
543 maint_notifications_pool_handler_copy = (
544 maint_notifications_pool_handler.get_handler_for_connection()
545 )
547 maint_notifications_pool_handler_copy.set_connection(self)
548 self._get_parser().set_node_moving_push_handler(
549 maint_notifications_pool_handler_copy.handle_notification
550 )
552 self._maint_notifications_pool_handler = maint_notifications_pool_handler_copy
554 # Update maintenance notification connection handler if it doesn't exist
555 if not self._maint_notifications_connection_handler:
556 self._maint_notifications_connection_handler = (
557 MaintNotificationsConnectionHandler(
558 self, maint_notifications_pool_handler.config
559 )
560 )
561 self._get_parser().set_maintenance_push_handler(
562 self._maint_notifications_connection_handler.handle_notification
563 )
564 else:
565 self._maint_notifications_connection_handler.config = (
566 maint_notifications_pool_handler.config
567 )
569 def set_maint_notifications_cluster_handler_for_connection(
570 self, oss_cluster_maint_notifications_handler: OSSMaintNotificationsHandler
571 ):
572 self._get_parser().set_oss_cluster_maint_push_handler(
573 oss_cluster_maint_notifications_handler.handle_notification
574 )
576 self._oss_cluster_maint_notifications_handler = (
577 oss_cluster_maint_notifications_handler
578 )
580 # Update maintenance notification connection handler if it doesn't exist
581 if not self._maint_notifications_connection_handler:
582 self._maint_notifications_connection_handler = (
583 MaintNotificationsConnectionHandler(
584 self, oss_cluster_maint_notifications_handler.config
585 )
586 )
587 self._get_parser().set_maintenance_push_handler(
588 self._maint_notifications_connection_handler.handle_notification
589 )
590 else:
591 self._maint_notifications_connection_handler.config = (
592 oss_cluster_maint_notifications_handler.config
593 )
595 def activate_maint_notifications_handling_if_enabled(self, check_health=True):
596 # Send maintenance notifications handshake if RESP3 is active
597 # and maintenance notifications are enabled
598 # and we have a host to determine the endpoint type from
599 # When the maint_notifications_config enabled mode is "auto",
600 # we just log a warning if the handshake fails
601 # When the mode is enabled=True, we raise an exception in case of failure
602 host = getattr(self, "host", None)
603 if (
604 self.get_protocol() not in [2, "2"]
605 and self.maint_notifications_config
606 and self.maint_notifications_config.enabled
607 and self._maint_notifications_connection_handler
608 and host is not None
609 ):
610 self._enable_maintenance_notifications(
611 maint_notifications_config=self.maint_notifications_config,
612 check_health=check_health,
613 )
615 def _enable_maintenance_notifications(
616 self, maint_notifications_config: MaintNotificationsConfig, check_health=True
617 ):
618 try:
619 host = getattr(self, "host", None)
620 if host is None:
621 raise ValueError(
622 "Cannot enable maintenance notifications for connection"
623 " object that doesn't have a host attribute."
624 )
625 else:
626 endpoint_type = maint_notifications_config.get_endpoint_type(host, self)
627 self.send_command(
628 "CLIENT",
629 "MAINT_NOTIFICATIONS",
630 "ON",
631 "moving-endpoint-type",
632 endpoint_type.value,
633 check_health=check_health,
634 )
635 response = self.read_response()
636 if not response or str_if_bytes(response) != "OK":
637 raise ResponseError(
638 "The server doesn't support maintenance notifications"
639 )
640 except Exception as e:
641 if (
642 isinstance(e, ResponseError)
643 and maint_notifications_config.enabled == "auto"
644 ):
645 # Log warning but don't fail the connection
646 import logging
648 logger = logging.getLogger(__name__)
649 logger.debug(f"Failed to enable maintenance notifications: {e}")
650 else:
651 raise
653 def get_resolved_ip(self) -> Optional[str]:
654 """
655 Extract the resolved IP address from an
656 established connection or resolve it from the host.
658 First tries to get the actual IP from the socket (most accurate),
659 then falls back to DNS resolution if needed.
661 Returns:
662 str: The resolved IP address, or None if it cannot be determined
663 """
665 # Method 1: Try to get the actual IP from the established socket connection
666 # This is most accurate as it shows the exact IP being used
667 try:
668 conn_socket = self._get_socket()
669 if conn_socket is not None:
670 peer_addr = conn_socket.getpeername()
671 if peer_addr and len(peer_addr) >= 1:
672 # For TCP sockets, peer_addr is typically (host, port) tuple
673 # Return just the host part
674 return peer_addr[0]
675 except (AttributeError, OSError):
676 # Socket might not be connected or getpeername() might fail
677 pass
679 # Method 2: Fallback to DNS resolution of the host
680 # This is less accurate but works when socket is not available
681 try:
682 host = getattr(self, "host", "localhost")
683 port = getattr(self, "port", 6379)
684 if host:
685 # Use getaddrinfo to resolve the hostname to IP
686 # This mimics what the connection would do during _connect()
687 addr_info = socket.getaddrinfo(
688 host, port, socket.AF_UNSPEC, socket.SOCK_STREAM
689 )
690 if addr_info:
691 # Return the IP from the first result
692 # addr_info[0] is (family, socktype, proto, canonname, sockaddr)
693 # sockaddr[0] is the IP address
694 return str(addr_info[0][4][0])
695 except (AttributeError, OSError, socket.gaierror):
696 # DNS resolution might fail
697 pass
699 return None
701 @property
702 def maintenance_state(self) -> MaintenanceState:
703 return self._maintenance_state
705 @maintenance_state.setter
706 def maintenance_state(self, state: "MaintenanceState"):
707 self._maintenance_state = state
709 def add_maint_start_notification(self, id: int):
710 self._processed_start_maint_notifications.add(id)
712 def get_processed_start_notifications(self) -> set:
713 return self._processed_start_maint_notifications
715 def add_skipped_end_notification(self, id: int):
716 self._skipped_end_maint_notifications.add(id)
718 def get_skipped_end_notifications(self) -> set:
719 return self._skipped_end_maint_notifications
721 def reset_received_notifications(self):
722 self._processed_start_maint_notifications.clear()
723 self._skipped_end_maint_notifications.clear()
725 def getpeername(self):
726 """
727 Returns the peer name of the connection.
728 """
729 conn_socket = self._get_socket()
730 if conn_socket:
731 return conn_socket.getpeername()[0]
732 return None
734 def update_current_socket_timeout(self, relaxed_timeout: Optional[float] = None):
735 conn_socket = self._get_socket()
736 if conn_socket:
737 timeout = relaxed_timeout if relaxed_timeout != -1 else self.socket_timeout
738 # if the current timeout is 0 it means we are in the middle of a can_read call
739 # in this case we don't want to change the timeout because the operation
740 # is non-blocking and should return immediately
741 # Changing the state from non-blocking to blocking in the middle of a read operation
742 # will lead to a deadlock
743 if conn_socket.gettimeout() != 0:
744 conn_socket.settimeout(timeout)
745 self.update_parser_timeout(timeout)
747 def update_parser_timeout(self, timeout: Optional[float] = None):
748 parser = self._get_parser()
749 if parser and parser._buffer:
750 if isinstance(parser, _RESP3Parser) and timeout:
751 parser._buffer.socket_timeout = timeout
752 elif isinstance(parser, _HiredisParser):
753 parser._socket_timeout = timeout
755 def set_tmp_settings(
756 self,
757 tmp_host_address: Optional[Union[str, object]] = SENTINEL,
758 tmp_relaxed_timeout: Optional[float] = None,
759 ):
760 """
761 The value of SENTINEL is used to indicate that the property should not be updated.
762 """
763 if tmp_host_address and tmp_host_address != SENTINEL:
764 self.host = str(tmp_host_address)
765 if tmp_relaxed_timeout != -1:
766 self.socket_timeout = tmp_relaxed_timeout
767 self.socket_connect_timeout = tmp_relaxed_timeout
769 def reset_tmp_settings(
770 self,
771 reset_host_address: bool = False,
772 reset_relaxed_timeout: bool = False,
773 ):
774 if reset_host_address:
775 self.host = self.orig_host_address
776 if reset_relaxed_timeout:
777 self.socket_timeout = self.orig_socket_timeout
778 self.socket_connect_timeout = self.orig_socket_connect_timeout
781class AbstractConnection(MaintNotificationsAbstractConnection, ConnectionInterface):
782 "Manages communication to and from a Redis server"
784 @deprecated_args(
785 args_to_warn=["lib_name", "lib_version"],
786 reason="Use 'driver_info' parameter instead. "
787 "lib_name and lib_version will be removed in a future version.",
788 )
789 def __init__(
790 self,
791 db: int = 0,
792 password: Optional[str] = None,
793 socket_timeout: Optional[float] = DEFAULT_SOCKET_TIMEOUT,
794 socket_connect_timeout: Optional[float] = DEFAULT_SOCKET_CONNECT_TIMEOUT,
795 retry_on_timeout: bool = False,
796 retry_on_error: Union[Iterable[Type[Exception]], object] = SENTINEL,
797 encoding: str = "utf-8",
798 encoding_errors: str = "strict",
799 decode_responses: bool = False,
800 parser_class=DefaultParser,
801 socket_read_size: int = DEFAULT_SOCKET_READ_SIZE,
802 health_check_interval: int = 0,
803 client_name: Optional[str] = None,
804 lib_name: Union[Optional[str], object] = SENTINEL,
805 lib_version: Union[Optional[str], object] = SENTINEL,
806 driver_info: Union[Optional[DriverInfo], object] = SENTINEL,
807 username: Optional[str] = None,
808 retry: Union[Any, None] = None,
809 redis_connect_func: Optional[Callable[[], None]] = None,
810 credential_provider: Optional[CredentialProvider] = None,
811 protocol: Optional[int] = None,
812 legacy_responses: bool = True,
813 command_packer: Optional[Callable[[], None]] = None,
814 event_dispatcher: Optional[EventDispatcher] = None,
815 maint_notifications_config: Optional[MaintNotificationsConfig] = None,
816 maint_notifications_pool_handler: Optional[
817 MaintNotificationsPoolHandler
818 ] = None,
819 maintenance_state: "MaintenanceState" = MaintenanceState.NONE,
820 maintenance_notification_hash: Optional[int] = None,
821 orig_host_address: Optional[str] = None,
822 orig_socket_timeout: Optional[float] = None,
823 orig_socket_connect_timeout: Optional[float] = None,
824 oss_cluster_maint_notifications_handler: Optional[
825 OSSMaintNotificationsHandler
826 ] = None,
827 ):
828 """
829 Initialize a new Connection.
831 To specify a retry policy for specific errors, first set
832 `retry_on_error` to a list of the error/s to retry on, then set
833 `retry` to a valid `Retry` object.
834 To retry on TimeoutError, `retry_on_timeout` can also be set to `True`.
836 Parameters
837 ----------
838 driver_info : DriverInfo, optional
839 Driver metadata for CLIENT SETINFO. If provided, lib_name and lib_version
840 are ignored. If not provided, a DriverInfo will be created from lib_name
841 and lib_version. Explicit None disables CLIENT SETINFO.
842 lib_name : str, optional
843 **Deprecated.** Use driver_info instead. Library name for CLIENT SETINFO.
844 lib_version : str, optional
845 **Deprecated.** Use driver_info instead. Library version for CLIENT SETINFO.
846 """
847 if (username or password) and credential_provider is not None:
848 raise DataError(
849 "'username' and 'password' cannot be passed along with 'credential_"
850 "provider'. Please provide only one of the following arguments: \n"
851 "1. 'password' and (optional) 'username'\n"
852 "2. 'credential_provider'"
853 )
854 if event_dispatcher is None:
855 self._event_dispatcher = EventDispatcher()
856 else:
857 self._event_dispatcher = event_dispatcher
858 self.pid = os.getpid()
859 self.db = db
860 self.client_name = client_name
862 # Handle driver_info: if provided, use it; otherwise create from lib_name/lib_version.
863 self.driver_info = resolve_driver_info(driver_info, lib_name, lib_version)
865 self.credential_provider = credential_provider
866 self.password = password
867 self.username = username
868 self._socket_timeout = socket_timeout
869 if socket_connect_timeout is None:
870 socket_connect_timeout = socket_timeout
871 self._socket_connect_timeout = socket_connect_timeout
872 self.retry_on_timeout = retry_on_timeout
873 if retry_on_error is SENTINEL:
874 retry_on_errors_list = []
875 else:
876 retry_on_errors_list = list(retry_on_error)
877 if retry_on_timeout:
878 # Add TimeoutError to the errors list to retry on
879 retry_on_errors_list.append(TimeoutError)
880 self.retry_on_error = retry_on_errors_list
881 if retry or self.retry_on_error:
882 if retry is None:
883 self.retry = Retry(NoBackoff(), 1)
884 else:
885 # deep-copy the Retry object as it is mutable
886 self.retry = copy.deepcopy(retry)
887 if self.retry_on_error:
888 # Update the retry's supported errors with the specified errors
889 self.retry.update_supported_errors(self.retry_on_error)
890 else:
891 self.retry = Retry(NoBackoff(), 0)
892 self.health_check_interval = health_check_interval
893 self.next_health_check = 0
894 self.redis_connect_func = redis_connect_func
895 self.encoder = Encoder(encoding, encoding_errors, decode_responses)
896 self.handshake_metadata = None
897 self._sock = None
898 self._socket_read_size = socket_read_size
899 self._connect_callbacks = []
900 self._buffer_cutoff = 6000
901 self._re_auth_token: Optional[TokenInterface] = None
902 try:
903 p = int(protocol)
904 except TypeError:
905 p = DEFAULT_RESP_VERSION
906 except ValueError:
907 raise ConnectionError("protocol must be an integer")
908 else:
909 if p < 2 or p > 3:
910 raise ConnectionError("protocol must be either 2 or 3")
911 self.protocol = p
912 self.legacy_responses = legacy_responses
913 if self.protocol == 3 and parser_class == _RESP2Parser:
914 # If the protocol is 3 but the parser is RESP2, change it to RESP3
915 # This is needed because the parser might be set before the protocol
916 # or might be provided as a kwarg to the constructor
917 # We need to react on discrepancy only for RESP2 and RESP3
918 # as hiredis supports both
919 parser_class = _RESP3Parser
920 self.set_parser(parser_class)
922 self._command_packer = self._construct_command_packer(command_packer)
923 self._should_reconnect = False
925 # Set up maintenance notifications
926 MaintNotificationsAbstractConnection.__init__(
927 self,
928 maint_notifications_config,
929 maint_notifications_pool_handler,
930 maintenance_state,
931 maintenance_notification_hash,
932 orig_host_address,
933 orig_socket_timeout,
934 orig_socket_connect_timeout,
935 oss_cluster_maint_notifications_handler,
936 self._parser,
937 event_dispatcher=self._event_dispatcher,
938 )
940 def __repr__(self):
941 repr_args = ",".join([f"{k}={v}" for k, v in self.repr_pieces()])
942 return f"<{self.__class__.__module__}.{self.__class__.__name__}({repr_args})>"
944 @abstractmethod
945 def repr_pieces(self):
946 pass
948 def __del__(self):
949 try:
950 self.disconnect()
951 except Exception:
952 pass
954 @property
955 def is_connected(self) -> bool:
956 return self._sock is not None
958 def _construct_command_packer(self, packer):
959 if packer is not None:
960 return packer
961 elif HIREDIS_AVAILABLE:
962 return HiredisRespSerializer()
963 else:
964 return PythonRespSerializer(self._buffer_cutoff, self.encoder.encode)
966 def register_connect_callback(self, callback):
967 """
968 Register a callback to be called when the connection is established either
969 initially or reconnected. This allows listeners to issue commands that
970 are ephemeral to the connection, for example pub/sub subscription or
971 key tracking. The callback must be a _method_ and will be kept as
972 a weak reference.
973 """
974 wm = weakref.WeakMethod(callback)
975 if wm not in self._connect_callbacks:
976 self._connect_callbacks.append(wm)
978 def deregister_connect_callback(self, callback):
979 """
980 De-register a previously registered callback. It will no-longer receive
981 notifications on connection events. Calling this is not required when the
982 listener goes away, since the callbacks are kept as weak methods.
983 """
984 try:
985 self._connect_callbacks.remove(weakref.WeakMethod(callback))
986 except ValueError:
987 pass
989 def set_parser(self, parser_class):
990 """
991 Creates a new instance of parser_class with socket size:
992 _socket_read_size and assigns it to the parser for the connection
993 :param parser_class: The required parser class
994 """
995 self._parser = parser_class(socket_read_size=self._socket_read_size)
997 def _get_parser(self) -> Union[_HiredisParser, _RESP3Parser, _RESP2Parser]:
998 return self._parser
1000 def connect(self):
1001 "Connects to the Redis server if not already connected"
1002 # try once the socket connect with the handshake, retry the whole
1003 # connect/handshake flow based on retry policy
1004 self.retry.call_with_retry(
1005 lambda: self.connect_check_health(
1006 check_health=True, retry_socket_connect=False
1007 ),
1008 lambda error: self.disconnect(error),
1009 )
1011 def connect_check_health(
1012 self, check_health: bool = True, retry_socket_connect: bool = True
1013 ):
1014 if self._sock:
1015 return
1016 # Track actual retry attempts for error reporting
1017 actual_retry_attempts = [0]
1019 def failure_callback(error, failure_count):
1020 actual_retry_attempts[0] = failure_count
1021 self.disconnect(error=error, failure_count=failure_count)
1023 try:
1024 if retry_socket_connect:
1025 sock = self.retry.call_with_retry(
1026 self._connect,
1027 failure_callback,
1028 with_failure_count=True,
1029 )
1030 else:
1031 sock = self._connect()
1032 except socket.timeout:
1033 e = TimeoutError("Timeout connecting to server")
1034 record_error_count(
1035 server_address=self.host,
1036 server_port=self.port,
1037 network_peer_address=self.host,
1038 network_peer_port=self.port,
1039 error_type=e,
1040 retry_attempts=actual_retry_attempts[0],
1041 )
1042 raise e
1043 except OSError as e:
1044 e = ConnectionError(self._error_message(e))
1045 record_error_count(
1046 server_address=getattr(self, "host", None),
1047 server_port=getattr(self, "port", None),
1048 network_peer_address=getattr(self, "host", None),
1049 network_peer_port=getattr(self, "port", None),
1050 error_type=e,
1051 retry_attempts=actual_retry_attempts[0],
1052 )
1053 raise e
1055 self._sock = sock
1056 try:
1057 if self.redis_connect_func is None:
1058 # Use the default on_connect function
1059 self.on_connect_check_health(check_health=check_health)
1060 else:
1061 # Use the passed function redis_connect_func
1062 self.redis_connect_func(self)
1063 except RedisError:
1064 # clean up after any error in on_connect
1065 self.disconnect()
1066 raise
1068 # run any user callbacks. right now the only internal callback
1069 # is for pubsub channel/pattern resubscription
1070 # first, remove any dead weakrefs
1071 self._connect_callbacks = [ref for ref in self._connect_callbacks if ref()]
1072 for ref in self._connect_callbacks:
1073 callback = ref()
1074 if callback:
1075 callback(self)
1077 @abstractmethod
1078 def _connect(self):
1079 pass
1081 @abstractmethod
1082 def _host_error(self):
1083 pass
1085 def _error_message(self, exception):
1086 return format_error_message(self._host_error(), exception)
1088 def on_connect(self):
1089 self.on_connect_check_health(check_health=True)
1091 def on_connect_check_health(self, check_health: bool = True):
1092 "Initialize the connection, authenticate and select a database"
1093 self._parser.on_connect(self)
1094 parser = self._parser
1096 auth_args = None
1097 # if credential provider or username and/or password are set, authenticate
1098 if self.credential_provider or (self.username or self.password):
1099 cred_provider = (
1100 self.credential_provider
1101 or UsernamePasswordCredentialProvider(self.username, self.password)
1102 )
1103 auth_args = cred_provider.get_credentials()
1105 # if resp version is specified and we have auth args,
1106 # we need to send them via HELLO
1107 if auth_args and self.protocol not in [2, "2"]:
1108 if isinstance(self._parser, _RESP2Parser):
1109 self.set_parser(_RESP3Parser)
1110 # update cluster exception classes
1111 self._parser.EXCEPTION_CLASSES = parser.EXCEPTION_CLASSES
1112 self._parser.on_connect(self)
1113 if len(auth_args) == 1:
1114 auth_args = ["default", auth_args[0]]
1115 # avoid checking health here -- PING will fail if we try
1116 # to check the health prior to the AUTH
1117 self.send_command(
1118 "HELLO", self.protocol, "AUTH", *auth_args, check_health=False
1119 )
1120 self.handshake_metadata = self.read_response()
1121 # if response.get(b"proto") != self.protocol and response.get(
1122 # "proto"
1123 # ) != self.protocol:
1124 # raise ConnectionError("Invalid RESP version")
1125 elif auth_args:
1126 # avoid checking health here -- PING will fail if we try
1127 # to check the health prior to the AUTH
1128 self.send_command("AUTH", *auth_args, check_health=False)
1130 try:
1131 auth_response = self.read_response()
1132 except AuthenticationWrongNumberOfArgsError:
1133 # a username and password were specified but the Redis
1134 # server seems to be < 6.0.0 which expects a single password
1135 # arg. retry auth with just the password.
1136 # https://github.com/andymccurdy/redis-py/issues/1274
1137 self.send_command("AUTH", auth_args[-1], check_health=False)
1138 auth_response = self.read_response()
1140 if str_if_bytes(auth_response) != "OK":
1141 raise AuthenticationError("Invalid Username or Password")
1143 # if resp version is specified, switch to it
1144 elif self.protocol not in [2, "2"]:
1145 if isinstance(self._parser, _RESP2Parser):
1146 self.set_parser(_RESP3Parser)
1147 # update cluster exception classes
1148 self._parser.EXCEPTION_CLASSES = parser.EXCEPTION_CLASSES
1149 self._parser.on_connect(self)
1150 self.send_command("HELLO", self.protocol, check_health=check_health)
1151 self.handshake_metadata = self.read_response()
1152 if (
1153 self.handshake_metadata.get(b"proto") != self.protocol
1154 and self.handshake_metadata.get("proto") != self.protocol
1155 ):
1156 raise ConnectionError("Invalid RESP version")
1158 # Activate maintenance notifications for this connection
1159 # if enabled in the configuration
1160 # This is a no-op if maintenance notifications are not enabled
1161 self.activate_maint_notifications_handling_if_enabled(check_health=check_health)
1163 # if a client_name is given, set it
1164 if self.client_name:
1165 self.send_command(
1166 "CLIENT",
1167 "SETNAME",
1168 self.client_name,
1169 check_health=check_health,
1170 )
1171 if str_if_bytes(self.read_response()) != "OK":
1172 raise ConnectionError("Error setting client name")
1174 # Set the library name and version from driver_info
1175 try:
1176 if self.driver_info and self.driver_info.formatted_name:
1177 self.send_command(
1178 "CLIENT",
1179 "SETINFO",
1180 "LIB-NAME",
1181 self.driver_info.formatted_name,
1182 check_health=check_health,
1183 )
1184 self.read_response()
1185 except ResponseError:
1186 pass
1188 try:
1189 if self.driver_info and self.driver_info.lib_version:
1190 self.send_command(
1191 "CLIENT",
1192 "SETINFO",
1193 "LIB-VER",
1194 self.driver_info.lib_version,
1195 check_health=check_health,
1196 )
1197 self.read_response()
1198 except ResponseError:
1199 pass
1201 # if a database is specified, switch to it
1202 if self.db:
1203 self.send_command("SELECT", self.db, check_health=check_health)
1204 if str_if_bytes(self.read_response()) != "OK":
1205 raise ConnectionError("Invalid Database")
1207 def disconnect(self, *args, **kwargs):
1208 "Disconnects from the Redis server"
1209 self._parser.on_disconnect()
1211 conn_sock = self._sock
1212 self._sock = None
1213 # reset the reconnect flag
1214 self.reset_should_reconnect()
1216 if conn_sock is None:
1217 return
1219 if os.getpid() == self.pid:
1220 try:
1221 conn_sock.shutdown(socket.SHUT_RDWR)
1222 except (OSError, TypeError):
1223 pass
1225 try:
1226 conn_sock.close()
1227 except OSError:
1228 pass
1230 error = kwargs.get("error")
1231 failure_count = kwargs.get("failure_count")
1232 health_check_failed = kwargs.get("health_check_failed")
1234 if error:
1235 if health_check_failed:
1236 close_reason = CloseReason.HEALTHCHECK_FAILED
1237 else:
1238 close_reason = CloseReason.ERROR
1240 if failure_count is not None and failure_count > self.retry.get_retries():
1241 record_error_count(
1242 server_address=self.host,
1243 server_port=self.port,
1244 network_peer_address=self.host,
1245 network_peer_port=self.port,
1246 error_type=error,
1247 retry_attempts=failure_count,
1248 )
1250 record_connection_closed(
1251 close_reason=close_reason,
1252 error_type=error,
1253 )
1254 else:
1255 record_connection_closed(
1256 close_reason=CloseReason.APPLICATION_CLOSE,
1257 )
1259 if self.maintenance_state == MaintenanceState.MAINTENANCE:
1260 # this block will be executed only if the connection was in maintenance state
1261 # and the connection was closed.
1262 # The state change won't be applied on connections that are in Moving state
1263 # because their state and configurations will be handled when the moving ttl expires.
1264 self.reset_tmp_settings(reset_relaxed_timeout=True)
1265 self.maintenance_state = MaintenanceState.NONE
1266 # reset the sets that keep track of received start maint
1267 # notifications and skipped end maint notifications
1268 self.reset_received_notifications()
1270 def mark_for_reconnect(self):
1271 self._should_reconnect = True
1273 def should_reconnect(self):
1274 return self._should_reconnect
1276 def reset_should_reconnect(self):
1277 self._should_reconnect = False
1279 def _send_ping(self):
1280 """Send PING, expect PONG in return"""
1281 self.send_command("PING", check_health=False)
1282 if str_if_bytes(self.read_response()) != "PONG":
1283 raise ConnectionError("Bad response from PING health check")
1285 def _ping_failed(self, error, failure_count):
1286 """Function to call when PING fails"""
1287 self.disconnect(
1288 error=error, failure_count=failure_count, health_check_failed=True
1289 )
1291 def check_health(self):
1292 """Check the health of the connection with a PING/PONG"""
1293 if self.health_check_interval and time.monotonic() > self.next_health_check:
1294 self.retry.call_with_retry(
1295 self._send_ping,
1296 self._ping_failed,
1297 with_failure_count=True,
1298 )
1300 def send_packed_command(self, command, check_health=True):
1301 """Send an already packed command to the Redis server"""
1302 if not self._sock:
1303 self.connect_check_health(check_health=False)
1304 # guard against health check recursion
1305 if check_health:
1306 self.check_health()
1307 try:
1308 if isinstance(command, str):
1309 command = [command]
1310 for item in command:
1311 self._sock.sendall(item)
1312 except socket.timeout:
1313 self.disconnect()
1314 raise TimeoutError("Timeout writing to socket")
1315 except OSError as e:
1316 self.disconnect()
1317 if len(e.args) == 1:
1318 errno, errmsg = "UNKNOWN", e.args[0]
1319 else:
1320 errno = e.args[0]
1321 errmsg = e.args[1]
1322 raise ConnectionError(f"Error {errno} while writing to socket. {errmsg}.")
1323 except BaseException:
1324 # BaseExceptions can be raised when a socket send operation is not
1325 # finished, e.g. due to a timeout. Ideally, a caller could then re-try
1326 # to send un-sent data. However, the send_packed_command() API
1327 # does not support it so there is no point in keeping the connection open.
1328 self.disconnect()
1329 raise
1331 def send_command(self, *args, **kwargs):
1332 """Pack and send a command to the Redis server"""
1333 self.send_packed_command(
1334 self._command_packer.pack(*args),
1335 check_health=kwargs.get("check_health", True),
1336 )
1338 def can_read(self, timeout: float = 0) -> bool:
1339 """Poll the socket to see if there's data that can be read."""
1340 # TODO: Rename this API; it detects pending data or dirty/closed
1341 # connection state, not only whether application data can be read.
1342 sock = self._sock
1343 if not sock:
1344 self.connect()
1346 host_error = self._host_error()
1348 try:
1349 return self._parser.can_read(timeout)
1351 except OSError as e:
1352 self.disconnect()
1353 raise ConnectionError(f"Error while reading from {host_error}: {e.args}")
1355 def read_response(
1356 self,
1357 disable_decoding=False,
1358 *,
1359 timeout: Union[float, object] = SENTINEL,
1360 disconnect_on_error=True,
1361 push_request=False,
1362 ):
1363 """Read the response from a previously sent command"""
1365 host_error = self._host_error()
1367 try:
1368 if self.protocol in ["3", 3]:
1369 response = self._parser.read_response(
1370 disable_decoding=disable_decoding,
1371 push_request=push_request,
1372 timeout=timeout,
1373 )
1374 else:
1375 response = self._parser.read_response(
1376 disable_decoding=disable_decoding, timeout=timeout
1377 )
1378 except socket.timeout:
1379 if disconnect_on_error:
1380 self.disconnect()
1381 raise TimeoutError(f"Timeout reading from {host_error}")
1382 except OSError as e:
1383 if disconnect_on_error:
1384 self.disconnect()
1385 raise ConnectionError(f"Error while reading from {host_error} : {e.args}")
1386 except BaseException:
1387 # Also by default close in case of BaseException. A lot of code
1388 # relies on this behaviour when doing Command/Response pairs.
1389 # See #1128.
1390 if disconnect_on_error:
1391 self.disconnect()
1392 raise
1394 if self.health_check_interval:
1395 self.next_health_check = time.monotonic() + self.health_check_interval
1397 if isinstance(response, ResponseError):
1398 try:
1399 raise response
1400 finally:
1401 del response # avoid creating ref cycles
1402 return response
1404 def pack_command(self, *args):
1405 """Pack a series of arguments into the Redis protocol"""
1406 return self._command_packer.pack(*args)
1408 def pack_commands(self, commands):
1409 """Pack multiple commands into the Redis protocol"""
1410 output = []
1411 pieces = []
1412 buffer_length = 0
1413 buffer_cutoff = self._buffer_cutoff
1415 for cmd in commands:
1416 for chunk in self._command_packer.pack(*cmd):
1417 chunklen = len(chunk)
1418 if (
1419 buffer_length > buffer_cutoff
1420 or chunklen > buffer_cutoff
1421 or isinstance(chunk, memoryview)
1422 ):
1423 if pieces:
1424 output.append(SYM_EMPTY.join(pieces))
1425 buffer_length = 0
1426 pieces = []
1428 if chunklen > buffer_cutoff or isinstance(chunk, memoryview):
1429 output.append(chunk)
1430 else:
1431 pieces.append(chunk)
1432 buffer_length += chunklen
1434 if pieces:
1435 output.append(SYM_EMPTY.join(pieces))
1436 return output
1438 def get_protocol(self) -> Union[int, str]:
1439 return self.protocol
1441 @property
1442 def handshake_metadata(self) -> Union[Dict[bytes, bytes], Dict[str, str]]:
1443 return self._handshake_metadata
1445 @handshake_metadata.setter
1446 def handshake_metadata(self, value: Union[Dict[bytes, bytes], Dict[str, str]]):
1447 self._handshake_metadata = value
1449 def set_re_auth_token(self, token: TokenInterface):
1450 self._re_auth_token = token
1452 def re_auth(self):
1453 if self._re_auth_token is not None:
1454 self.send_command(
1455 "AUTH",
1456 self._re_auth_token.try_get("oid"),
1457 self._re_auth_token.get_value(),
1458 )
1459 self.read_response()
1460 self._re_auth_token = None
1462 def _get_socket(self) -> Optional[socket.socket]:
1463 return self._sock
1465 @property
1466 def socket_timeout(self) -> Optional[Union[float, int]]:
1467 return self._socket_timeout
1469 @socket_timeout.setter
1470 def socket_timeout(self, value: Optional[Union[float, int]]):
1471 self._socket_timeout = value
1473 @property
1474 def socket_connect_timeout(self) -> Optional[Union[float, int]]:
1475 return self._socket_connect_timeout
1477 @socket_connect_timeout.setter
1478 def socket_connect_timeout(self, value: Optional[Union[float, int]]):
1479 self._socket_connect_timeout = value
1481 def extract_connection_details(self) -> str:
1482 socket_address = None
1483 if self._sock is None:
1484 return "not connected"
1485 try:
1486 socket_address = self._sock.getsockname() if self._sock else None
1487 socket_address = socket_address[1] if socket_address else None
1488 except (AttributeError, OSError):
1489 pass
1491 return f"connected to ip {self.get_resolved_ip()}, local socket port: {socket_address}"
1494class Connection(AbstractConnection):
1495 "Manages TCP communication to and from a Redis server"
1497 def __init__(
1498 self,
1499 host="localhost",
1500 port=6379,
1501 socket_keepalive=True,
1502 socket_keepalive_options=SENTINEL,
1503 socket_type=0,
1504 **kwargs,
1505 ):
1506 """
1507 Initialize a TCP connection.
1509 Parameters
1510 ----------
1511 socket_keepalive : bool
1512 If `True`, TCP keepalive is enabled for TCP socket connections.
1513 socket_keepalive_options : Mapping[int, int | bytes] | object | None
1514 Mapping of TCP keepalive socket option constants to values, for
1515 example `{socket.TCP_KEEPIDLE: 30}`. If left unspecified, redis-py
1516 uses TCP keepalive defaults when `socket_keepalive` is enabled:
1517 idle 30 seconds, interval 5 seconds, and 3 probes. Platform-specific
1518 options that are not available are skipped. Pass `None` or `{}` to
1519 avoid setting additional TCP keepalive options.
1520 """
1521 self._host = host
1522 self.port = int(port)
1523 self.socket_keepalive = socket_keepalive
1524 if socket_keepalive_options is SENTINEL:
1525 socket_keepalive_options = get_default_socket_keepalive_options()
1526 self.socket_keepalive_options = socket_keepalive_options or {}
1527 self.socket_type = socket_type
1528 super().__init__(**kwargs)
1530 def repr_pieces(self):
1531 pieces = [("host", self.host), ("port", self.port), ("db", self.db)]
1532 if self.client_name:
1533 pieces.append(("client_name", self.client_name))
1534 return pieces
1536 def _connect(self):
1537 "Create a TCP socket connection"
1538 # we want to mimic what socket.create_connection does to support
1539 # ipv4/ipv6, but we want to set options prior to calling
1540 # socket.connect()
1541 err = None
1543 for res in socket.getaddrinfo(
1544 self.host, self.port, self.socket_type, socket.SOCK_STREAM
1545 ):
1546 family, socktype, proto, canonname, socket_address = res
1547 sock = None
1548 try:
1549 sock = socket.socket(family, socktype, proto)
1550 # TCP_NODELAY
1551 sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
1553 # TCP_KEEPALIVE
1554 if self.socket_keepalive:
1555 sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
1556 for k, v in self.socket_keepalive_options.items():
1557 sock.setsockopt(socket.IPPROTO_TCP, k, v)
1559 # set the socket_connect_timeout before we connect
1560 sock.settimeout(self.socket_connect_timeout)
1562 # connect
1563 sock.connect(socket_address)
1565 # set the socket_timeout now that we're connected
1566 sock.settimeout(self.socket_timeout)
1567 return sock
1569 except OSError as _:
1570 err = _
1571 if sock is not None:
1572 try:
1573 sock.shutdown(socket.SHUT_RDWR) # ensure a clean close
1574 except OSError:
1575 pass
1576 sock.close()
1578 if err is not None:
1579 raise err
1580 raise OSError("socket.getaddrinfo returned an empty list")
1582 def _host_error(self):
1583 return f"{self.host}:{self.port}"
1585 @property
1586 def host(self) -> str:
1587 return self._host
1589 @host.setter
1590 def host(self, value: str):
1591 self._host = value
1594class CacheProxyConnection(MaintNotificationsAbstractConnection, ConnectionInterface):
1595 DUMMY_CACHE_VALUE = b"foo"
1596 MIN_ALLOWED_VERSION = "7.4.0"
1597 DEFAULT_SERVER_NAME = "redis"
1599 def __init__(
1600 self,
1601 conn: ConnectionInterface,
1602 cache: CacheInterface,
1603 pool_lock: threading.RLock,
1604 ):
1605 self.pid = os.getpid()
1606 self._conn = conn
1607 self.retry = self._conn.retry
1608 self.host = self._conn.host
1609 self.port = self._conn.port
1610 self.db = self._conn.db
1611 self._event_dispatcher = self._conn._event_dispatcher
1612 self.credential_provider = conn.credential_provider
1613 self._pool_lock = pool_lock
1614 self._cache = cache
1615 self._cache_lock = threading.RLock()
1616 self._current_command_cache_key = None
1617 self._current_options = None
1618 self.register_connect_callback(self._enable_tracking_callback)
1620 if isinstance(self._conn, MaintNotificationsAbstractConnection):
1621 MaintNotificationsAbstractConnection.__init__(
1622 self,
1623 self._conn.maint_notifications_config,
1624 self._conn._maint_notifications_pool_handler,
1625 self._conn.maintenance_state,
1626 self._conn.maintenance_notification_hash,
1627 self._conn.host,
1628 self._conn.socket_timeout,
1629 self._conn.socket_connect_timeout,
1630 self._conn._oss_cluster_maint_notifications_handler,
1631 self._conn._get_parser(),
1632 event_dispatcher=self._conn.event_dispatcher,
1633 )
1635 def repr_pieces(self):
1636 return self._conn.repr_pieces()
1638 @property
1639 def is_connected(self) -> bool:
1640 return self._conn.is_connected
1642 def register_connect_callback(self, callback):
1643 self._conn.register_connect_callback(callback)
1645 def deregister_connect_callback(self, callback):
1646 self._conn.deregister_connect_callback(callback)
1648 def set_parser(self, parser_class):
1649 self._conn.set_parser(parser_class)
1651 def set_maint_notifications_pool_handler_for_connection(
1652 self, maint_notifications_pool_handler
1653 ):
1654 if isinstance(self._conn, MaintNotificationsAbstractConnection):
1655 self._conn.set_maint_notifications_pool_handler_for_connection(
1656 maint_notifications_pool_handler
1657 )
1659 def set_maint_notifications_cluster_handler_for_connection(
1660 self, oss_cluster_maint_notifications_handler
1661 ):
1662 if isinstance(self._conn, MaintNotificationsAbstractConnection):
1663 self._conn.set_maint_notifications_cluster_handler_for_connection(
1664 oss_cluster_maint_notifications_handler
1665 )
1667 def get_protocol(self):
1668 return self._conn.get_protocol()
1670 def connect(self):
1671 self._conn.connect()
1673 server_name = self._conn.handshake_metadata.get(b"server", None)
1674 if server_name is None:
1675 server_name = self._conn.handshake_metadata.get("server", None)
1676 server_ver = self._conn.handshake_metadata.get(b"version", None)
1677 if server_ver is None:
1678 server_ver = self._conn.handshake_metadata.get("version", None)
1679 if server_ver is None or server_name is None:
1680 raise ConnectionError("Cannot retrieve information about server version")
1682 server_ver = ensure_string(server_ver)
1683 server_name = ensure_string(server_name)
1685 if (
1686 server_name != self.DEFAULT_SERVER_NAME
1687 or compare_versions(server_ver, self.MIN_ALLOWED_VERSION) == 1
1688 ):
1689 raise ConnectionError(
1690 "To maximize compatibility with all Redis products, client-side caching is supported by Redis 7.4 or later" # noqa: E501
1691 )
1693 def on_connect(self):
1694 self._conn.on_connect()
1696 def disconnect(self, *args, **kwargs):
1697 with self._cache_lock:
1698 self._cache.flush()
1699 self._conn.disconnect(*args, **kwargs)
1701 def check_health(self):
1702 self._conn.check_health()
1704 def send_packed_command(self, command, check_health=True):
1705 # TODO: Investigate if it's possible to unpack command
1706 # or extract keys from packed command
1707 self._conn.send_packed_command(command)
1709 def send_command(self, *args, **kwargs):
1710 self._process_pending_invalidations()
1712 with self._cache_lock:
1713 # Command is write command or not allowed
1714 # to be cached.
1715 if not self._cache.is_cachable(
1716 CacheKey(command=args[0], redis_keys=(), redis_args=())
1717 ):
1718 self._current_command_cache_key = None
1719 self._conn.send_command(*args, **kwargs)
1720 return
1722 if kwargs.get("keys") is None:
1723 raise ValueError("Cannot create cache key.")
1725 # Creates cache key.
1726 self._current_command_cache_key = CacheKey(
1727 command=args[0], redis_keys=tuple(kwargs.get("keys")), redis_args=args
1728 )
1730 with self._cache_lock:
1731 # We have to trigger invalidation processing in case if
1732 # it was cached by another connection to avoid
1733 # queueing invalidations in stale connections.
1734 if self._cache.get(self._current_command_cache_key):
1735 entry = self._cache.get(self._current_command_cache_key)
1737 with self._pool_lock:
1738 while entry.connection_ref.can_read():
1739 try:
1740 entry.connection_ref.read_response(
1741 push_request=True,
1742 timeout=0,
1743 disconnect_on_error=False,
1744 )
1745 except TimeoutError:
1746 break
1748 # Re-check: if the entry was invalidated during the drain,
1749 # fall through to send the command over the network.
1750 if self._cache.get(self._current_command_cache_key):
1751 return
1753 # Set temporary entry value to prevent
1754 # race condition from another connection.
1755 self._cache.set(
1756 CacheEntry(
1757 cache_key=self._current_command_cache_key,
1758 cache_value=self.DUMMY_CACHE_VALUE,
1759 status=CacheEntryStatus.IN_PROGRESS,
1760 connection_ref=self._conn,
1761 )
1762 )
1764 # Send command over socket only if it's allowed
1765 # read-only command that not yet cached.
1766 self._conn.send_command(*args, **kwargs)
1768 def can_read(self, timeout: float = 0) -> bool:
1769 # TODO: Rename this API; it detects pending data or dirty/closed
1770 # connection state, not only whether application data can be read.
1771 return self._conn.can_read(timeout)
1773 def read_response(
1774 self,
1775 disable_decoding=False,
1776 *,
1777 timeout: Union[float, object] = SENTINEL,
1778 disconnect_on_error=True,
1779 push_request=False,
1780 ):
1781 with self._cache_lock:
1782 # Check if command response exists in a cache and it's not in progress.
1783 if self._current_command_cache_key is not None:
1784 if (
1785 self._cache.get(self._current_command_cache_key) is not None
1786 and self._cache.get(self._current_command_cache_key).status
1787 != CacheEntryStatus.IN_PROGRESS
1788 ):
1789 res = copy.deepcopy(
1790 self._cache.get(self._current_command_cache_key).cache_value
1791 )
1792 self._current_command_cache_key = None
1793 record_csc_request(
1794 result=CSCResult.HIT,
1795 )
1796 record_csc_network_saved(
1797 bytes_saved=len(res) if hasattr(res, "__len__") else 0,
1798 )
1799 return res
1800 record_csc_request(
1801 result=CSCResult.MISS,
1802 )
1804 response = self._conn.read_response(
1805 disable_decoding=disable_decoding,
1806 timeout=timeout,
1807 disconnect_on_error=disconnect_on_error,
1808 push_request=push_request,
1809 )
1811 with self._cache_lock:
1812 # Prevent not-allowed command from caching.
1813 if self._current_command_cache_key is None:
1814 return response
1815 # If response is None prevent from caching.
1816 if response is None:
1817 self._cache.delete_by_cache_keys([self._current_command_cache_key])
1818 return response
1820 cache_entry = self._cache.get(self._current_command_cache_key)
1822 # Cache only responses that still valid
1823 # and wasn't invalidated by another connection in meantime.
1824 if cache_entry is not None:
1825 cache_entry.status = CacheEntryStatus.VALID
1826 cache_entry.cache_value = response
1827 self._cache.set(cache_entry)
1829 self._current_command_cache_key = None
1831 return response
1833 def pack_command(self, *args):
1834 return self._conn.pack_command(*args)
1836 def pack_commands(self, commands):
1837 return self._conn.pack_commands(commands)
1839 @property
1840 def handshake_metadata(self) -> Union[Dict[bytes, bytes], Dict[str, str]]:
1841 return self._conn.handshake_metadata
1843 def set_re_auth_token(self, token: TokenInterface):
1844 self._conn.set_re_auth_token(token)
1846 def re_auth(self):
1847 self._conn.re_auth()
1849 def mark_for_reconnect(self):
1850 self._conn.mark_for_reconnect()
1852 def should_reconnect(self):
1853 return self._conn.should_reconnect()
1855 def reset_should_reconnect(self):
1856 self._conn.reset_should_reconnect()
1858 @property
1859 def host(self) -> str:
1860 return self._conn.host
1862 @host.setter
1863 def host(self, value: str):
1864 self._conn.host = value
1866 @property
1867 def socket_timeout(self) -> Optional[Union[float, int]]:
1868 return self._conn.socket_timeout
1870 @socket_timeout.setter
1871 def socket_timeout(self, value: Optional[Union[float, int]]):
1872 self._conn.socket_timeout = value
1874 @property
1875 def socket_connect_timeout(self) -> Optional[Union[float, int]]:
1876 return self._conn.socket_connect_timeout
1878 @socket_connect_timeout.setter
1879 def socket_connect_timeout(self, value: Optional[Union[float, int]]):
1880 self._conn.socket_connect_timeout = value
1882 @property
1883 def _maint_notifications_connection_handler(
1884 self,
1885 ) -> Optional[MaintNotificationsConnectionHandler]:
1886 if isinstance(self._conn, MaintNotificationsAbstractConnection):
1887 return self._conn._maint_notifications_connection_handler
1889 @_maint_notifications_connection_handler.setter
1890 def _maint_notifications_connection_handler(
1891 self, value: Optional[MaintNotificationsConnectionHandler]
1892 ):
1893 self._conn._maint_notifications_connection_handler = value
1895 def _get_socket(self) -> Optional[socket.socket]:
1896 if isinstance(self._conn, MaintNotificationsAbstractConnection):
1897 return self._conn._get_socket()
1898 else:
1899 raise NotImplementedError(
1900 "Maintenance notifications are not supported by this connection type"
1901 )
1903 def _get_maint_notifications_connection_instance(
1904 self, connection
1905 ) -> MaintNotificationsAbstractConnection:
1906 """
1907 Validate that connection instance supports maintenance notifications.
1908 With this helper method we ensure that we are working
1909 with the correct connection type.
1910 After twe validate that connection instance supports maintenance notifications
1911 we can safely return the connection instance
1912 as MaintNotificationsAbstractConnection.
1913 """
1914 if not isinstance(connection, MaintNotificationsAbstractConnection):
1915 raise NotImplementedError(
1916 "Maintenance notifications are not supported by this connection type"
1917 )
1918 else:
1919 return connection
1921 @property
1922 def maintenance_state(self) -> MaintenanceState:
1923 con = self._get_maint_notifications_connection_instance(self._conn)
1924 return con.maintenance_state
1926 @maintenance_state.setter
1927 def maintenance_state(self, state: MaintenanceState):
1928 con = self._get_maint_notifications_connection_instance(self._conn)
1929 con.maintenance_state = state
1931 def getpeername(self):
1932 con = self._get_maint_notifications_connection_instance(self._conn)
1933 return con.getpeername()
1935 def get_resolved_ip(self):
1936 con = self._get_maint_notifications_connection_instance(self._conn)
1937 return con.get_resolved_ip()
1939 def update_current_socket_timeout(self, relaxed_timeout: Optional[float] = None):
1940 con = self._get_maint_notifications_connection_instance(self._conn)
1941 con.update_current_socket_timeout(relaxed_timeout)
1943 def set_tmp_settings(
1944 self,
1945 tmp_host_address: Optional[str] = None,
1946 tmp_relaxed_timeout: Optional[float] = None,
1947 ):
1948 con = self._get_maint_notifications_connection_instance(self._conn)
1949 con.set_tmp_settings(tmp_host_address, tmp_relaxed_timeout)
1951 def reset_tmp_settings(
1952 self,
1953 reset_host_address: bool = False,
1954 reset_relaxed_timeout: bool = False,
1955 ):
1956 con = self._get_maint_notifications_connection_instance(self._conn)
1957 con.reset_tmp_settings(reset_host_address, reset_relaxed_timeout)
1959 def _connect(self):
1960 self._conn._connect()
1962 def _host_error(self):
1963 self._conn._host_error()
1965 def _enable_tracking_callback(self, conn: ConnectionInterface) -> None:
1966 conn.send_command("CLIENT", "TRACKING", "ON")
1967 conn.read_response()
1968 conn._parser.set_invalidation_push_handler(self._on_invalidation_callback)
1970 def _process_pending_invalidations(self):
1971 while self.can_read():
1972 try:
1973 self._conn.read_response(
1974 push_request=True, timeout=0, disconnect_on_error=False
1975 )
1976 except TimeoutError:
1977 break
1979 def _on_invalidation_callback(self, data: List[Union[str, Optional[List[bytes]]]]):
1980 with self._cache_lock:
1981 # Flush cache when DB flushed on server-side
1982 if data[1] is None:
1983 self._cache.flush()
1984 else:
1985 keys_deleted = self._cache.delete_by_redis_keys(data[1])
1987 if len(keys_deleted) > 0:
1988 record_csc_eviction(
1989 count=len(keys_deleted),
1990 reason=CSCReason.INVALIDATION,
1991 )
1993 def extract_connection_details(self) -> str:
1994 return self._conn.extract_connection_details()
1997class SSLConnection(Connection):
1998 """Manages SSL connections to and from the Redis server(s).
1999 This class extends the Connection class, adding SSL functionality, and making
2000 use of ssl.SSLContext (https://docs.python.org/3/library/ssl.html#ssl.SSLContext)
2001 """ # noqa
2003 def __init__(
2004 self,
2005 ssl_keyfile=None,
2006 ssl_certfile=None,
2007 ssl_cert_reqs="required",
2008 ssl_include_verify_flags: Optional[List["VerifyFlags"]] = None,
2009 ssl_exclude_verify_flags: Optional[List["VerifyFlags"]] = None,
2010 ssl_ca_certs=None,
2011 ssl_ca_data=None,
2012 ssl_check_hostname=True,
2013 ssl_ca_path=None,
2014 ssl_password=None,
2015 ssl_validate_ocsp=False,
2016 ssl_validate_ocsp_stapled=False,
2017 ssl_ocsp_context=None,
2018 ssl_ocsp_expected_cert=None,
2019 ssl_min_version=None,
2020 ssl_ciphers=None,
2021 **kwargs,
2022 ):
2023 """Constructor
2025 Args:
2026 ssl_keyfile: Path to an ssl private key. Defaults to None.
2027 ssl_certfile: Path to an ssl certificate. Defaults to None.
2028 ssl_cert_reqs: The string value for the SSLContext.verify_mode (none, optional, required),
2029 or an ssl.VerifyMode. Defaults to "required".
2030 ssl_include_verify_flags: A list of flags to be included in the SSLContext.verify_flags. Defaults to None.
2031 ssl_exclude_verify_flags: A list of flags to be excluded from the SSLContext.verify_flags. Defaults to None.
2032 ssl_ca_certs: The path to a file of concatenated CA certificates in PEM format. Defaults to None.
2033 ssl_ca_data: Either an ASCII string of one or more PEM-encoded certificates or a bytes-like object of DER-encoded certificates.
2034 ssl_check_hostname: If set, match the hostname during the SSL handshake. Defaults to True.
2035 ssl_ca_path: The path to a directory containing several CA certificates in PEM format. Defaults to None.
2036 ssl_password: Password for unlocking an encrypted private key. Defaults to None.
2038 ssl_validate_ocsp: If set, perform a full ocsp validation (i.e not a stapled verification)
2039 ssl_validate_ocsp_stapled: If set, perform a validation on a stapled ocsp response
2040 ssl_ocsp_context: A fully initialized OpenSSL.SSL.Context object to be used in verifying the ssl_ocsp_expected_cert
2041 ssl_ocsp_expected_cert: A PEM armoured string containing the expected certificate to be returned from the ocsp verification service.
2042 ssl_min_version: The lowest supported SSL version. It affects the supported SSL versions of the SSLContext. None leaves the default provided by ssl module.
2043 ssl_ciphers: A string listing the ciphers that are allowed to be used. Defaults to None, which means that the default ciphers are used. See https://docs.python.org/3/library/ssl.html#ssl.SSLContext.set_ciphers for more information.
2045 Raises:
2046 RedisError
2047 """ # noqa
2048 if not SSL_AVAILABLE:
2049 raise RedisError("Python wasn't built with SSL support")
2051 self.keyfile = ssl_keyfile
2052 self.certfile = ssl_certfile
2053 if ssl_cert_reqs is None:
2054 ssl_cert_reqs = ssl.CERT_NONE
2055 elif isinstance(ssl_cert_reqs, str):
2056 CERT_REQS = { # noqa: N806
2057 "none": ssl.CERT_NONE,
2058 "optional": ssl.CERT_OPTIONAL,
2059 "required": ssl.CERT_REQUIRED,
2060 }
2061 if ssl_cert_reqs not in CERT_REQS:
2062 raise RedisError(
2063 f"Invalid SSL Certificate Requirements Flag: {ssl_cert_reqs}"
2064 )
2065 ssl_cert_reqs = CERT_REQS[ssl_cert_reqs]
2066 self.cert_reqs = ssl_cert_reqs
2067 self.ssl_include_verify_flags = ssl_include_verify_flags
2068 self.ssl_exclude_verify_flags = ssl_exclude_verify_flags
2069 self.ca_certs = ssl_ca_certs
2070 self.ca_data = ssl_ca_data
2071 self.ca_path = ssl_ca_path
2072 self.check_hostname = (
2073 ssl_check_hostname if self.cert_reqs != ssl.CERT_NONE else False
2074 )
2075 self.certificate_password = ssl_password
2076 self.ssl_validate_ocsp = ssl_validate_ocsp
2077 self.ssl_validate_ocsp_stapled = ssl_validate_ocsp_stapled
2078 self.ssl_ocsp_context = ssl_ocsp_context
2079 self.ssl_ocsp_expected_cert = ssl_ocsp_expected_cert
2080 self.ssl_min_version = ssl_min_version
2081 self.ssl_ciphers = ssl_ciphers
2082 super().__init__(**kwargs)
2084 def _connect(self):
2085 """
2086 Wrap the socket with SSL support, handling potential errors.
2087 """
2088 sock = super()._connect()
2089 try:
2090 return self._wrap_socket_with_ssl(sock)
2091 except (OSError, RedisError):
2092 sock.close()
2093 raise
2095 def _wrap_socket_with_ssl(self, sock):
2096 """
2097 Wraps the socket with SSL support.
2099 Args:
2100 sock: The plain socket to wrap with SSL.
2102 Returns:
2103 An SSL wrapped socket.
2104 """
2105 context = ssl.create_default_context()
2106 context.check_hostname = self.check_hostname
2107 context.verify_mode = self.cert_reqs
2108 if self.ssl_include_verify_flags:
2109 for flag in self.ssl_include_verify_flags:
2110 context.verify_flags |= flag
2111 if self.ssl_exclude_verify_flags:
2112 for flag in self.ssl_exclude_verify_flags:
2113 context.verify_flags &= ~flag
2114 if self.certfile or self.keyfile:
2115 context.load_cert_chain(
2116 certfile=self.certfile,
2117 keyfile=self.keyfile,
2118 password=self.certificate_password,
2119 )
2120 if (
2121 self.ca_certs is not None
2122 or self.ca_path is not None
2123 or self.ca_data is not None
2124 ):
2125 context.load_verify_locations(
2126 cafile=self.ca_certs, capath=self.ca_path, cadata=self.ca_data
2127 )
2128 if self.ssl_min_version is not None:
2129 context.minimum_version = self.ssl_min_version
2130 if self.ssl_ciphers:
2131 context.set_ciphers(self.ssl_ciphers)
2132 if self.ssl_validate_ocsp is True and CRYPTOGRAPHY_AVAILABLE is False:
2133 raise RedisError("cryptography is not installed.")
2135 if self.ssl_validate_ocsp_stapled and self.ssl_validate_ocsp:
2136 raise RedisError(
2137 "Either an OCSP staple or pure OCSP connection must be validated "
2138 "- not both."
2139 )
2141 sslsock = context.wrap_socket(sock, server_hostname=self.host)
2143 # validation for the stapled case
2144 if self.ssl_validate_ocsp_stapled:
2145 import OpenSSL
2147 from .ocsp import ocsp_staple_verifier
2149 # if a context is provided use it - otherwise, a basic context
2150 if self.ssl_ocsp_context is None:
2151 staple_ctx = OpenSSL.SSL.Context(OpenSSL.SSL.SSLv23_METHOD)
2152 staple_ctx.use_certificate_file(self.certfile)
2153 staple_ctx.use_privatekey_file(self.keyfile)
2154 else:
2155 staple_ctx = self.ssl_ocsp_context
2157 staple_ctx.set_ocsp_client_callback(
2158 ocsp_staple_verifier, self.ssl_ocsp_expected_cert
2159 )
2161 # need another socket
2162 con = OpenSSL.SSL.Connection(staple_ctx, socket.socket())
2163 con.request_ocsp()
2164 con.connect((self.host, self.port))
2165 con.do_handshake()
2166 con.shutdown()
2167 return sslsock
2169 # pure ocsp validation
2170 if self.ssl_validate_ocsp is True and CRYPTOGRAPHY_AVAILABLE:
2171 from .ocsp import OCSPVerifier
2173 o = OCSPVerifier(sslsock, self.host, self.port, self.ca_certs)
2174 if o.is_valid():
2175 return sslsock
2176 else:
2177 raise ConnectionError("ocsp validation error")
2178 return sslsock
2181class UnixDomainSocketConnection(AbstractConnection):
2182 "Manages UDS communication to and from a Redis server"
2184 def __init__(self, path="", socket_timeout=DEFAULT_SOCKET_TIMEOUT, **kwargs):
2185 super().__init__(**kwargs)
2186 self.path = path
2187 self.socket_timeout = socket_timeout
2189 def repr_pieces(self):
2190 pieces = [("path", self.path), ("db", self.db)]
2191 if self.client_name:
2192 pieces.append(("client_name", self.client_name))
2193 return pieces
2195 def _connect(self):
2196 "Create a Unix domain socket connection"
2197 sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
2198 sock.settimeout(self.socket_connect_timeout)
2199 try:
2200 sock.connect(self.path)
2201 except OSError:
2202 # Prevent ResourceWarnings for unclosed sockets.
2203 try:
2204 sock.shutdown(socket.SHUT_RDWR) # ensure a clean close
2205 except OSError:
2206 pass
2207 sock.close()
2208 raise
2209 sock.settimeout(self.socket_timeout)
2210 return sock
2212 def _host_error(self):
2213 return self.path
2216FALSE_STRINGS = ("0", "F", "FALSE", "N", "NO")
2219def to_bool(value):
2220 if value is None or value == "":
2221 return None
2222 if isinstance(value, str) and value.upper() in FALSE_STRINGS:
2223 return False
2224 return bool(value)
2227def parse_ssl_verify_flags(value):
2228 # flags are passed in as a string representation of a list,
2229 # e.g. VERIFY_X509_STRICT, VERIFY_X509_PARTIAL_CHAIN
2230 verify_flags_str = value.replace("[", "").replace("]", "")
2232 verify_flags = []
2233 for flag in verify_flags_str.split(","):
2234 flag = flag.strip()
2235 if not hasattr(VerifyFlags, flag):
2236 raise ValueError(f"Invalid ssl verify flag: {flag}")
2237 verify_flags.append(getattr(VerifyFlags, flag))
2238 return verify_flags
2241URL_QUERY_ARGUMENT_PARSERS = {
2242 "db": int,
2243 "socket_timeout": float,
2244 "socket_connect_timeout": float,
2245 "socket_read_size": int,
2246 "socket_keepalive": to_bool,
2247 "retry_on_timeout": to_bool,
2248 "retry_on_error": list,
2249 "max_connections": int,
2250 "health_check_interval": int,
2251 "ssl_check_hostname": to_bool,
2252 "ssl_include_verify_flags": parse_ssl_verify_flags,
2253 "ssl_exclude_verify_flags": parse_ssl_verify_flags,
2254 "ssl_min_version": int,
2255 "timeout": float,
2256 "protocol": int,
2257 "legacy_responses": to_bool,
2258}
2261def parse_url(url):
2262 if not (
2263 url.startswith("redis://")
2264 or url.startswith("rediss://")
2265 or url.startswith("unix://")
2266 ):
2267 raise ValueError(
2268 "Redis URL must specify one of the following "
2269 "schemes (redis://, rediss://, unix://)"
2270 )
2272 url = urlparse(url)
2273 kwargs = {}
2275 for name, value in parse_qs(url.query).items():
2276 if value and len(value) > 0:
2277 value = unquote(value[0])
2278 parser = URL_QUERY_ARGUMENT_PARSERS.get(name)
2279 if parser:
2280 try:
2281 kwargs[name] = parser(value)
2282 except (TypeError, ValueError):
2283 raise ValueError(f"Invalid value for '{name}' in connection URL.")
2284 else:
2285 kwargs[name] = value
2287 if url.username:
2288 kwargs["username"] = unquote(url.username)
2289 if url.password:
2290 kwargs["password"] = unquote(url.password)
2292 # We only support redis://, rediss:// and unix:// schemes.
2293 if url.scheme == "unix":
2294 if url.path:
2295 kwargs["path"] = unquote(url.path)
2296 kwargs["connection_class"] = UnixDomainSocketConnection
2298 else: # implied: url.scheme in ("redis", "rediss"):
2299 if url.hostname:
2300 kwargs["host"] = unquote(url.hostname)
2301 if url.port:
2302 kwargs["port"] = int(url.port)
2304 # If there's a path argument, use it as the db argument if a
2305 # querystring value wasn't specified
2306 if url.path and "db" not in kwargs:
2307 try:
2308 kwargs["db"] = int(unquote(url.path).replace("/", ""))
2309 except (AttributeError, ValueError):
2310 pass
2312 if url.scheme == "rediss":
2313 kwargs["connection_class"] = SSLConnection
2315 return kwargs
2318_CP = TypeVar("_CP", bound="ConnectionPool")
2321class ConnectionPoolInterface(ABC):
2322 @abstractmethod
2323 def get_protocol(self):
2324 pass
2326 @abstractmethod
2327 def reset(self):
2328 pass
2330 @abstractmethod
2331 @deprecated_args(
2332 args_to_warn=["*"],
2333 reason="Use get_connection() without args instead",
2334 version="5.3.0",
2335 )
2336 def get_connection(
2337 self, command_name: Optional[str], *keys, **options
2338 ) -> ConnectionInterface:
2339 pass
2341 @abstractmethod
2342 def get_encoder(self):
2343 pass
2345 @abstractmethod
2346 def release(self, connection: ConnectionInterface):
2347 pass
2349 @abstractmethod
2350 def disconnect(self, inuse_connections: bool = True):
2351 pass
2353 @abstractmethod
2354 def close(self):
2355 pass
2357 @abstractmethod
2358 def set_retry(self, retry: Retry):
2359 pass
2361 @abstractmethod
2362 def re_auth_callback(self, token: TokenInterface):
2363 pass
2365 @abstractmethod
2366 def get_connection_count(self) -> list[tuple[int, dict]]:
2367 """
2368 Returns a connection count (both idle and in use).
2369 """
2370 pass
2373class MaintNotificationsAbstractConnectionPool:
2374 """
2375 Abstract class for handling maintenance notifications logic.
2376 This class is mixed into the ConnectionPool classes.
2378 This class is not intended to be used directly!
2380 All logic related to maintenance notifications and
2381 connection pool handling is encapsulated in this class.
2382 """
2384 def __init__(
2385 self,
2386 maint_notifications_config: Optional[MaintNotificationsConfig] = None,
2387 oss_cluster_maint_notifications_handler: Optional[
2388 OSSMaintNotificationsHandler
2389 ] = None,
2390 **kwargs,
2391 ):
2392 # Initialize maintenance notifications
2393 is_protocol_supported = check_protocol_version(kwargs.get("protocol"), 3)
2395 if maint_notifications_config is None and is_protocol_supported:
2396 maint_notifications_config = MaintNotificationsConfig()
2398 if maint_notifications_config and maint_notifications_config.enabled:
2399 if not is_protocol_supported:
2400 raise RedisError(
2401 "Maintenance notifications handlers on connection are only supported with RESP version 3"
2402 )
2404 self._event_dispatcher = kwargs.get("event_dispatcher", None)
2405 if self._event_dispatcher is None:
2406 self._event_dispatcher = EventDispatcher()
2408 self._maint_notifications_pool_handler = MaintNotificationsPoolHandler(
2409 self, maint_notifications_config
2410 )
2411 if oss_cluster_maint_notifications_handler:
2412 self._oss_cluster_maint_notifications_handler = (
2413 oss_cluster_maint_notifications_handler
2414 )
2415 self._update_connection_kwargs_for_maint_notifications(
2416 oss_cluster_maint_notifications_handler=self._oss_cluster_maint_notifications_handler
2417 )
2418 self._maint_notifications_pool_handler = None
2419 else:
2420 self._oss_cluster_maint_notifications_handler = None
2421 self._maint_notifications_pool_handler = MaintNotificationsPoolHandler(
2422 self, maint_notifications_config
2423 )
2425 self._update_connection_kwargs_for_maint_notifications(
2426 maint_notifications_pool_handler=self._maint_notifications_pool_handler
2427 )
2428 else:
2429 self._maint_notifications_pool_handler = None
2430 self._oss_cluster_maint_notifications_handler = None
2432 @property
2433 @abstractmethod
2434 def connection_kwargs(self) -> Dict[str, Any]:
2435 pass
2437 @connection_kwargs.setter
2438 @abstractmethod
2439 def connection_kwargs(self, value: Dict[str, Any]):
2440 pass
2442 @abstractmethod
2443 def _get_pool_lock(self) -> threading.RLock:
2444 pass
2446 @abstractmethod
2447 def _get_free_connections(self) -> Iterable["MaintNotificationsAbstractConnection"]:
2448 pass
2450 @abstractmethod
2451 def _get_in_use_connections(
2452 self,
2453 ) -> Iterable["MaintNotificationsAbstractConnection"]:
2454 pass
2456 def maint_notifications_enabled(self):
2457 """
2458 Returns:
2459 True if the maintenance notifications are enabled, False otherwise.
2460 The maintenance notifications config is stored in the pool handler.
2461 If the pool handler is not set, the maintenance notifications are not enabled.
2462 """
2463 if self._oss_cluster_maint_notifications_handler:
2464 maint_notifications_config = (
2465 self._oss_cluster_maint_notifications_handler.config
2466 )
2467 else:
2468 maint_notifications_config = (
2469 self._maint_notifications_pool_handler.config
2470 if self._maint_notifications_pool_handler
2471 else None
2472 )
2474 return maint_notifications_config and maint_notifications_config.enabled
2476 def update_maint_notifications_config(
2477 self,
2478 maint_notifications_config: MaintNotificationsConfig,
2479 oss_cluster_maint_notifications_handler: Optional[
2480 OSSMaintNotificationsHandler
2481 ] = None,
2482 ):
2483 """
2484 Updates the maintenance notifications configuration.
2485 This method should be called only if the pool was created
2486 without enabling the maintenance notifications and
2487 in a later point in time maintenance notifications
2488 are requested to be enabled.
2489 """
2490 if (
2491 self.maint_notifications_enabled()
2492 and not maint_notifications_config.enabled
2493 ):
2494 raise ValueError(
2495 "Cannot disable maintenance notifications after enabling them"
2496 )
2497 if oss_cluster_maint_notifications_handler:
2498 self._oss_cluster_maint_notifications_handler = (
2499 oss_cluster_maint_notifications_handler
2500 )
2501 else:
2502 # first update pool settings
2503 if not self._maint_notifications_pool_handler:
2504 self._maint_notifications_pool_handler = MaintNotificationsPoolHandler(
2505 self, maint_notifications_config
2506 )
2507 else:
2508 self._maint_notifications_pool_handler.config = (
2509 maint_notifications_config
2510 )
2512 # then update connection kwargs and existing connections
2513 self._update_connection_kwargs_for_maint_notifications(
2514 maint_notifications_pool_handler=self._maint_notifications_pool_handler,
2515 oss_cluster_maint_notifications_handler=self._oss_cluster_maint_notifications_handler,
2516 )
2517 self._update_maint_notifications_configs_for_connections(
2518 maint_notifications_pool_handler=self._maint_notifications_pool_handler,
2519 oss_cluster_maint_notifications_handler=self._oss_cluster_maint_notifications_handler,
2520 )
2522 def _update_connection_kwargs_for_maint_notifications(
2523 self,
2524 maint_notifications_pool_handler: Optional[
2525 MaintNotificationsPoolHandler
2526 ] = None,
2527 oss_cluster_maint_notifications_handler: Optional[
2528 OSSMaintNotificationsHandler
2529 ] = None,
2530 ):
2531 """
2532 Update the connection kwargs for all future connections.
2533 """
2534 if not self.maint_notifications_enabled():
2535 return
2536 if maint_notifications_pool_handler:
2537 self.connection_kwargs.update(
2538 {
2539 "maint_notifications_pool_handler": maint_notifications_pool_handler,
2540 "maint_notifications_config": maint_notifications_pool_handler.config,
2541 }
2542 )
2543 if oss_cluster_maint_notifications_handler:
2544 self.connection_kwargs.update(
2545 {
2546 "oss_cluster_maint_notifications_handler": oss_cluster_maint_notifications_handler,
2547 "maint_notifications_config": oss_cluster_maint_notifications_handler.config,
2548 }
2549 )
2551 # Store original connection parameters for maintenance notifications.
2552 if self.connection_kwargs.get("orig_host_address", None) is None:
2553 # If orig_host_address is None it means we haven't
2554 # configured the original values yet
2555 self.connection_kwargs.update(
2556 {
2557 "orig_host_address": self.connection_kwargs.get("host"),
2558 "orig_socket_timeout": self.connection_kwargs.get(
2559 "socket_timeout", DEFAULT_SOCKET_TIMEOUT
2560 ),
2561 "orig_socket_connect_timeout": self.connection_kwargs.get(
2562 "socket_connect_timeout", DEFAULT_SOCKET_CONNECT_TIMEOUT
2563 ),
2564 }
2565 )
2567 def _update_maint_notifications_configs_for_connections(
2568 self,
2569 maint_notifications_pool_handler: Optional[
2570 MaintNotificationsPoolHandler
2571 ] = None,
2572 oss_cluster_maint_notifications_handler: Optional[
2573 OSSMaintNotificationsHandler
2574 ] = None,
2575 ):
2576 """Update the maintenance notifications config for all connections in the pool."""
2577 with self._get_pool_lock():
2578 for conn in self._get_free_connections():
2579 if oss_cluster_maint_notifications_handler:
2580 # set cluster handler for conn
2581 conn.set_maint_notifications_cluster_handler_for_connection(
2582 oss_cluster_maint_notifications_handler
2583 )
2584 conn.maint_notifications_config = (
2585 oss_cluster_maint_notifications_handler.config
2586 )
2587 elif maint_notifications_pool_handler:
2588 conn.set_maint_notifications_pool_handler_for_connection(
2589 maint_notifications_pool_handler
2590 )
2591 conn.maint_notifications_config = (
2592 maint_notifications_pool_handler.config
2593 )
2594 else:
2595 raise ValueError(
2596 "Either maint_notifications_pool_handler or oss_cluster_maint_notifications_handler must be set"
2597 )
2598 conn.disconnect()
2599 for conn in self._get_in_use_connections():
2600 if oss_cluster_maint_notifications_handler:
2601 conn.maint_notifications_config = (
2602 oss_cluster_maint_notifications_handler.config
2603 )
2604 conn._configure_maintenance_notifications(
2605 oss_cluster_maint_notifications_handler=oss_cluster_maint_notifications_handler
2606 )
2607 elif maint_notifications_pool_handler:
2608 conn.set_maint_notifications_pool_handler_for_connection(
2609 maint_notifications_pool_handler
2610 )
2611 conn.maint_notifications_config = (
2612 maint_notifications_pool_handler.config
2613 )
2614 else:
2615 raise ValueError(
2616 "Either maint_notifications_pool_handler or oss_cluster_maint_notifications_handler must be set"
2617 )
2618 conn.mark_for_reconnect()
2620 def _should_update_connection(
2621 self,
2622 conn: "MaintNotificationsAbstractConnection",
2623 matching_pattern: Literal[
2624 "connected_address", "configured_address", "notification_hash"
2625 ] = "connected_address",
2626 matching_address: Optional[str] = None,
2627 matching_notification_hash: Optional[int] = None,
2628 ) -> bool:
2629 """
2630 Check if the connection should be updated based on the matching criteria.
2631 """
2632 if matching_pattern == "connected_address":
2633 if matching_address and conn.getpeername() != matching_address:
2634 return False
2635 elif matching_pattern == "configured_address":
2636 if matching_address and conn.host != matching_address:
2637 return False
2638 elif matching_pattern == "notification_hash":
2639 if (
2640 matching_notification_hash
2641 and conn.maintenance_notification_hash != matching_notification_hash
2642 ):
2643 return False
2644 return True
2646 def update_connection_settings(
2647 self,
2648 conn: "MaintNotificationsAbstractConnection",
2649 state: Optional["MaintenanceState"] = None,
2650 maintenance_notification_hash: Optional[int] = None,
2651 host_address: Optional[str] = None,
2652 relaxed_timeout: Optional[float] = None,
2653 update_notification_hash: bool = False,
2654 reset_host_address: bool = False,
2655 reset_relaxed_timeout: bool = False,
2656 ):
2657 """
2658 Update the settings for a single connection.
2659 """
2660 if state:
2661 conn.maintenance_state = state
2663 if update_notification_hash:
2664 # update the notification hash only if requested
2665 conn.maintenance_notification_hash = maintenance_notification_hash
2667 if host_address is not None:
2668 conn.set_tmp_settings(tmp_host_address=host_address)
2670 if relaxed_timeout is not None:
2671 conn.set_tmp_settings(tmp_relaxed_timeout=relaxed_timeout)
2673 if reset_relaxed_timeout or reset_host_address:
2674 conn.reset_tmp_settings(
2675 reset_host_address=reset_host_address,
2676 reset_relaxed_timeout=reset_relaxed_timeout,
2677 )
2679 conn.update_current_socket_timeout(relaxed_timeout)
2681 def update_connections_settings(
2682 self,
2683 state: Optional["MaintenanceState"] = None,
2684 maintenance_notification_hash: Optional[int] = None,
2685 host_address: Optional[str] = None,
2686 relaxed_timeout: Optional[float] = None,
2687 matching_address: Optional[str] = None,
2688 matching_notification_hash: Optional[int] = None,
2689 matching_pattern: Literal[
2690 "connected_address", "configured_address", "notification_hash"
2691 ] = "connected_address",
2692 update_notification_hash: bool = False,
2693 reset_host_address: bool = False,
2694 reset_relaxed_timeout: bool = False,
2695 include_free_connections: bool = True,
2696 ):
2697 """
2698 Update the settings for all matching connections in the pool.
2700 This method does not create new connections.
2701 This method does not affect the connection kwargs.
2703 :param state: The maintenance state to set for the connection.
2704 :param maintenance_notification_hash: The hash of the maintenance notification
2705 to set for the connection.
2706 :param host_address: The host address to set for the connection.
2707 :param relaxed_timeout: The relaxed timeout to set for the connection.
2708 :param matching_address: The address to match for the connection.
2709 :param matching_notification_hash: The notification hash to match for the connection.
2710 :param matching_pattern: The pattern to match for the connection.
2711 :param update_notification_hash: Whether to update the notification hash for the connection.
2712 :param reset_host_address: Whether to reset the host address to the original address.
2713 :param reset_relaxed_timeout: Whether to reset the relaxed timeout to the original timeout.
2714 :param include_free_connections: Whether to include free/available connections.
2715 """
2716 with self._get_pool_lock():
2717 for conn in self._get_in_use_connections():
2718 if self._should_update_connection(
2719 conn,
2720 matching_pattern,
2721 matching_address,
2722 matching_notification_hash,
2723 ):
2724 self.update_connection_settings(
2725 conn,
2726 state=state,
2727 maintenance_notification_hash=maintenance_notification_hash,
2728 host_address=host_address,
2729 relaxed_timeout=relaxed_timeout,
2730 update_notification_hash=update_notification_hash,
2731 reset_host_address=reset_host_address,
2732 reset_relaxed_timeout=reset_relaxed_timeout,
2733 )
2735 if include_free_connections:
2736 for conn in self._get_free_connections():
2737 if self._should_update_connection(
2738 conn,
2739 matching_pattern,
2740 matching_address,
2741 matching_notification_hash,
2742 ):
2743 self.update_connection_settings(
2744 conn,
2745 state=state,
2746 maintenance_notification_hash=maintenance_notification_hash,
2747 host_address=host_address,
2748 relaxed_timeout=relaxed_timeout,
2749 update_notification_hash=update_notification_hash,
2750 reset_host_address=reset_host_address,
2751 reset_relaxed_timeout=reset_relaxed_timeout,
2752 )
2754 def update_connection_kwargs(
2755 self,
2756 **kwargs,
2757 ):
2758 """
2759 Update the connection kwargs for all future connections.
2761 This method updates the connection kwargs for all future connections created by the pool.
2762 Existing connections are not affected.
2763 """
2764 self.connection_kwargs.update(kwargs)
2766 def update_active_connections_for_reconnect(
2767 self,
2768 moving_address_src: Optional[str] = None,
2769 ):
2770 """
2771 Mark all active connections for reconnect.
2772 This is used when a cluster node is migrated to a different address.
2774 :param moving_address_src: The address of the node that is being moved.
2775 """
2776 with self._get_pool_lock():
2777 for conn in self._get_in_use_connections():
2778 if self._should_update_connection(
2779 conn, "connected_address", moving_address_src
2780 ):
2781 conn.mark_for_reconnect()
2783 def disconnect_free_connections(
2784 self,
2785 moving_address_src: Optional[str] = None,
2786 ):
2787 """
2788 Disconnect all free/available connections.
2789 This is used when a cluster node is migrated to a different address.
2791 :param moving_address_src: The address of the node that is being moved.
2792 """
2793 with self._get_pool_lock():
2794 for conn in self._get_free_connections():
2795 if self._should_update_connection(
2796 conn, "connected_address", moving_address_src
2797 ):
2798 conn.disconnect()
2801class ConnectionPool(MaintNotificationsAbstractConnectionPool, ConnectionPoolInterface):
2802 """
2803 Create a connection pool. ``If max_connections`` is set, then this
2804 object raises :py:class:`~redis.exceptions.ConnectionError` when the pool's
2805 limit is reached.
2807 By default, TCP connections are created unless ``connection_class``
2808 is specified. Use class:`.UnixDomainSocketConnection` for
2809 unix sockets.
2810 :py:class:`~redis.SSLConnection` can be used for SSL enabled connections.
2812 If ``maint_notifications_config`` is provided, the connection pool will support
2813 maintenance notifications.
2814 Maintenance notifications are supported only with RESP3.
2815 If the ``maint_notifications_config`` is not provided but the ``protocol`` is 3,
2816 the maintenance notifications will be enabled by default.
2818 Any additional keyword arguments are passed to the constructor of
2819 ``connection_class``.
2820 """
2822 @classmethod
2823 def from_url(cls: Type[_CP], url: str, **kwargs) -> _CP:
2824 """
2825 Return a connection pool configured from the given URL.
2827 For example::
2829 redis://[[username]:[password]]@localhost:6379/0
2830 rediss://[[username]:[password]]@localhost:6379/0
2831 unix://[username@]/path/to/socket.sock?db=0[&password=password]
2833 Three URL schemes are supported:
2835 - `redis://` creates a TCP socket connection. See more at:
2836 <https://www.iana.org/assignments/uri-schemes/prov/redis>
2837 - `rediss://` creates a SSL wrapped TCP socket connection. See more at:
2838 <https://www.iana.org/assignments/uri-schemes/prov/rediss>
2839 - ``unix://``: creates a Unix Domain Socket connection.
2841 The username, password, hostname, path and all querystring values
2842 are passed through urllib.parse.unquote in order to replace any
2843 percent-encoded values with their corresponding characters.
2845 There are several ways to specify a database number. The first value
2846 found will be used:
2848 1. A ``db`` querystring option, e.g. redis://localhost?db=0
2849 2. If using the redis:// or rediss:// schemes, the path argument
2850 of the url, e.g. redis://localhost/0
2851 3. A ``db`` keyword argument to this function.
2853 If none of these options are specified, the default db=0 is used.
2855 All querystring options are cast to their appropriate Python types.
2856 Boolean arguments can be specified with string values "True"/"False"
2857 or "Yes"/"No". Values that cannot be properly cast cause a
2858 ``ValueError`` to be raised. Once parsed, the querystring arguments
2859 and keyword arguments are passed to the ``ConnectionPool``'s
2860 class initializer. In the case of conflicting arguments, querystring
2861 arguments always win.
2862 """
2863 url_options = parse_url(url)
2865 if "connection_class" in kwargs:
2866 url_options["connection_class"] = kwargs["connection_class"]
2868 kwargs.update(url_options)
2869 return cls(**kwargs)
2871 def __init__(
2872 self,
2873 connection_class=Connection,
2874 max_connections: Optional[int] = None,
2875 cache_factory: Optional[CacheFactoryInterface] = None,
2876 maint_notifications_config: Optional[MaintNotificationsConfig] = None,
2877 **connection_kwargs,
2878 ):
2879 max_connections = max_connections or 100
2880 if not isinstance(max_connections, int) or max_connections < 0:
2881 raise ValueError('"max_connections" must be a positive integer')
2883 self.connection_class = connection_class
2884 self._connection_kwargs = connection_kwargs
2885 self.max_connections = max_connections
2886 self.cache = None
2887 self._cache_factory = cache_factory
2889 try:
2890 supports_maint_notifications = issubclass(
2891 connection_class, MaintNotificationsAbstractConnection
2892 )
2893 is_unix_domain_socket_connection = issubclass(
2894 connection_class, UnixDomainSocketConnection
2895 )
2896 except TypeError:
2897 supports_maint_notifications = False
2898 is_unix_domain_socket_connection = False
2900 if is_unix_domain_socket_connection or not supports_maint_notifications:
2901 if (
2902 maint_notifications_config
2903 and maint_notifications_config.enabled is True
2904 ):
2905 raise RedisError(
2906 "Maintenance notifications are not supported with "
2907 f"{connection_class}"
2908 )
2909 maint_notifications_config = MaintNotificationsConfig(enabled=False)
2911 self._event_dispatcher = self._connection_kwargs.get("event_dispatcher", None)
2912 if self._event_dispatcher is None:
2913 self._event_dispatcher = EventDispatcher()
2915 if connection_kwargs.get("cache_config") or connection_kwargs.get("cache"):
2916 if not check_protocol_version(self._connection_kwargs.get("protocol"), 3):
2917 raise RedisError("Client caching is only supported with RESP version 3")
2919 cache = self._connection_kwargs.get("cache")
2921 if cache is not None:
2922 if not isinstance(cache, CacheInterface):
2923 raise ValueError("Cache must implement CacheInterface")
2925 self.cache = cache
2926 else:
2927 if self._cache_factory is not None:
2928 self.cache = CacheProxy(self._cache_factory.get_cache())
2929 else:
2930 self.cache = CacheFactory(
2931 self._connection_kwargs.get("cache_config")
2932 ).get_cache()
2934 init_csc_items()
2935 register_csc_items_callback(
2936 callback=lambda: self.cache.size,
2937 pool_name=get_pool_name(self),
2938 )
2940 connection_kwargs.pop("cache", None)
2941 connection_kwargs.pop("cache_config", None)
2943 # a lock to protect the critical section in _checkpid().
2944 # this lock is acquired when the process id changes, such as
2945 # after a fork. during this time, multiple threads in the child
2946 # process could attempt to acquire this lock. the first thread
2947 # to acquire the lock will reset the data structures and lock
2948 # object of this pool. subsequent threads acquiring this lock
2949 # will notice the first thread already did the work and simply
2950 # release the lock.
2952 self._fork_lock = threading.RLock()
2953 self._lock = threading.RLock()
2955 # Generate unique pool ID for observability (matches go-redis behavior)
2956 import secrets
2958 self._pool_id = secrets.token_hex(4)
2960 MaintNotificationsAbstractConnectionPool.__init__(
2961 self,
2962 maint_notifications_config=maint_notifications_config,
2963 **connection_kwargs,
2964 )
2966 self.reset()
2968 # Keys that should be redacted in __repr__ to avoid exposing sensitive information
2969 SENSITIVE_REPR_KEYS = frozenset(
2970 {
2971 "password",
2972 "username",
2973 "ssl_password",
2974 "credential_provider",
2975 }
2976 )
2978 def __repr__(self) -> str:
2979 conn_kwargs = ",".join(
2980 [
2981 f"{k}={'<REDACTED>' if k in self.SENSITIVE_REPR_KEYS else v}"
2982 for k, v in self.connection_kwargs.items()
2983 ]
2984 )
2985 return (
2986 f"<{self.__class__.__module__}.{self.__class__.__name__}"
2987 f"(<{self.connection_class.__module__}.{self.connection_class.__name__}"
2988 f"({conn_kwargs})>)>"
2989 )
2991 @property
2992 def connection_kwargs(self) -> Dict[str, Any]:
2993 return self._connection_kwargs
2995 @connection_kwargs.setter
2996 def connection_kwargs(self, value: Dict[str, Any]):
2997 self._connection_kwargs = value
2999 def get_protocol(self):
3000 """
3001 Returns:
3002 The RESP protocol version, or ``None`` if the protocol is not specified,
3003 in which case the server default will be used.
3004 """
3005 return self.connection_kwargs.get("protocol", None)
3007 def reset(self) -> None:
3008 # Record metrics for connections being removed before clearing
3009 # (only if attributes exist - they won't during __init__)
3010 if hasattr(self, "_available_connections") and hasattr(
3011 self, "_in_use_connections"
3012 ):
3013 with self._lock:
3014 idle_count = len(self._available_connections)
3015 in_use_count = len(self._in_use_connections)
3016 if idle_count > 0 or in_use_count > 0:
3017 pool_name = get_pool_name(self)
3018 if idle_count > 0:
3019 record_connection_count(
3020 pool_name=pool_name,
3021 connection_state=ConnectionState.IDLE,
3022 counter=-idle_count,
3023 )
3024 if in_use_count > 0:
3025 record_connection_count(
3026 pool_name=pool_name,
3027 connection_state=ConnectionState.USED,
3028 counter=-in_use_count,
3029 )
3031 self._created_connections = 0
3032 self._available_connections = []
3033 self._in_use_connections = set()
3035 # this must be the last operation in this method. while reset() is
3036 # called when holding _fork_lock, other threads in this process
3037 # can call _checkpid() which compares self.pid and os.getpid() without
3038 # holding any lock (for performance reasons). keeping this assignment
3039 # as the last operation ensures that those other threads will also
3040 # notice a pid difference and block waiting for the first thread to
3041 # release _fork_lock. when each of these threads eventually acquire
3042 # _fork_lock, they will notice that another thread already called
3043 # reset() and they will immediately release _fork_lock and continue on.
3044 self.pid = os.getpid()
3046 def __del__(self) -> None:
3047 """Clean up connection pool and record metrics when garbage collected."""
3048 try:
3049 if not hasattr(self, "_available_connections") or not hasattr(
3050 self, "_in_use_connections"
3051 ):
3052 return
3053 # Record metrics for all connections being removed
3054 idle_count = len(self._available_connections)
3055 in_use_count = len(self._in_use_connections)
3056 if idle_count > 0 or in_use_count > 0:
3057 pool_name = get_pool_name(self)
3058 if idle_count > 0:
3059 record_connection_count(
3060 pool_name=pool_name,
3061 connection_state=ConnectionState.IDLE,
3062 counter=-idle_count,
3063 )
3064 if in_use_count > 0:
3065 record_connection_count(
3066 pool_name=pool_name,
3067 connection_state=ConnectionState.USED,
3068 counter=-in_use_count,
3069 )
3070 except Exception:
3071 pass
3073 def _checkpid(self) -> None:
3074 # _checkpid() attempts to keep ConnectionPool fork-safe on modern
3075 # systems. this is called by all ConnectionPool methods that
3076 # manipulate the pool's state such as get_connection() and release().
3077 #
3078 # _checkpid() determines whether the process has forked by comparing
3079 # the current process id to the process id saved on the ConnectionPool
3080 # instance. if these values are the same, _checkpid() simply returns.
3081 #
3082 # when the process ids differ, _checkpid() assumes that the process
3083 # has forked and that we're now running in the child process. the child
3084 # process cannot use the parent's file descriptors (e.g., sockets).
3085 # therefore, when _checkpid() sees the process id change, it calls
3086 # reset() in order to reinitialize the child's ConnectionPool. this
3087 # will cause the child to make all new connection objects.
3088 #
3089 # _checkpid() is protected by self._fork_lock to ensure that multiple
3090 # threads in the child process do not call reset() multiple times.
3091 #
3092 # there is an extremely small chance this could fail in the following
3093 # scenario:
3094 # 1. process A calls _checkpid() for the first time and acquires
3095 # self._fork_lock.
3096 # 2. while holding self._fork_lock, process A forks (the fork()
3097 # could happen in a different thread owned by process A)
3098 # 3. process B (the forked child process) inherits the
3099 # ConnectionPool's state from the parent. that state includes
3100 # a locked _fork_lock. process B will not be notified when
3101 # process A releases the _fork_lock and will thus never be
3102 # able to acquire the _fork_lock.
3103 #
3104 # to mitigate this possible deadlock, _checkpid() will only wait 5
3105 # seconds to acquire _fork_lock. if _fork_lock cannot be acquired in
3106 # that time it is assumed that the child is deadlocked and a
3107 # redis.ChildDeadlockedError error is raised.
3108 if self.pid != os.getpid():
3109 acquired = self._fork_lock.acquire(timeout=5)
3110 if not acquired:
3111 raise ChildDeadlockedError
3112 # reset() the instance for the new process if another thread
3113 # hasn't already done so
3114 try:
3115 if self.pid != os.getpid():
3116 self.reset()
3117 finally:
3118 self._fork_lock.release()
3120 @deprecated_args(
3121 args_to_warn=["*"],
3122 reason="Use get_connection() without args instead",
3123 version="5.3.0",
3124 )
3125 def get_connection(self, command_name=None, *keys, **options) -> "Connection":
3126 "Get a connection from the pool"
3128 # Start timing for observability
3129 self._checkpid()
3130 is_created = False
3132 with self._lock:
3133 try:
3134 connection = self._available_connections.pop()
3135 except IndexError:
3136 # Start timing for observability
3137 start_time_created = time.monotonic()
3139 connection = self.make_connection()
3140 is_created = True
3141 self._in_use_connections.add(connection)
3143 # Record state transition: IDLE -> USED
3144 # (make_connection already recorded IDLE +1 for new connections)
3145 # This ensures counters stay balanced if connect() fails and release() is called
3146 pool_name = get_pool_name(self)
3147 record_connection_count(
3148 pool_name=pool_name,
3149 connection_state=ConnectionState.IDLE,
3150 counter=-1,
3151 )
3152 record_connection_count(
3153 pool_name=pool_name,
3154 connection_state=ConnectionState.USED,
3155 counter=1,
3156 )
3158 try:
3159 # ensure this connection is connected to Redis
3160 connection.connect()
3161 # connections that the pool provides should be ready to send
3162 # a command. if not, the connection was either returned to the
3163 # pool before all data has been read or the socket has been
3164 # closed. either way, reconnect and verify everything is good.
3165 try:
3166 if (
3167 connection.can_read()
3168 and self.cache is None
3169 and not self.maint_notifications_enabled()
3170 ):
3171 raise ConnectionError("Connection has data")
3172 except (ConnectionError, TimeoutError, OSError):
3173 connection.disconnect()
3174 connection.connect()
3175 if connection.can_read():
3176 raise ConnectionError("Connection not ready")
3177 except BaseException:
3178 # release the connection back to the pool so that we don't
3179 # leak it
3180 self.release(connection)
3181 raise
3183 if is_created:
3184 record_connection_create_time(
3185 connection_pool=self,
3186 duration_seconds=time.monotonic() - start_time_created,
3187 )
3189 return connection
3191 def get_encoder(self) -> Encoder:
3192 "Return an encoder based on encoding settings"
3193 kwargs = self.connection_kwargs
3194 return Encoder(
3195 encoding=kwargs.get("encoding", "utf-8"),
3196 encoding_errors=kwargs.get("encoding_errors", "strict"),
3197 decode_responses=kwargs.get("decode_responses", False),
3198 )
3200 def make_connection(self) -> "ConnectionInterface":
3201 "Create a new connection"
3202 if self._created_connections >= self.max_connections:
3203 raise MaxConnectionsError("Too many connections")
3204 self._created_connections += 1
3206 kwargs = dict(self.connection_kwargs)
3208 # Create the connection first, then record metrics only on success
3209 if self.cache is not None:
3210 connection = CacheProxyConnection(
3211 self.connection_class(**kwargs), self.cache, self._lock
3212 )
3213 else:
3214 connection = self.connection_class(**kwargs)
3216 # Record new connection created (starts as IDLE) - only after successful construction
3217 record_connection_count(
3218 pool_name=get_pool_name(self),
3219 connection_state=ConnectionState.IDLE,
3220 counter=1,
3221 )
3223 return connection
3225 def release(self, connection: "Connection") -> None:
3226 "Releases the connection back to the pool"
3227 self._checkpid()
3228 with self._lock:
3229 try:
3230 self._in_use_connections.remove(connection)
3231 except KeyError:
3232 # Gracefully fail when a connection is returned to this pool
3233 # that the pool doesn't actually own
3234 return
3236 if self.owns_connection(connection):
3237 if connection.should_reconnect():
3238 connection.disconnect()
3239 self._available_connections.append(connection)
3240 self._event_dispatcher.dispatch(
3241 AfterConnectionReleasedEvent(connection)
3242 )
3244 # Record state transition: USED -> IDLE
3245 pool_name = get_pool_name(self)
3246 record_connection_count(
3247 pool_name=pool_name,
3248 connection_state=ConnectionState.USED,
3249 counter=-1,
3250 )
3251 record_connection_count(
3252 pool_name=pool_name,
3253 connection_state=ConnectionState.IDLE,
3254 counter=1,
3255 )
3256 else:
3257 # Pool doesn't own this connection, do not add it back
3258 # to the pool.
3259 # The created connections count should not be changed,
3260 # because the connection was not created by the pool.
3261 # Still need to decrement USED since it was counted in get_connection()
3262 connection.disconnect()
3263 record_connection_count(
3264 pool_name="unknown_pool",
3265 connection_state=ConnectionState.USED,
3266 counter=-1,
3267 )
3268 return
3270 def owns_connection(self, connection: "Connection") -> int:
3271 return connection.pid == self.pid
3273 def disconnect(self, inuse_connections: bool = True) -> None:
3274 """
3275 Disconnects connections in the pool
3277 If ``inuse_connections`` is True, disconnect connections that are
3278 currently in use, potentially by other threads. Otherwise only disconnect
3279 connections that are idle in the pool.
3280 """
3281 self._checkpid()
3282 with self._lock:
3283 if inuse_connections:
3284 connections = chain(
3285 self._available_connections, self._in_use_connections
3286 )
3287 else:
3288 connections = self._available_connections
3290 for connection in connections:
3291 connection.disconnect()
3293 def close(self) -> None:
3294 """Close the pool, disconnecting all connections"""
3295 self.disconnect()
3297 def __enter__(self: _CP) -> _CP:
3298 return self
3300 def __exit__(self, exc_type, exc_value, traceback) -> None:
3301 self.close()
3303 def set_retry(self, retry: Retry) -> None:
3304 self.connection_kwargs.update({"retry": retry})
3305 for conn in self._available_connections:
3306 conn.retry = retry
3307 for conn in self._in_use_connections:
3308 conn.retry = retry
3310 def re_auth_callback(self, token: TokenInterface):
3311 with self._lock:
3312 for conn in self._available_connections:
3313 conn.retry.call_with_retry(
3314 lambda: conn.send_command(
3315 "AUTH", token.try_get("oid"), token.get_value()
3316 ),
3317 lambda error: self._mock(error),
3318 )
3319 conn.retry.call_with_retry(
3320 lambda: conn.read_response(), lambda error: self._mock(error)
3321 )
3322 for conn in self._in_use_connections:
3323 conn.set_re_auth_token(token)
3325 def _get_pool_lock(self):
3326 return self._lock
3328 def _get_free_connections(self):
3329 with self._lock:
3330 return list(self._available_connections)
3332 def _get_in_use_connections(self):
3333 with self._lock:
3334 return set(self._in_use_connections)
3336 def _mock(self, error: RedisError):
3337 """
3338 Dummy functions, needs to be passed as error callback to retry object.
3339 :param error:
3340 :return:
3341 """
3342 pass
3344 def get_connection_count(self) -> List[tuple[int, dict]]:
3345 from redis.observability.attributes import get_pool_name
3347 attributes = AttributeBuilder.build_base_attributes()
3348 attributes[DB_CLIENT_CONNECTION_POOL_NAME] = get_pool_name(self)
3349 free_connections_attributes = attributes.copy()
3350 in_use_connections_attributes = attributes.copy()
3352 free_connections_attributes[DB_CLIENT_CONNECTION_STATE] = (
3353 ConnectionState.IDLE.value
3354 )
3355 in_use_connections_attributes[DB_CLIENT_CONNECTION_STATE] = (
3356 ConnectionState.USED.value
3357 )
3359 return [
3360 (len(self._get_free_connections()), free_connections_attributes),
3361 (len(self._get_in_use_connections()), in_use_connections_attributes),
3362 ]
3365class BlockingConnectionPool(ConnectionPool):
3366 """
3367 Thread-safe blocking connection pool::
3369 >>> from redis.client import Redis
3370 >>> client = Redis(connection_pool=BlockingConnectionPool())
3372 It performs the same function as the default
3373 :py:class:`~redis.ConnectionPool` implementation, in that,
3374 it maintains a pool of reusable connections that can be shared by
3375 multiple redis clients (safely across threads if required).
3377 The difference is that, in the event that a client tries to get a
3378 connection from the pool when all of connections are in use, rather than
3379 raising a :py:class:`~redis.ConnectionError` (as the default
3380 :py:class:`~redis.ConnectionPool` implementation does), it
3381 makes the client wait ("blocks") for a specified number of seconds until
3382 a connection becomes available.
3384 Use ``max_connections`` to increase / decrease the pool size::
3386 >>> pool = BlockingConnectionPool(max_connections=10)
3388 Use ``timeout`` to tell it either how many seconds to wait for a connection
3389 to become available, or to block forever:
3391 >>> # Block forever.
3392 >>> pool = BlockingConnectionPool(timeout=None)
3394 >>> # Raise a ``ConnectionError`` after five seconds if a connection is
3395 >>> # not available.
3396 >>> pool = BlockingConnectionPool(timeout=5)
3397 """
3399 def __init__(
3400 self,
3401 max_connections=50,
3402 timeout=20,
3403 connection_class=Connection,
3404 queue_class=LifoQueue,
3405 **connection_kwargs,
3406 ):
3407 self.queue_class = queue_class
3408 self.timeout = timeout
3409 self._in_maintenance = False
3410 self._locked = False
3411 super().__init__(
3412 connection_class=connection_class,
3413 max_connections=max_connections,
3414 **connection_kwargs,
3415 )
3417 def reset(self):
3418 # Create and fill up a thread safe queue with ``None`` values.
3419 try:
3420 if self._in_maintenance:
3421 self._lock.acquire()
3422 self._locked = True
3424 # Record metrics for connections being removed before clearing
3425 # Note: Access pool.queue directly to avoid deadlock since we may
3426 # already hold self._lock (which is non-reentrant)
3427 if (
3428 hasattr(self, "_connections")
3429 and self._connections
3430 and hasattr(self, "pool")
3431 ):
3432 with self._lock:
3433 connections_in_queue = {conn for conn in self.pool.queue if conn}
3434 idle_count = len(connections_in_queue)
3435 in_use_count = len(self._connections) - idle_count
3436 if idle_count > 0 or in_use_count > 0:
3437 pool_name = get_pool_name(self)
3438 if idle_count > 0:
3439 record_connection_count(
3440 pool_name=pool_name,
3441 connection_state=ConnectionState.IDLE,
3442 counter=-idle_count,
3443 )
3444 if in_use_count > 0:
3445 record_connection_count(
3446 pool_name=pool_name,
3447 connection_state=ConnectionState.USED,
3448 counter=-in_use_count,
3449 )
3451 self.pool = self.queue_class(self.max_connections)
3452 while True:
3453 try:
3454 self.pool.put_nowait(None)
3455 except Full:
3456 break
3458 # Keep a list of actual connection instances so that we can
3459 # disconnect them later.
3460 self._connections = []
3461 finally:
3462 if self._locked:
3463 try:
3464 self._lock.release()
3465 except Exception:
3466 pass
3467 self._locked = False
3469 # this must be the last operation in this method. while reset() is
3470 # called when holding _fork_lock, other threads in this process
3471 # can call _checkpid() which compares self.pid and os.getpid() without
3472 # holding any lock (for performance reasons). keeping this assignment
3473 # as the last operation ensures that those other threads will also
3474 # notice a pid difference and block waiting for the first thread to
3475 # release _fork_lock. when each of these threads eventually acquire
3476 # _fork_lock, they will notice that another thread already called
3477 # reset() and they will immediately release _fork_lock and continue on.
3478 self.pid = os.getpid()
3480 def __del__(self) -> None:
3481 """Clean up connection pool and record metrics when garbage collected."""
3482 try:
3483 # Note: Access pool.queue directly to avoid potential deadlock
3484 # if GC runs while the lock is held by the same thread
3485 if (
3486 hasattr(self, "_connections")
3487 and self._connections
3488 and hasattr(self, "pool")
3489 ):
3490 connections_in_queue = {conn for conn in self.pool.queue if conn}
3491 idle_count = len(connections_in_queue)
3492 in_use_count = len(self._connections) - idle_count
3493 if idle_count > 0 or in_use_count > 0:
3494 pool_name = get_pool_name(self)
3495 if idle_count > 0:
3496 record_connection_count(
3497 pool_name=pool_name,
3498 connection_state=ConnectionState.IDLE,
3499 counter=-idle_count,
3500 )
3501 if in_use_count > 0:
3502 record_connection_count(
3503 pool_name=pool_name,
3504 connection_state=ConnectionState.USED,
3505 counter=-in_use_count,
3506 )
3507 except Exception:
3508 pass
3510 def make_connection(self):
3511 "Make a fresh connection."
3512 try:
3513 if self._in_maintenance:
3514 self._lock.acquire()
3515 self._locked = True
3517 if self.cache is not None:
3518 connection = CacheProxyConnection(
3519 self.connection_class(**self.connection_kwargs),
3520 self.cache,
3521 self._lock,
3522 )
3523 else:
3524 connection = self.connection_class(**self.connection_kwargs)
3525 self._connections.append(connection)
3527 # Record new connection created (starts as IDLE)
3528 record_connection_count(
3529 pool_name=get_pool_name(self),
3530 connection_state=ConnectionState.IDLE,
3531 counter=1,
3532 )
3534 return connection
3535 finally:
3536 if self._locked:
3537 try:
3538 self._lock.release()
3539 except Exception:
3540 pass
3541 self._locked = False
3543 @deprecated_args(
3544 args_to_warn=["*"],
3545 reason="Use get_connection() without args instead",
3546 version="5.3.0",
3547 )
3548 def get_connection(self, command_name=None, *keys, **options):
3549 """
3550 Get a connection, blocking for ``self.timeout`` until a connection
3551 is available from the pool.
3553 If the connection returned is ``None`` then creates a new connection.
3554 Because we use a last-in first-out queue, the existing connections
3555 (having been returned to the pool after the initial ``None`` values
3556 were added) will be returned before ``None`` values. This means we only
3557 create new connections when we need to, i.e.: the actual number of
3558 connections will only increase in response to demand.
3559 """
3560 start_time_acquired = time.monotonic()
3561 # Make sure we haven't changed process.
3562 self._checkpid()
3563 is_created = False
3565 # Try and get a connection from the pool. If one isn't available within
3566 # self.timeout then raise a ``ConnectionError``.
3567 connection = None
3568 try:
3569 if self._in_maintenance:
3570 self._lock.acquire()
3571 self._locked = True
3572 try:
3573 connection = self.pool.get(block=True, timeout=self.timeout)
3574 except Empty:
3575 # Note that this is not caught by the redis client and will be
3576 # raised unless handled by application code. If you want never to
3577 raise ConnectionError("No connection available.")
3579 # If the ``connection`` is actually ``None`` then that's a cue to make
3580 # a new connection to add to the pool.
3581 if connection is None:
3582 # Start timing for observability
3583 start_time_created = time.monotonic()
3584 connection = self.make_connection()
3585 is_created = True
3586 finally:
3587 if self._locked:
3588 try:
3589 self._lock.release()
3590 except Exception:
3591 pass
3592 self._locked = False
3594 # Record state transition: IDLE -> USED
3595 # (make_connection already recorded IDLE +1 for new connections)
3596 # This ensures counters stay balanced if connect() fails and release() is called
3597 pool_name = get_pool_name(self)
3598 record_connection_count(
3599 pool_name=pool_name,
3600 connection_state=ConnectionState.IDLE,
3601 counter=-1,
3602 )
3603 record_connection_count(
3604 pool_name=pool_name,
3605 connection_state=ConnectionState.USED,
3606 counter=1,
3607 )
3609 try:
3610 # ensure this connection is connected to Redis
3611 connection.connect()
3612 # connections that the pool provides should be ready to send
3613 # a command. if not, the connection was either returned to the
3614 # pool before all data has been read or the socket has been
3615 # closed. either way, reconnect and verify everything is good.
3616 try:
3617 if connection.can_read():
3618 raise ConnectionError("Connection has data")
3619 except (ConnectionError, TimeoutError, OSError):
3620 connection.disconnect()
3621 connection.connect()
3622 if connection.can_read():
3623 raise ConnectionError("Connection not ready")
3624 except BaseException:
3625 # release the connection back to the pool so that we don't leak it
3626 self.release(connection)
3627 raise
3629 if is_created:
3630 record_connection_create_time(
3631 connection_pool=self,
3632 duration_seconds=time.monotonic() - start_time_created,
3633 )
3635 record_connection_wait_time(
3636 pool_name=pool_name,
3637 duration_seconds=time.monotonic() - start_time_acquired,
3638 )
3640 return connection
3642 def release(self, connection):
3643 "Releases the connection back to the pool."
3644 # Make sure we haven't changed process.
3645 self._checkpid()
3647 try:
3648 if self._in_maintenance:
3649 self._lock.acquire()
3650 self._locked = True
3651 if not self.owns_connection(connection):
3652 # pool doesn't own this connection. do not add it back
3653 # to the pool. instead add a None value which is a placeholder
3654 # that will cause the pool to recreate the connection if
3655 # its needed.
3656 connection.disconnect()
3657 self.pool.put_nowait(None)
3658 # Still need to decrement USED since it was counted in get_connection()
3659 record_connection_count(
3660 pool_name="unknown_pool",
3661 connection_state=ConnectionState.USED,
3662 counter=-1,
3663 )
3664 return
3665 if connection.should_reconnect():
3666 connection.disconnect()
3667 # Put the connection back into the pool.
3668 pool_name = get_pool_name(self)
3669 try:
3670 self.pool.put_nowait(connection)
3672 # Record state transition: USED -> IDLE
3673 record_connection_count(
3674 pool_name=pool_name,
3675 connection_state=ConnectionState.USED,
3676 counter=-1,
3677 )
3678 record_connection_count(
3679 pool_name=pool_name,
3680 connection_state=ConnectionState.IDLE,
3681 counter=1,
3682 )
3683 except Full:
3684 pass
3685 finally:
3686 if self._locked:
3687 try:
3688 self._lock.release()
3689 except Exception:
3690 pass
3691 self._locked = False
3693 def disconnect(self, inuse_connections: bool = True):
3694 """
3695 Disconnects either all connections in the pool or just the free connections.
3696 """
3697 self._checkpid()
3698 try:
3699 if self._in_maintenance:
3700 self._lock.acquire()
3701 self._locked = True
3703 if inuse_connections:
3704 connections = self._connections
3705 else:
3706 connections = self._get_free_connections()
3708 for connection in connections:
3709 connection.disconnect()
3710 finally:
3711 if self._locked:
3712 try:
3713 self._lock.release()
3714 except Exception:
3715 pass
3716 self._locked = False
3718 def _get_free_connections(self):
3719 with self._lock:
3720 return {conn for conn in self.pool.queue if conn}
3722 def _get_in_use_connections(self):
3723 with self._lock:
3724 # free connections
3725 connections_in_queue = {conn for conn in self.pool.queue if conn}
3726 # in self._connections we keep all created connections
3727 # so the ones that are not in the queue are the in use ones
3728 return {
3729 conn for conn in self._connections if conn not in connections_in_queue
3730 }
3732 def set_in_maintenance(self, in_maintenance: bool):
3733 """
3734 Sets a flag that this Blocking ConnectionPool is in maintenance mode.
3736 This is used to prevent new connections from being created while we are in maintenance mode.
3737 The pool will be in maintenance mode only when we are processing a MOVING notification.
3738 """
3739 self._in_maintenance = in_maintenance