diff --git a/elasticsearch/helpers/__init__.py b/elasticsearch/helpers/__init__.py index df4e1b77..ee852dd5 100644 --- a/elasticsearch/helpers/__init__.py +++ b/elasticsearch/helpers/__init__.py @@ -2,6 +2,7 @@ from __future__ import unicode_literals import logging from operator import methodcaller +import time from ..exceptions import ElasticsearchException, TransportError from ..compat import map, string_types, Queue @@ -133,8 +134,10 @@ def _process_bulk_chunk(client, bulk_actions, bulk_data, raise_on_exception=True raise BulkIndexError('%i document(s) failed to index.' % len(errors), errors) def streaming_bulk(client, actions, chunk_size=500, max_chunk_bytes=100 * 1024 * 1024, - raise_on_error=True, expand_action_callback=expand_action, - raise_on_exception=True, **kwargs): + raise_on_error=True, expand_action_callback=expand_action, + raise_on_exception=True, max_retries=0, initial_backoff=2, + max_backoff=600, yield_ok=True, **kwargs): + """ Streaming bulk consumes actions from the iterable passed in and yields results per action. For non-streaming usecases use @@ -154,12 +157,58 @@ def streaming_bulk(client, actions, chunk_size=500, max_chunk_bytes=100 * 1024 * :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 max_retries: maximum number of times a document will be retried when + ``429`` is received, set to 0 (default) for no retires on ``429`` + :arg initial_backoff: number of seconds we should wait before the first + retry. Any subsequent retries will be powers of ``inittial_backoff * + 2**retry_number`` + :arg max_backoff: maximum number of seconds a retry will wait + """ actions = map(expand_action_callback, actions) - for bulk_data, bulk_actions in _chunk_actions(actions, chunk_size, max_chunk_bytes, client.transport.serializer): - for result in _process_bulk_chunk(client, bulk_actions, bulk_data, raise_on_exception, raise_on_error, **kwargs): - yield result + for bulk_data, bulk_actions in _chunk_actions(actions, chunk_size, + max_chunk_bytes, + client.transport.serializer): + + for attempt in range(max_retries + 1): + to_retry, to_retry_data = [], [] + if attempt: + time.sleep(min(max_backoff, initial_backoff * 2**(attempt-1))) + + try: + for data, (ok, info) in zip( + bulk_data, + _process_bulk_chunk(client, bulk_actions, bulk_data, + raise_on_exception=raise_on_exception, + raise_on_error=raise_on_error, **kwargs) + ): + + if not ok: + action, info = info.popitem() + # retry if retries enabled, we get 429, and we are not + # in the last attempt + if max_retries and info['status'] == 429 and (attempt+1) <= max_retries: + to_retry.extend(data) + to_retry_data.append(data) + else: + if action != 'delete': + info['data'] = data[1] + yield ok, {action: info} + elif yield_ok: + yield ok, info + + except TransportError as e: + # suppress 429 errors since we will retry them + if not max_retries or e.status_code != 429: + raise + else: + if not to_retry: + break + # retry only subset of documents that didn't succeed + bulk_actions = to_retry + bulk_data = to_retry_data + def bulk(client, actions, stats_only=False, **kwargs): """ @@ -190,6 +239,8 @@ def bulk(client, actions, stats_only=False, **kwargs): # list of errors to be collected is not stats_only errors = [] + # make streaming_bulk yield successful results so we can count them + kwargs['yield_ok'] = True for ok, item in streaming_bulk(client, actions, **kwargs): # go through request-reponse pairs and detect failures if not ok: diff --git a/elasticsearch/helpers/backoff.py b/elasticsearch/helpers/backoff.py deleted file mode 100644 index 9a12bdee..00000000 --- a/elasticsearch/helpers/backoff.py +++ /dev/null @@ -1,70 +0,0 @@ -import time - -from . import expand_action, _chunk_actions, _process_bulk_chunk -from ..exceptions import TransportError - - -def backoff_bulk(client, actions, chunk_size=500, max_chunk_bytes=100 * 1024 * 1024, - expand_action_callback=expand_action, max_retries=-1, - initial_backoff=2, max_backoff=600, **kwargs): - """ - Bulk helper that implements proper retry strategy when dealing with - rejecttions - when the target cluster is overloaded, returning ``429`` - responses, the bulk actions will be retried after a backoff period to allow - the cluster to recover. - - This function is a generator yielding all documents that failed to index. - - :arg client: instance of :class:`~elasticsearch.Elasticsearch` to use - :arg actions: iterable containing the actions to be executed - :arg chunk_size: number of docs in one chunk sent to es (default: 500) - :arg max_chunk_bytes: the maximum size of the request in bytes (default: 100MB) - :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 max_retries: maximum number of times a document will be retried - :arg initial_backoff: number of seconds we should wait before the first - retry. Any subsequent retries will be powers of ``inittial_backoff * - 2**retry_number`` - :arg max_backoff: maximum number of seconds a retry will wait - """ - actions = map(expand_action_callback, actions) - - for bulk_data, bulk_actions in _chunk_actions(actions, chunk_size, - max_chunk_bytes, - client.transport.serializer): - retry = 0 - while max_retries == -1 or retry <= max_retries: - to_retry, to_retry_data = [], [] - if retry: - time.sleep(min(max_backoff, initial_backoff * 2**(retry-1))) - - try: - for data, (ok, info) in zip(bulk_data, - _process_bulk_chunk(client, - bulk_actions, - bulk_data, - raise_on_exception=True, - raise_on_error=False, - **kwargs)): - - if not ok: - action, info = info.popitem() - if info['status'] == 429 and (retry+1) <= max_retries: - to_retry.extend(data) - to_retry_data.append(data) - else: - info['data'] = data - yield {action: info} - - except TransportError as e: - if e.status_code != 429: - raise - retry += 1 - else: - if not to_retry: - break - bulk_actions = to_retry - bulk_data = to_retry_data - - retry += 1