Use enum instead of function for status retries

Author: pquentin

elasticsearch/_async/helpers.py

@@ -158,12 +158,6 @@ async def azip(
         pass
 
 
-def _retry_for_status(status: int) -> bool:
-    if status == 429:
-        return True
-    return False
-
-
 async def async_streaming_bulk(
     client: AsyncElasticsearch,
     actions: Union[Iterable[_TYPE_BULK_ACTION], AsyncIterable[_TYPE_BULK_ACTION]],
@@ -173,13 +167,13 @@ async def async_streaming_bulk(
     expand_action_callback: Callable[
         [_TYPE_BULK_ACTION], _TYPE_BULK_ACTION_HEADER_AND_BODY
     ] = expand_action,
-    retry_for_status_callback: Callable[[int], bool] = _retry_for_status,
     raise_on_exception: bool = True,
     max_retries: int = 0,
     initial_backoff: float = 2,
     max_backoff: float = 600,
     yield_ok: bool = True,
     ignore_status: Union[int, Collection[int]] = (),
+    retry_on_status: Union[int, Collection[int]] = (429,),
     *args: Any,
     **kwargs: Any,
 ) -> AsyncIterable[Tuple[bool, Dict[str, Any]]]:
@@ -191,7 +185,7 @@ async def async_streaming_bulk(
     entire input is consumed and sent.
 
     If you specify ``max_retries`` it will also retry any documents that were
-    rejected with a ``429`` status code. Use ``retry_for_status_callback`` to
+    rejected with a ``429`` status code. Use ``retry_on_status`` to
     configure which status codes will be retried. To do this it will wait
     (**by calling time.sleep which will block**) for ``initial_backoff`` seconds
     and then, every subsequent rejection for the same chunk, for double the time
@@ -208,12 +202,11 @@ async def async_streaming_bulk(
     :arg expand_action_callback: callback executed on each action passed in,
         should return a tuple containing the action line and the data line
         (`None` if data line should be omitted).
-    :arg retry_for_status_callback: callback executed on each item's status,
-        should return a True if the status require a retry and False if not.
+    :arg retry_on_status: HTTP status code that will trigger a retry.
         (if `None` is specified only status 429 will retry).
     :arg max_retries: maximum number of times a document will be retried when
-        retry_for_status_callback (defaulting to ``429``) is received,
-        set to 0 (default) for no retries on retry_for_status_callback
+        retry_on_status (defaulting to ``429``) is received,
+        set to 0 (default) for no retries
     :arg initial_backoff: number of seconds we should wait before the first
         retry. Any subsequent retries will be powers of ``initial_backoff *
         2**retry_number``
@@ -225,6 +218,9 @@ async def async_streaming_bulk(
     client = client.options()
     client._client_meta = (("h", "bp"),)
 
+    if isinstance(retry_on_status, int):
+        retry_on_status = (retry_on_status,)
+
     async def map_actions() -> AsyncIterable[_TYPE_BULK_ACTION_HEADER_AND_BODY]:
         async for item in aiter(actions):
             yield expand_action_callback(item)
@@ -277,10 +273,10 @@ async def map_actions() -> AsyncIterable[_TYPE_BULK_ACTION_HEADER_AND_BODY]:
                     if not ok:
                         action, info = info.popitem()
                         # retry if retries enabled, we are not in the last attempt,
-                        # and retry_for_status_callback is true (defaulting to 429)
+                        # and status not in retry_on_status (defaulting to 429)
                         if (
                             max_retries
-                            and retry_for_status_callback(info["status"])
+                            and info["status"] in retry_on_status
                             and (attempt + 1) <= max_retries
                         ):
                             # _process_bulk_chunk expects strings so we need to
@@ -293,11 +289,9 @@ async def map_actions() -> AsyncIterable[_TYPE_BULK_ACTION_HEADER_AND_BODY]:
                         yield ok, info
 
             except ApiError as e:
-                # suppress any status which retry_for_status_callback is true (defaulting to 429)
+                # suppress any status in retry_on_status (429 by default)
                 # since we will retry them
-                if attempt == max_retries or not retry_for_status_callback(
-                    e.status_code
-                ):
+                if attempt == max_retries or e.status_code not in retry_on_status:
                     raise
             else:
                 if not to_retry:

elasticsearch/helpers/actions.py

@@ -359,12 +359,6 @@ def _process_bulk_chunk(
         yield from gen
 
 
-def _retry_for_status(status: int) -> bool:
-    if status == 429:
-        return True
-    return False
-
-
 def streaming_bulk(
     client: Elasticsearch,
     actions: Iterable[_TYPE_BULK_ACTION],
@@ -374,13 +368,13 @@ def streaming_bulk(
     expand_action_callback: Callable[
         [_TYPE_BULK_ACTION], _TYPE_BULK_ACTION_HEADER_AND_BODY
     ] = expand_action,
-    retry_for_status_callback: Callable[[int], bool] = _retry_for_status,
     raise_on_exception: bool = True,
     max_retries: int = 0,
     initial_backoff: float = 2,
     max_backoff: float = 600,
     yield_ok: bool = True,
     ignore_status: Union[int, Collection[int]] = (),
+    retry_on_status: Union[int, Collection[int]] = (429,),
     span_name: str = "helpers.streaming_bulk",
     *args: Any,
     **kwargs: Any,
@@ -393,7 +387,7 @@ def streaming_bulk(
     entire input is consumed and sent.
 
     If you specify ``max_retries`` it will also retry any documents that were
-    rejected with a ``429`` status code. Use ``retry_for_status_callback`` to
+    rejected with a ``429`` status code. Use ``retry_on_status`` to
     configure which status codes will be retried. To do this it will wait
     (**by calling time.sleep which will block**) for ``initial_backoff`` seconds
     and then, every subsequent rejection for the same chunk, for double the time
@@ -410,12 +404,11 @@ def streaming_bulk(
     :arg expand_action_callback: callback executed on each action passed in,
         should return a tuple containing the action line and the data line
         (`None` if data line should be omitted).
-    :arg retry_for_status_callback: callback executed on each item's status,
-        should return a True if the status require a retry and False if not.
+    :arg retry_on_status: HTTP status code that will trigger a retry.
         (if `None` is specified only status 429 will retry).
     :arg max_retries: maximum number of times a document will be retried when
-        retry_for_status_callback (defaulting to ``429``) is received,
-        set to 0 (default) for no retries on retry_for_status_callback
+        retry_on_status (defaulting to ``429``) is received,
+        set to 0 (default) for no retries
     :arg initial_backoff: number of seconds we should wait before the first
         retry. Any subsequent retries will be powers of ``initial_backoff *
         2**retry_number``
@@ -427,6 +420,9 @@ def streaming_bulk(
         client = client.options()
         client._client_meta = (("h", "bp"),)
 
+        if isinstance(retry_on_status, int):
+            retry_on_status = (retry_on_status,)
+
         serializer = client.transport.serializers.get_serializer("application/json")
 
         bulk_data: List[
@@ -471,10 +467,10 @@ def streaming_bulk(
                         if not ok:
                             action, info = info.popitem()
                             # retry if retries enabled, we are not in the last attempt,
-                            # and retry_for_status_callback is true (defaulting to 429)
+                            # and status not in retry_on_status (defaulting to 429)
                             if (
                                 max_retries
-                                and retry_for_status_callback(info["status"])
+                                and info["status"] in retry_on_status
                                 and (attempt + 1) <= max_retries
                             ):
                                 # _process_bulk_chunk expects bytes so we need to
@@ -487,11 +483,9 @@ def streaming_bulk(
                             yield ok, info
 
                 except ApiError as e:
-                    # suppress any status which retry_for_status_callback is true (defaulting to 429)
+                    # suppress any status in retry_on_status (429 by default)
                     # since we will retry them
-                    if attempt == max_retries or not retry_for_status_callback(
-                        e.status_code
-                    ):
+                    if attempt == max_retries or e.status_code not in retry_on_status:
                         raise
                 else:
                     if not to_retry:

test_elasticsearch/test_async/test_server/test_helpers.py

@@ -312,11 +312,6 @@ async def test_connection_timeout_is_retried_with_retry_status_callback(
             {"_index": "i", "_id": 42, "f": "v"},
         ]
 
-        def _retry_for_connection_timeout(status):
-            if status == 522:
-                return True
-            return False
-
         results = [
             x
             async for x in helpers.async_streaming_bulk(
@@ -325,7 +320,7 @@ def _retry_for_connection_timeout(status):
                 raise_on_exception=False,
                 raise_on_error=False,
                 chunk_size=1,
-                retry_for_status_callback=_retry_for_connection_timeout,
+                retry_on_status=522,
                 max_retries=1,
                 initial_backoff=0,
             )

test_elasticsearch/test_server/test_helpers.py

@@ -305,11 +305,6 @@ def test_connection_timeout_is_retried_with_retry_status_callback(sync_client):
         {"_index": "i", "_id": 42, "f": "v"},
     ]
 
-    def _retry_for_connection_timeout(status):
-        if status == 522:
-            return True
-        return False
-
     results = list(
         helpers.streaming_bulk(
             failing_client,
@@ -318,7 +313,7 @@ def _retry_for_connection_timeout(status):
             raise_on_exception=False,
             raise_on_error=False,
             chunk_size=1,
-            retry_for_status_callback=_retry_for_connection_timeout,
+            retry_on_status=522,
             max_retries=1,
             initial_backoff=0,
         )