Make retries config consistent in sync and async

Author: elemoine

redis/asyncio/client.py

@@ -158,6 +158,7 @@ def __init__(
         encoding_errors: str = "strict",
         decode_responses: bool = False,
         retry_on_timeout: bool = False,
+        retry_on_error: Optional[list] = None,
         ssl: bool = False,
         ssl_keyfile: Optional[str] = None,
         ssl_certfile: Optional[str] = None,
@@ -176,8 +177,10 @@ def __init__(
     ):
         """
         Initialize a new Redis client.
-        To specify a retry policy, first set `retry_on_timeout` to `True`
-        then set `retry` to a valid `Retry` object
+        To specify a retry policy for specific errors, first set
+        `retry_on_error` to a list of the error/s to retry on, then set
+        `retry` to a valid `Retry` object.
+        To retry on TimeoutError, `retry_on_timeout` can also be set to `True`.
         """
         kwargs: Dict[str, Any]
         # auto_close_connection_pool only has an effect if connection_pool is
@@ -188,6 +191,10 @@ def __init__(
             auto_close_connection_pool if connection_pool is None else False
         )
         if not connection_pool:
+            if not retry_on_error:
+                retry_on_error = []
+            if retry_on_timeout is True:
+                retry_on_error.append(TimeoutError)
             kwargs = {
                 "db": db,
                 "username": username,
@@ -197,6 +204,7 @@ def __init__(
                 "encoding_errors": encoding_errors,
                 "decode_responses": decode_responses,
                 "retry_on_timeout": retry_on_timeout,
+                "retry_on_error": retry_on_error,
                 "retry": copy.deepcopy(retry),
                 "max_connections": max_connections,
                 "health_check_interval": health_check_interval,

redis/asyncio/connection.py

@@ -578,6 +578,7 @@ class Connection:
         "socket_type",
         "redis_connect_func",
         "retry_on_timeout",
+        "retry_on_error",
         "health_check_interval",
         "next_health_check",
         "last_active_at",
@@ -606,6 +607,7 @@ def __init__(
         socket_keepalive_options: Optional[Mapping[int, Union[int, bytes]]] = None,
         socket_type: int = 0,
         retry_on_timeout: bool = False,
+        retry_on_error: Union[list, _Sentinel] = SENTINEL,
         encoding: str = "utf-8",
         encoding_errors: str = "strict",
         decode_responses: bool = False,
@@ -631,12 +633,19 @@ def __init__(
         self.socket_keepalive_options = socket_keepalive_options or {}
         self.socket_type = socket_type
         self.retry_on_timeout = retry_on_timeout
+        if retry_on_error is SENTINEL:
+            retry_on_error = []
         if retry_on_timeout:
+            retry_on_error.append(TimeoutError)
+        self.retry_on_error = retry_on_error
+        if retry_on_error:
             if not retry:
                 self.retry = Retry(NoBackoff(), 1)
             else:
                 # deep-copy the Retry object as it is mutable
                 self.retry = copy.deepcopy(retry)
+            # Update the retry's supported errors with the specified errors
+            self.retry.update_supported_errors(retry_on_error)
         else:
             self.retry = Retry(NoBackoff(), 0)
         self.health_check_interval = health_check_interval
@@ -1169,6 +1178,7 @@ def __init__(
         encoding_errors: str = "strict",
         decode_responses: bool = False,
         retry_on_timeout: bool = False,
+        retry_on_error: Union[list, _Sentinel] = SENTINEL,
         parser_class: Type[BaseParser] = DefaultParser,
         socket_read_size: int = 65536,
         health_check_interval: float = 0.0,
@@ -1190,12 +1200,18 @@ def __init__(
         self.socket_timeout = socket_timeout
         self.socket_connect_timeout = socket_connect_timeout or socket_timeout or None
         self.retry_on_timeout = retry_on_timeout
+        if retry_on_error is SENTINEL:
+            retry_on_error = []
         if retry_on_timeout:
+            retry_on_error.append(TimeoutError)
+        if retry_on_error:
             if retry is None:
                 self.retry = Retry(NoBackoff(), 1)
             else:
                 # deep-copy the Retry object as it is mutable
                 self.retry = copy.deepcopy(retry)
+            # Update the retry's supported errors with the specified errors
+            self.retry.update_supported_errors(retry_on_error)
         else:
             self.retry = Retry(NoBackoff(), 0)
         self.health_check_interval = health_check_interval

redis/asyncio/retry.py

@@ -35,6 +35,14 @@ def __init__(
         self._retries = retries
         self._supported_errors = supported_errors
 
+    def update_supported_errors(self, specified_errors: list):
+        """
+        Updates the supported errors with the specified error types
+        """
+        self._supported_errors = tuple(
+            set(self._supported_errors + tuple(specified_errors))
+        )
+
     async def call_with_retry(
         self, do: Callable[[], Awaitable[T]], fail: Callable[[RedisError], Any]
     ) -> T: