[7.x] Add ignore_status option to bulk helpers

This commit is contained in:
Vladimir Kaspar
2021-04-20 10:19:38 -05:00
committed by Seth Michael Larson
parent e4cdd4544d
commit a8cf8d93ee
6 changed files with 121 additions and 12 deletions
+16 -2
View File
@@ -59,12 +59,16 @@ async def _process_bulk_chunk(
bulk_data, bulk_data,
raise_on_exception=True, raise_on_exception=True,
raise_on_error=True, raise_on_error=True,
ignore_status=(),
*args, *args,
**kwargs **kwargs
): ):
""" """
Send a bulk request to elasticsearch and process the output. Send a bulk request to elasticsearch and process the output.
""" """
if not isinstance(ignore_status, (list, tuple)):
ignore_status = (ignore_status,)
try: try:
# send the actual request # send the actual request
resp = await client.bulk("\n".join(bulk_actions) + "\n", *args, **kwargs) 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( gen = _process_bulk_chunk_error(
error=e, error=e,
bulk_data=bulk_data, bulk_data=bulk_data,
ignore_status=ignore_status,
raise_on_exception=raise_on_exception, raise_on_exception=raise_on_exception,
raise_on_error=raise_on_error, raise_on_error=raise_on_error,
) )
else: else:
gen = _process_bulk_chunk_success( 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: for item in gen:
yield item yield item
@@ -121,6 +129,7 @@ async def async_streaming_bulk(
initial_backoff=2, initial_backoff=2,
max_backoff=600, max_backoff=600,
yield_ok=True, yield_ok=True,
ignore_status=(),
*args, *args,
**kwargs **kwargs
): ):
@@ -156,6 +165,7 @@ async def async_streaming_bulk(
2**retry_number`` 2**retry_number``
:arg max_backoff: maximum number of seconds a retry will wait :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 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(): async def map_actions():
@@ -182,6 +192,7 @@ async def async_streaming_bulk(
bulk_data, bulk_data,
raise_on_exception, raise_on_exception,
raise_on_error, raise_on_error,
ignore_status,
*args, *args,
**kwargs **kwargs
), ),
@@ -218,7 +229,9 @@ async def async_streaming_bulk(
bulk_actions, bulk_data = to_retry, to_retry_data 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 Helper for the :meth:`~elasticsearch.AsyncElasticsearch.bulk` api that provides
a more human friendly interface - it consumes an iterator of actions and 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 actions: iterator containing the actions
:arg stats_only: if `True` only report number of successful/failed :arg stats_only: if `True` only report number of successful/failed
operations instead of just number of successful and a list of error responses 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 Any additional keyword arguments will be passed to
:func:`~elasticsearch.helpers.async_streaming_bulk` which is used to execute :func:`~elasticsearch.helpers.async_streaming_bulk` which is used to execute
+3
View File
@@ -48,6 +48,7 @@ def _process_bulk_chunk(
bulk_data: Any, bulk_data: Any,
raise_on_exception: bool = ..., raise_on_exception: bool = ...,
raise_on_error: bool = ..., raise_on_error: bool = ...,
ignore_status: Optional[Union[int, Collection[int]]] = ...,
*args: Any, *args: Any,
**kwargs: Any **kwargs: Any
) -> AsyncGenerator[Tuple[bool, Any], None]: ... ) -> AsyncGenerator[Tuple[bool, Any], None]: ...
@@ -67,6 +68,7 @@ def async_streaming_bulk(
initial_backoff: Union[float, int] = ..., initial_backoff: Union[float, int] = ...,
max_backoff: Union[float, int] = ..., max_backoff: Union[float, int] = ...,
yield_ok: bool = ..., yield_ok: bool = ...,
ignore_status: Optional[Union[int, Collection[int]]] = ...,
*args: Any, *args: Any,
**kwargs: Any **kwargs: Any
) -> AsyncGenerator[Tuple[bool, Any], None]: ... ) -> AsyncGenerator[Tuple[bool, Any], None]: ...
@@ -74,6 +76,7 @@ async def async_bulk(
client: AsyncElasticsearch, client: AsyncElasticsearch,
actions: Union[Iterable[Any], AsyncIterable[Any]], actions: Union[Iterable[Any], AsyncIterable[Any]],
stats_only: bool = ..., stats_only: bool = ...,
ignore_status: Optional[Union[int, Collection[int]]] = ...,
*args: Any, *args: Any,
**kwargs: Any **kwargs: Any
) -> Tuple[int, Union[int, List[Any]]]: ... ) -> Tuple[int, Union[int, List[Any]]]: ...
+33 -10
View File
@@ -161,7 +161,7 @@ def _chunk_actions(actions, chunk_size, max_chunk_bytes, serializer):
yield ret 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 # if raise on error is set, we need to collect errors per chunk before raising them
errors = [] errors = []
@@ -169,8 +169,10 @@ def _process_bulk_chunk_success(resp, bulk_data, raise_on_error=True):
for data, (op_type, item) in zip( for data, (op_type, item) in zip(
bulk_data, map(methodcaller("popitem"), resp["items"]) bulk_data, map(methodcaller("popitem"), resp["items"])
): ):
ok = 200 <= item.get("status", 500) < 300 status_code = item.get("status", 500)
if not ok and raise_on_error:
ok = 200 <= status_code < 300
if not ok and raise_on_error and status_code not in ignore_status:
# include original document source # include original document source
if len(data) > 1: if len(data) > 1:
item["data"] = 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( 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 # default behavior - just propagate exception
if raise_on_exception: if raise_on_exception and error.status_code not in ignore_status:
raise error raise error
# if we are not propagating, mark all actions in current chunk as failed # 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}) exc_errors.append({op_type: info})
# emulate standard behavior for failed actions # emulate standard behavior for failed actions
if raise_on_error: if raise_on_error and error.status_code not in ignore_status:
raise BulkIndexError( raise BulkIndexError(
"%i document(s) failed to index." % len(exc_errors), exc_errors "%i document(s) failed to index." % len(exc_errors), exc_errors
) )
@@ -221,6 +223,7 @@ def _process_bulk_chunk(
bulk_data, bulk_data,
raise_on_exception=True, raise_on_exception=True,
raise_on_error=True, raise_on_error=True,
ignore_status=(),
*args, *args,
**kwargs **kwargs
): ):
@@ -229,6 +232,9 @@ def _process_bulk_chunk(
""" """
kwargs = _add_helper_meta_to_kwargs(kwargs, "bp") kwargs = _add_helper_meta_to_kwargs(kwargs, "bp")
if not isinstance(ignore_status, (list, tuple)):
ignore_status = (ignore_status,)
try: try:
# send the actual request # send the actual request
resp = client.bulk("\n".join(bulk_actions) + "\n", *args, **kwargs) resp = client.bulk("\n".join(bulk_actions) + "\n", *args, **kwargs)
@@ -236,12 +242,16 @@ def _process_bulk_chunk(
gen = _process_bulk_chunk_error( gen = _process_bulk_chunk_error(
error=e, error=e,
bulk_data=bulk_data, bulk_data=bulk_data,
ignore_status=ignore_status,
raise_on_exception=raise_on_exception, raise_on_exception=raise_on_exception,
raise_on_error=raise_on_error, raise_on_error=raise_on_error,
) )
else: else:
gen = _process_bulk_chunk_success( 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: for item in gen:
yield item yield item
@@ -266,6 +276,7 @@ def streaming_bulk(
initial_backoff=2, initial_backoff=2,
max_backoff=600, max_backoff=600,
yield_ok=True, yield_ok=True,
ignore_status=(),
*args, *args,
**kwargs **kwargs
): ):
@@ -301,6 +312,7 @@ def streaming_bulk(
2**retry_number`` 2**retry_number``
:arg max_backoff: maximum number of seconds a retry will wait :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 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) actions = map(expand_action_callback, actions)
@@ -322,6 +334,7 @@ def streaming_bulk(
bulk_data, bulk_data,
raise_on_exception, raise_on_exception,
raise_on_error, raise_on_error,
ignore_status,
*args, *args,
**kwargs **kwargs
), ),
@@ -358,7 +371,7 @@ def streaming_bulk(
bulk_actions, bulk_data = to_retry, to_retry_data 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 Helper for the :meth:`~elasticsearch.Elasticsearch.bulk` api that provides
a more human friendly interface - it consumes an iterator of actions and 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 actions: iterator containing the actions
:arg stats_only: if `True` only report number of successful/failed :arg stats_only: if `True` only report number of successful/failed
operations instead of just number of successful and a list of error responses 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 Any additional keyword arguments will be passed to
:func:`~elasticsearch.helpers.streaming_bulk` which is used to execute :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 # make streaming_bulk yield successful results so we can count them
kwargs["yield_ok"] = True 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 # go through request-response pairs and detect failures
if not ok: if not ok:
if not stats_only: if not stats_only:
@@ -413,6 +429,7 @@ def parallel_bulk(
max_chunk_bytes=100 * 1024 * 1024, max_chunk_bytes=100 * 1024 * 1024,
queue_size=4, queue_size=4,
expand_action_callback=expand_action, expand_action_callback=expand_action,
ignore_status=(),
*args, *args,
**kwargs **kwargs
): ):
@@ -433,6 +450,7 @@ def parallel_bulk(
(`None` if data line should be omitted). (`None` if data line should be omitted).
:arg queue_size: size of the task queue between the main thread (producing :arg queue_size: size of the task queue between the main thread (producing
chunks to send) and the processing threads. 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 # Avoid importing multiprocessing unless parallel_bulk is used
# to avoid exceptions on restricted environments like App Engine # to avoid exceptions on restricted environments like App Engine
@@ -454,7 +472,12 @@ def parallel_bulk(
for result in pool.imap( for result in pool.imap(
lambda bulk_chunk: list( lambda bulk_chunk: list(
_process_bulk_chunk( _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( _chunk_actions(
+3
View File
@@ -61,6 +61,7 @@ def streaming_bulk(
initial_backoff: Union[float, int] = ..., initial_backoff: Union[float, int] = ...,
max_backoff: Union[float, int] = ..., max_backoff: Union[float, int] = ...,
yield_ok: bool = ..., yield_ok: bool = ...,
ignore_status: Optional[Union[int, Collection[int]]] = ...,
*args: Any, *args: Any,
**kwargs: Any **kwargs: Any
) -> Generator[Tuple[bool, Any], None, None]: ... ) -> Generator[Tuple[bool, Any], None, None]: ...
@@ -68,6 +69,7 @@ def bulk(
client: Elasticsearch, client: Elasticsearch,
actions: Iterable[Any], actions: Iterable[Any],
stats_only: bool = ..., stats_only: bool = ...,
ignore_status: Optional[Union[int, Collection[int]]] = ...,
*args: Any, *args: Any,
**kwargs: Any **kwargs: Any
) -> Tuple[int, Union[int, List[Any]]]: ... ) -> Tuple[int, Union[int, List[Any]]]: ...
@@ -79,6 +81,7 @@ def parallel_bulk(
max_chunk_bytes: int = ..., max_chunk_bytes: int = ...,
queue_size: int = ..., queue_size: int = ...,
expand_action_callback: Callable[[Any], Tuple[Dict[str, Any], Optional[Any]]] = ..., expand_action_callback: Callable[[Any], Tuple[Dict[str, Any], Optional[Any]]] = ...,
ignore_status: Optional[Union[int, Collection[int]]] = ...,
*args: Any, *args: Any,
**kwargs: Any **kwargs: Any
) -> Generator[Tuple[bool, Any], None, None]: ... ) -> Generator[Tuple[bool, Any], None, None]: ...
@@ -336,6 +336,39 @@ class TestBulk(object):
with pytest.raises(helpers.BulkIndexError): with pytest.raises(helpers.BulkIndexError):
await helpers.async_bulk(async_client, [{"a": 42}, {"a": "c"}], index="i") 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): async def test_errors_are_collected_properly(self, async_client):
await async_client.indices.create( await async_client.indices.create(
"i", "i",
@@ -303,6 +303,39 @@ class TestBulk(ElasticsearchTestCase):
index="i", 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): def test_errors_are_collected_properly(self):
self.client.indices.create( self.client.indices.create(
"i", "i",