[7.x] Add ignore_status option to bulk helpers

This commit is contained in:
Vladimir Kaspar
2021-04-20 14:35:22 +02: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,
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
+3
View File
@@ -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]]]: ...
+33 -10
View File
@@ -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(
+3
View File
@@ -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]: ...
@@ -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",
@@ -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",