[7.x] Add ignore_status option to bulk helpers

Author: Sparkycz

elasticsearch/_async/helpers.py

@@ -59,25 +59,33 @@ async def _process_bulk_chunk(
     bulk_data,
     raise_on_exception=True,
     raise_on_error=True,
+    ignore_status=(),
     *args,
     **kwargs
 ):
     """
     Send a bulk request to elasticsearch and process the output.
     """
+    if not isinstance(ignore_status, (list, tuple)):
+        ignore_status = (ignore_status,)
+
     try:
         # send the actual request
         resp = await client.bulk("\n".join(bulk_actions) + "\n", *args, **kwargs)
     except TransportError as e:
         gen = _process_bulk_chunk_error(
             error=e,
             bulk_data=bulk_data,
+            ignore_status=ignore_status,
             raise_on_exception=raise_on_exception,
             raise_on_error=raise_on_error,
         )
     else:
         gen = _process_bulk_chunk_success(
-            resp=resp, bulk_data=bulk_data, raise_on_error=raise_on_error
+            resp=resp,
+            bulk_data=bulk_data,
+            ignore_status=ignore_status,
+            raise_on_error=raise_on_error,
         )
     for item in gen:
         yield item
@@ -121,6 +129,7 @@ async def async_streaming_bulk(
     initial_backoff=2,
     max_backoff=600,
     yield_ok=True,
+    ignore_status=(),
     *args,
     **kwargs
 ):
@@ -156,6 +165,7 @@ async def async_streaming_bulk(
         2**retry_number``
     :arg max_backoff: maximum number of seconds a retry will wait
     :arg yield_ok: if set to False will skip successful documents in the output
+    :arg ignore_status: list of HTTP status code that you want to ignore
     """
 
     async def map_actions():
@@ -182,6 +192,7 @@ async def map_actions():
                         bulk_data,
                         raise_on_exception,
                         raise_on_error,
+                        ignore_status,
                         *args,
                         **kwargs
                     ),
@@ -218,7 +229,9 @@ async def map_actions():
                 bulk_actions, bulk_data = to_retry, to_retry_data
 
 
-async def async_bulk(client, actions, stats_only=False, *args, **kwargs):
+async def async_bulk(
+    client, actions, stats_only=False, ignore_status=(), *args, **kwargs
+):
     """
     Helper for the :meth:`~elasticsearch.AsyncElasticsearch.bulk` api that provides
     a more human friendly interface - it consumes an iterator of actions and
@@ -240,6 +253,7 @@ async def async_bulk(client, actions, stats_only=False, *args, **kwargs):
     :arg actions: iterator containing the actions
     :arg stats_only: if `True` only report number of successful/failed
         operations instead of just number of successful and a list of error responses
+    :arg ignore_status: list of HTTP status code that you want to ignore
 
     Any additional keyword arguments will be passed to
     :func:`~elasticsearch.helpers.async_streaming_bulk` which is used to execute

elasticsearch/_async/helpers.pyi

@@ -48,6 +48,7 @@ def _process_bulk_chunk(
     bulk_data: Any,
     raise_on_exception: bool = ...,
     raise_on_error: bool = ...,
+    ignore_status: Optional[Union[int, Collection[int]]] = ...,
     *args: Any,
     **kwargs: Any
 ) -> AsyncGenerator[Tuple[bool, Any], None]: ...
@@ -67,13 +68,15 @@ def async_streaming_bulk(
     initial_backoff: Union[float, int] = ...,
     max_backoff: Union[float, int] = ...,
     yield_ok: bool = ...,
+    ignore_status: Optional[Union[int, Collection[int]]] = ...,
     *args: Any,
     **kwargs: Any
 ) -> AsyncGenerator[Tuple[bool, Any], None]: ...
 async def async_bulk(
     client: AsyncElasticsearch,
     actions: Union[Iterable[Any], AsyncIterable[Any]],
     stats_only: bool = ...,
+    ignore_status: Optional[Union[int, Collection[int]]] = ...,
     *args: Any,
     **kwargs: Any
 ) -> Tuple[int, Union[int, List[Any]]]: ...

elasticsearch/helpers/actions.py

@@ -161,16 +161,18 @@ def _chunk_actions(actions, chunk_size, max_chunk_bytes, serializer):
         yield ret
 
 
-def _process_bulk_chunk_success(resp, bulk_data, raise_on_error=True):
+def _process_bulk_chunk_success(resp, bulk_data, ignore_status, raise_on_error=True):
     # if raise on error is set, we need to collect errors per chunk before raising them
     errors = []
 
     # go through request-response pairs and detect failures
     for data, (op_type, item) in zip(
         bulk_data, map(methodcaller("popitem"), resp["items"])
     ):
-        ok = 200 <= item.get("status", 500) < 300
-        if not ok and raise_on_error:
+        status_code = item.get("status", 500)
+
+        ok = 200 <= status_code < 300
+        if not ok and raise_on_error and status_code not in ignore_status:
             # include original document source
             if len(data) > 1:
                 item["data"] = data[1]
@@ -186,10 +188,10 @@ def _process_bulk_chunk_success(resp, bulk_data, raise_on_error=True):
 
 
 def _process_bulk_chunk_error(
-    error, bulk_data, raise_on_exception=True, raise_on_error=True
+    error, bulk_data, ignore_status, raise_on_exception=True, raise_on_error=True
 ):
     # default behavior - just propagate exception
-    if raise_on_exception:
+    if raise_on_exception and error.status_code not in ignore_status:
         raise error
 
     # if we are not propagating, mark all actions in current chunk as failed
@@ -206,7 +208,7 @@ def _process_bulk_chunk_error(
         exc_errors.append({op_type: info})
 
     # emulate standard behavior for failed actions
-    if raise_on_error:
+    if raise_on_error and error.status_code not in ignore_status:
         raise BulkIndexError(
             "%i document(s) failed to index." % len(exc_errors), exc_errors
         )
@@ -221,6 +223,7 @@ def _process_bulk_chunk(
     bulk_data,
     raise_on_exception=True,
     raise_on_error=True,
+    ignore_status=(),
     *args,
     **kwargs
 ):
@@ -229,19 +232,26 @@ def _process_bulk_chunk(
     """
     kwargs = _add_helper_meta_to_kwargs(kwargs, "bp")
 
+    if not isinstance(ignore_status, (list, tuple)):
+        ignore_status = (ignore_status,)
+
     try:
         # send the actual request
         resp = client.bulk("\n".join(bulk_actions) + "\n", *args, **kwargs)
     except TransportError as e:
         gen = _process_bulk_chunk_error(
             error=e,
             bulk_data=bulk_data,
+            ignore_status=ignore_status,
             raise_on_exception=raise_on_exception,
             raise_on_error=raise_on_error,
         )
     else:
         gen = _process_bulk_chunk_success(
-            resp=resp, bulk_data=bulk_data, raise_on_error=raise_on_error
+            resp=resp,
+            bulk_data=bulk_data,
+            ignore_status=ignore_status,
+            raise_on_error=raise_on_error,
         )
     for item in gen:
         yield item
@@ -266,6 +276,7 @@ def streaming_bulk(
     initial_backoff=2,
     max_backoff=600,
     yield_ok=True,
+    ignore_status=(),
     *args,
     **kwargs
 ):
@@ -301,6 +312,7 @@ def streaming_bulk(
         2**retry_number``
     :arg max_backoff: maximum number of seconds a retry will wait
     :arg yield_ok: if set to False will skip successful documents in the output
+    :arg ignore_status: list of HTTP status code that you want to ignore
     """
     actions = map(expand_action_callback, actions)
 
@@ -322,6 +334,7 @@ def streaming_bulk(
                         bulk_data,
                         raise_on_exception,
                         raise_on_error,
+                        ignore_status,
                         *args,
                         **kwargs
                     ),
@@ -358,7 +371,7 @@ def streaming_bulk(
                 bulk_actions, bulk_data = to_retry, to_retry_data
 
 
-def bulk(client, actions, stats_only=False, *args, **kwargs):
+def bulk(client, actions, stats_only=False, ignore_status=(), *args, **kwargs):
     """
     Helper for the :meth:`~elasticsearch.Elasticsearch.bulk` api that provides
     a more human friendly interface - it consumes an iterator of actions and
@@ -380,6 +393,7 @@ def bulk(client, actions, stats_only=False, *args, **kwargs):
     :arg actions: iterator containing the actions
     :arg stats_only: if `True` only report number of successful/failed
         operations instead of just number of successful and a list of error responses
+    :arg ignore_status: list of HTTP status code that you want to ignore
 
     Any additional keyword arguments will be passed to
     :func:`~elasticsearch.helpers.streaming_bulk` which is used to execute
@@ -393,7 +407,9 @@ def bulk(client, actions, stats_only=False, *args, **kwargs):
 
     # make streaming_bulk yield successful results so we can count them
     kwargs["yield_ok"] = True
-    for ok, item in streaming_bulk(client, actions, *args, **kwargs):
+    for ok, item in streaming_bulk(
+        client, actions, ignore_status=ignore_status, *args, **kwargs
+    ):
         # go through request-response pairs and detect failures
         if not ok:
             if not stats_only:
@@ -413,6 +429,7 @@ def parallel_bulk(
     max_chunk_bytes=100 * 1024 * 1024,
     queue_size=4,
     expand_action_callback=expand_action,
+    ignore_status=(),
     *args,
     **kwargs
 ):
@@ -433,6 +450,7 @@ def parallel_bulk(
         (`None` if data line should be omitted).
     :arg queue_size: size of the task queue between the main thread (producing
         chunks to send) and the processing threads.
+    :arg ignore_status: list of HTTP status code that you want to ignore
     """
     # Avoid importing multiprocessing unless parallel_bulk is used
     # to avoid exceptions on restricted environments like App Engine
@@ -454,7 +472,12 @@ def _setup_queues(self):
         for result in pool.imap(
             lambda bulk_chunk: list(
                 _process_bulk_chunk(
-                    client, bulk_chunk[1], bulk_chunk[0], *args, **kwargs
+                    client,
+                    bulk_chunk[1],
+                    bulk_chunk[0],
+                    ignore_status=ignore_status,
+                    *args,
+                    **kwargs
                 )
             ),
             _chunk_actions(

elasticsearch/helpers/actions.pyi

@@ -61,13 +61,15 @@ def streaming_bulk(
     initial_backoff: Union[float, int] = ...,
     max_backoff: Union[float, int] = ...,
     yield_ok: bool = ...,
+    ignore_status: Optional[Union[int, Collection[int]]] = ...,
     *args: Any,
     **kwargs: Any
 ) -> Generator[Tuple[bool, Any], None, None]: ...
 def bulk(
     client: Elasticsearch,
     actions: Iterable[Any],
     stats_only: bool = ...,
+    ignore_status: Optional[Union[int, Collection[int]]] = ...,
     *args: Any,
     **kwargs: Any
 ) -> Tuple[int, Union[int, List[Any]]]: ...
@@ -79,6 +81,7 @@ def parallel_bulk(
     max_chunk_bytes: int = ...,
     queue_size: int = ...,
     expand_action_callback: Callable[[Any], Tuple[Dict[str, Any], Optional[Any]]] = ...,
+    ignore_status: Optional[Union[int, Collection[int]]] = ...,
     *args: Any,
     **kwargs: Any
 ) -> Generator[Tuple[bool, Any], None, None]: ...

test_elasticsearch/test_async/test_server/test_helpers.py

@@ -336,6 +336,39 @@ async def test_error_is_raised(self, async_client):
         with pytest.raises(helpers.BulkIndexError):
             await helpers.async_bulk(async_client, [{"a": 42}, {"a": "c"}], index="i")
 
+    async def test_ignore_error_if_raised(self, async_client):
+        # ignore the status code 400 in tuple
+        await helpers.async_bulk(
+            async_client, [{"a": 42}, {"a": "c"}], index="i", ignore_status=(400,)
+        )
+
+        # ignore the status code 400 in list
+        await helpers.async_bulk(
+            async_client,
+            [{"a": 42}, {"a": "c"}],
+            index="i",
+            ignore_status=[
+                400,
+            ],
+        )
+
+        # ignore the status code 400
+        await helpers.async_bulk(
+            async_client, [{"a": 42}, {"a": "c"}], index="i", ignore_status=400
+        )
+
+        # ignore only the status code in the `ignore_status` argument
+        with pytest.raises(helpers.BulkIndexError):
+            await helpers.async_bulk(
+                async_client, [{"a": 42}, {"a": "c"}], index="i", ignore_status=(444,)
+            )
+
+        # ignore transport error exception
+        failing_client = FailingBulkClient(async_client)
+        await helpers.async_bulk(
+            failing_client, [{"a": 42}], index="i", ignore_status=(599,)
+        )
+
     async def test_errors_are_collected_properly(self, async_client):
         await async_client.indices.create(
             "i",

test_elasticsearch/test_server/test_helpers.py

@@ -303,6 +303,39 @@ def test_error_is_raised(self):
             index="i",
         )
 
+    def test_ignore_error_if_raised(self):
+        # ignore the status code 400 in tuple
+        helpers.bulk(
+            self.client, [{"a": 42}, {"a": "c"}], index="i", ignore_status=(400,)
+        )
+
+        # ignore the status code 400 in list
+        helpers.bulk(
+            self.client,
+            [{"a": 42}, {"a": "c"}],
+            index="i",
+            ignore_status=[
+                400,
+            ],
+        )
+
+        # ignore the status code 400
+        helpers.bulk(self.client, [{"a": 42}, {"a": "c"}], index="i", ignore_status=400)
+
+        # ignore only the status code in the `ignore_status` argument
+        self.assertRaises(
+            helpers.BulkIndexError,
+            helpers.bulk,
+            self.client,
+            [{"a": 42}, {"a": "c"}],
+            index="i",
+            ignore_status=(444,),
+        )
+
+        # ignore transport error exception
+        failing_client = FailingBulkClient(self.client)
+        helpers.bulk(failing_client, [{"a": 42}], index="i", ignore_status=(599,))
+
     def test_errors_are_collected_properly(self):
         self.client.indices.create(
             "i",