From a8cf8d93ee704d3721735ac7173cecdbe0bba4c9 Mon Sep 17 00:00:00 2001 From: Vladimir Kaspar Date: Tue, 20 Apr 2021 14:35:22 +0200 Subject: [PATCH] [7.x] Add ignore_status option to bulk helpers --- elasticsearch/_async/helpers.py | 18 +++++++- elasticsearch/_async/helpers.pyi | 3 ++ elasticsearch/helpers/actions.py | 43 ++++++++++++++----- elasticsearch/helpers/actions.pyi | 3 ++ .../test_async/test_server/test_helpers.py | 33 ++++++++++++++ .../test_server/test_helpers.py | 33 ++++++++++++++ 6 files changed, 121 insertions(+), 12 deletions(-) diff --git a/elasticsearch/_async/helpers.py b/elasticsearch/_async/helpers.py index b84bfae5..fefb3f91 100644 --- a/elasticsearch/_async/helpers.py +++ b/elasticsearch/_async/helpers.py @@ -59,12 +59,16 @@ 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) @@ -72,12 +76,16 @@ async def _process_bulk_chunk( 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 async_streaming_bulk( bulk_data, raise_on_exception, raise_on_error, + ignore_status, *args, **kwargs ), @@ -218,7 +229,9 @@ async def async_streaming_bulk( 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 diff --git a/elasticsearch/_async/helpers.pyi b/elasticsearch/_async/helpers.pyi index 6028a022..f116e009 100644 --- a/elasticsearch/_async/helpers.pyi +++ b/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,6 +68,7 @@ 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]: ... @@ -74,6 +76,7 @@ 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]]]: ... diff --git a/elasticsearch/helpers/actions.py b/elasticsearch/helpers/actions.py index 228b339d..a8c7712c 100644 --- a/elasticsearch/helpers/actions.py +++ b/elasticsearch/helpers/actions.py @@ -161,7 +161,7 @@ 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 = [] @@ -169,8 +169,10 @@ def _process_bulk_chunk_success(resp, bulk_data, raise_on_error=True): 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,6 +232,9 @@ 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) @@ -236,12 +242,16 @@ def _process_bulk_chunk( 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 parallel_bulk( 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( diff --git a/elasticsearch/helpers/actions.pyi b/elasticsearch/helpers/actions.pyi index 90060c1f..2bc13314 100644 --- a/elasticsearch/helpers/actions.pyi +++ b/elasticsearch/helpers/actions.pyi @@ -61,6 +61,7 @@ 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]: ... @@ -68,6 +69,7 @@ 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]: ... diff --git a/test_elasticsearch/test_async/test_server/test_helpers.py b/test_elasticsearch/test_async/test_server/test_helpers.py index ef98728b..46927b9c 100644 --- a/test_elasticsearch/test_async/test_server/test_helpers.py +++ b/test_elasticsearch/test_async/test_server/test_helpers.py @@ -336,6 +336,39 @@ class TestBulk(object): 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", diff --git a/test_elasticsearch/test_server/test_helpers.py b/test_elasticsearch/test_server/test_helpers.py index 43b0efe2..d9451062 100644 --- a/test_elasticsearch/test_server/test_helpers.py +++ b/test_elasticsearch/test_server/test_helpers.py @@ -303,6 +303,39 @@ class TestBulk(ElasticsearchTestCase): 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",