diff --git a/elasticsearch/helpers/__init__.py b/elasticsearch/helpers/__init__.py index 41639bbe..a59437da 100644 --- a/elasticsearch/helpers/__init__.py +++ b/elasticsearch/helpers/__init__.py @@ -92,7 +92,7 @@ def _process_bulk_chunk(client, bulk_actions, bulk_data, raise_on_exception=True try: # send the actual request - resp = client.bulk('\n'.join(map(client.transport.serializer.dumps, bulk_actions)) + '\n', **kwargs) + resp = client.bulk('\n'.join(bulk_actions) + '\n', **kwargs) except TransportError as e: # default behavior - just propagate exception if raise_on_exception: @@ -104,7 +104,7 @@ def _process_bulk_chunk(client, bulk_actions, bulk_data, raise_on_exception=True for data in bulk_data: # collect all the information about failed actions - op_type, action = data[0].popitem() + op_type, action = data[0].copy().popitem() info = {"error": err_message, "status": e.status_code, "exception": e} if op_type != 'delete': info['data'] = data[1] @@ -199,7 +199,9 @@ def streaming_bulk(client, actions, chunk_size=500, max_chunk_bytes=100 * 1024 * if max_retries \ and info['status'] == 429 \ and (attempt+1) <= max_retries: - to_retry.extend(data) + # _process_bulk_chunk expects strings so we need to + # re-serialize the data + to_retry.extend(map(client.transport.serializer.dumps, data)) to_retry_data.append(data) else: yield ok, {action: info} diff --git a/test_elasticsearch/test_server/test_helpers.py b/test_elasticsearch/test_server/test_helpers.py index df6f8099..be80e721 100644 --- a/test_elasticsearch/test_server/test_helpers.py +++ b/test_elasticsearch/test_server/test_helpers.py @@ -5,16 +5,17 @@ from ..test_cases import SkipTest class FailingBulkClient(object): - def __init__(self, client, fail_at=1): + def __init__(self, client, fail_at=(2, ), fail_with=TransportError(599, "Error!", {})): self.client = client - self._called = -1 + self._called = 0 self._fail_at = fail_at self.transport = client.transport + self._fail_with = fail_with def bulk(self, *args, **kwargs): self._called += 1 - if self._called == self._fail_at: - raise TransportError(599, "Error!", {}) + if self._called in self._fail_at: + raise self._fail_with return self.client.bulk(*args, **kwargs) class TestStreamingBulk(ElasticsearchTestCase): @@ -87,7 +88,6 @@ class TestStreamingBulk(ElasticsearchTestCase): '_index': 'i', '_type': 't', '_id': 45, - 'data': {'f': 'v'}, 'error': "TransportError(599, 'Error!')", 'status': 599 @@ -96,6 +96,46 @@ class TestStreamingBulk(ElasticsearchTestCase): results[1][1] ) + def test_rejected_documents_are_retried(self): + failing_client = FailingBulkClient(self.client, fail_with=TransportError(429, 'Rejected!', {})) + docs = [ + {'_index': 'i', '_type': 't', '_id': 47, 'f': 'v'}, + {'_index': 'i', '_type': 't', '_id': 45, 'f': 'v'}, + {'_index': 'i', '_type': 't', '_id': 42, 'f': 'v'}, + ] + results = list(helpers.streaming_bulk(failing_client, docs, + raise_on_exception=False, + raise_on_error=False, + chunk_size=1, max_retries=1, + initial_backoff=0)) + self.assertEquals(3, len(results)) + self.assertEquals([True, True, True], [r[0] for r in results]) + self.client.indices.refresh(index='i') + res = self.client.search(index='i') + self.assertEquals(3, res['hits']['total']) + self.assertEquals(4, failing_client._called) + + def test_rejected_documents_are_retried_at_most_max_retries_times(self): + failing_client = FailingBulkClient(self.client, fail_at=(1, 2, ), + fail_with=TransportError(429, 'Rejected!', {})) + + docs = [ + {'_index': 'i', '_type': 't', '_id': 47, 'f': 'v'}, + {'_index': 'i', '_type': 't', '_id': 45, 'f': 'v'}, + {'_index': 'i', '_type': 't', '_id': 42, 'f': 'v'}, + ] + results = list(helpers.streaming_bulk(failing_client, docs, + raise_on_exception=False, + raise_on_error=False, + chunk_size=1, max_retries=1, + initial_backoff=0)) + self.assertEquals(3, len(results)) + self.assertEquals([False, True, True], [r[0] for r in results]) + self.client.indices.refresh(index='i') + res = self.client.search(index='i') + self.assertEquals(2, res['hits']['total']) + self.assertEquals(4, failing_client._called) + class TestBulk(ElasticsearchTestCase): def test_bulk_works_with_single_item(self):