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: