Fixed infinitely recursive health checks.

Author: excitoon

redis/asyncio/connection.py

@@ -282,7 +282,7 @@ def set_parser(self, parser_class: Type[BaseParser]) -> None:
         """
         self._parser = parser_class(socket_read_size=self._socket_read_size)
 
-    async def connect(self):
+    async def connect(self, check_health: bool = True):
         """Connects to the Redis server if not already connected"""
         if self.is_connected:
             return
@@ -302,7 +302,7 @@ async def connect(self):
         try:
             if not self.redis_connect_func:
                 # Use the default on_connect function
-                await self.on_connect()
+                await self.on_connect(check_health=check_health)
             else:
                 # Use the passed function redis_connect_func
                 (
@@ -339,7 +339,7 @@ def _error_message(self, exception: BaseException) -> str:
     def get_protocol(self):
         return self.protocol
 
-    async def on_connect(self) -> None:
+    async def on_connect(self, check_health: bool = True) -> None:
         """Initialize the connection, authenticate and select a database"""
         self._parser.on_connect(self)
         parser = self._parser
@@ -398,7 +398,7 @@ async def on_connect(self) -> None:
                 # update cluster exception classes
                 self._parser.EXCEPTION_CLASSES = parser.EXCEPTION_CLASSES
                 self._parser.on_connect(self)
-            await self.send_command("HELLO", self.protocol)
+            await self.send_command("HELLO", self.protocol, check_health=check_health)
             response = await self.read_response()
             # if response.get(b"proto") != self.protocol and response.get(
             #     "proto"
@@ -407,18 +407,35 @@ async def on_connect(self) -> None:
 
         # if a client_name is given, set it
         if self.client_name:
-            await self.send_command("CLIENT", "SETNAME", self.client_name)
+            await self.send_command(
+                "CLIENT",
+                "SETNAME",
+                self.client_name,
+                check_health=check_health,
+            )
             if str_if_bytes(await self.read_response()) != "OK":
                 raise ConnectionError("Error setting client name")
 
         # set the library name and version, pipeline for lower startup latency
         if self.lib_name:
-            await self.send_command("CLIENT", "SETINFO", "LIB-NAME", self.lib_name)
+            await self.send_command(
+                "CLIENT",
+                "SETINFO",
+                "LIB-NAME",
+                self.lib_name,
+                check_health=check_health,
+            )
         if self.lib_version:
-            await self.send_command("CLIENT", "SETINFO", "LIB-VER", self.lib_version)
+            await self.send_command(
+                "CLIENT",
+                "SETINFO",
+                "LIB-VER",
+                self.lib_version,
+                check_health=check_health,
+            )
         # if a database is specified, switch to it. Also pipeline this
         if self.db:
-            await self.send_command("SELECT", self.db)
+            await self.send_command("SELECT", self.db, check_health=check_health)
 
         # read responses from pipeline
         for _ in (sent for sent in (self.lib_name, self.lib_version) if sent):
@@ -480,8 +497,8 @@ async def send_packed_command(
         self, command: Union[bytes, str, Iterable[bytes]], check_health: bool = True
     ) -> None:
         if not self.is_connected:
-            await self.connect()
-        elif check_health:
+            await self.connect(check_health=False)
+        if check_health:
             await self.check_health()
 
         try:
@@ -1134,7 +1151,7 @@ def make_connection(self):
 
     async def ensure_connection(self, connection: AbstractConnection):
         """Ensure that the connection object is connected and valid"""
-        await connection.connect()
+        await connection.connect(check_health=True)
         # connections that the pool provides should be ready to send
         # a command. if not, the connection was either returned to the
         # pool before all data has been read or the socket has been
@@ -1144,7 +1161,7 @@ async def ensure_connection(self, connection: AbstractConnection):
                 raise ConnectionError("Connection has data") from None
         except (ConnectionError, OSError):
             await connection.disconnect()
-            await connection.connect()
+            await connection.connect(check_health=True)
             if await connection.can_read_destructive():
                 raise ConnectionError("Connection not ready") from None
 

redis/connection.py

@@ -159,11 +159,11 @@ def get_protocol(self):
         pass
 
     @abstractmethod
-    def connect(self):
+    def connect(self, check_health=True):
         pass
 
     @abstractmethod
-    def on_connect(self):
+    def on_connect(self, check_health=True):
         pass
 
     @abstractmethod
@@ -370,13 +370,14 @@ def set_parser(self, parser_class):
         """
         self._parser = parser_class(socket_read_size=self._socket_read_size)
 
-    def connect(self):
+    def connect(self, check_health: bool = True):
         "Connects to the Redis server if not already connected"
         if self._sock:
             return
         try:
             sock = self.retry.call_with_retry(
-                lambda: self._connect(), lambda error: self.disconnect(error)
+                lambda: self._connect(),
+                lambda error: self.disconnect(error),
             )
         except socket.timeout:
             raise TimeoutError("Timeout connecting to server")
@@ -387,7 +388,7 @@ def connect(self):
         try:
             if self.redis_connect_func is None:
                 # Use the default on_connect function
-                self.on_connect()
+                self.on_connect(check_health=check_health)
             else:
                 # Use the passed function redis_connect_func
                 self.redis_connect_func(self)
@@ -416,7 +417,7 @@ def _host_error(self):
     def _error_message(self, exception):
         return format_error_message(self._host_error(), exception)
 
-    def on_connect(self):
+    def on_connect(self, check_health: bool = True):
         "Initialize the connection, authenticate and select a database"
         self._parser.on_connect(self)
         parser = self._parser
@@ -475,7 +476,7 @@ def on_connect(self):
                 # update cluster exception classes
                 self._parser.EXCEPTION_CLASSES = parser.EXCEPTION_CLASSES
                 self._parser.on_connect(self)
-            self.send_command("HELLO", self.protocol)
+            self.send_command("HELLO", self.protocol, check_health=check_health)
             self.handshake_metadata = self.read_response()
             if (
                 self.handshake_metadata.get(b"proto") != self.protocol
@@ -485,28 +486,45 @@ def on_connect(self):
 
         # if a client_name is given, set it
         if self.client_name:
-            self.send_command("CLIENT", "SETNAME", self.client_name)
+            self.send_command(
+                "CLIENT",
+                "SETNAME",
+                self.client_name,
+                check_health=check_health,
+            )
             if str_if_bytes(self.read_response()) != "OK":
                 raise ConnectionError("Error setting client name")
 
         try:
             # set the library name and version
             if self.lib_name:
-                self.send_command("CLIENT", "SETINFO", "LIB-NAME", self.lib_name)
+                self.send_command(
+                    "CLIENT",
+                    "SETINFO",
+                    "LIB-NAME",
+                    self.lib_name,
+                    check_health=check_health,
+                )
                 self.read_response()
         except ResponseError:
             pass
 
         try:
             if self.lib_version:
-                self.send_command("CLIENT", "SETINFO", "LIB-VER", self.lib_version)
+                self.send_command(
+                    "CLIENT",
+                    "SETINFO",
+                    "LIB-VER",
+                    self.lib_version,
+                    check_health=check_health,
+                )
                 self.read_response()
         except ResponseError:
             pass
 
         # if a database is specified, switch to it
         if self.db:
-            self.send_command("SELECT", self.db)
+            self.send_command("SELECT", self.db, check_health=check_health)
             if str_if_bytes(self.read_response()) != "OK":
                 raise ConnectionError("Invalid Database")
 
@@ -548,7 +566,7 @@ def check_health(self):
     def send_packed_command(self, command, check_health=True):
         """Send an already packed command to the Redis server"""
         if not self._sock:
-            self.connect()
+            self.connect(check_health=False)
         # guard against health check recursion
         if check_health:
             self.check_health()
@@ -587,7 +605,7 @@ def can_read(self, timeout=0):
         """Poll the socket to see if there's data that can be read."""
         sock = self._sock
         if not sock:
-            self.connect()
+            self.connect(check_health=True)
 
         host_error = self._host_error()
 
@@ -806,8 +824,8 @@ def deregister_connect_callback(self, callback):
     def set_parser(self, parser_class):
         self._conn.set_parser(parser_class)
 
-    def connect(self):
-        self._conn.connect()
+    def connect(self, check_health: bool = True):
+        self._conn.connect(check_health=check_health)
 
         server_name = self._conn.handshake_metadata.get(b"server", None)
         if server_name is None:
@@ -829,8 +847,8 @@ def connect(self):
                 "To maximize compatibility with all Redis products, client-side caching is supported by Redis 7.4 or later"  # noqa: E501
             )
 
-    def on_connect(self):
-        self._conn.on_connect()
+    def on_connect(self, check_health: bool = True):
+        self._conn.on_connect(check_health=check_health)
 
     def disconnect(self, *args):
         with self._cache_lock:
@@ -1482,7 +1500,7 @@ def get_connection(self, command_name=None, *keys, **options) -> "Connection":
 
         try:
             # ensure this connection is connected to Redis
-            connection.connect()
+            connection.connect(check_health=True)
             # connections that the pool provides should be ready to send
             # a command. if not, the connection was either returned to the
             # pool before all data has been read or the socket has been
@@ -1492,7 +1510,7 @@ def get_connection(self, command_name=None, *keys, **options) -> "Connection":
                     raise ConnectionError("Connection has data")
             except (ConnectionError, OSError):
                 connection.disconnect()
-                connection.connect()
+                connection.connect(check_health=True)
                 if connection.can_read():
                     raise ConnectionError("Connection not ready")
         except BaseException:
@@ -1729,7 +1747,7 @@ def get_connection(self, command_name=None, *keys, **options):
 
         try:
             # ensure this connection is connected to Redis
-            connection.connect()
+            connection.connect(check_health=True)
             # connections that the pool provides should be ready to send
             # a command. if not, the connection was either returned to the
             # pool before all data has been read or the socket has been
@@ -1739,7 +1757,7 @@ def get_connection(self, command_name=None, *keys, **options):
                     raise ConnectionError("Connection has data")
             except (ConnectionError, OSError):
                 connection.disconnect()
-                connection.connect()
+                connection.connect(check_health=True)
                 if connection.can_read():
                     raise ConnectionError("Connection not ready")
         except BaseException:

tests/test_asyncio/test_connection_pool.py

@@ -98,7 +98,7 @@ def __init__(self, **kwargs):
     def repr_pieces(self):
         return [("id", id(self)), ("kwargs", self.kwargs)]
 
-    async def connect(self):
+    async def connect(self, **kwargs):
         pass
 
     async def disconnect(self):

tests/test_connection_pool.py

@@ -27,7 +27,7 @@ def __init__(self, **kwargs):
         self.kwargs = kwargs
         self.pid = os.getpid()
 
-    def connect(self):
+    def connect(self, **kwargs):
         pass
 
     def can_read(self):