diff --git a/elasticsearch/helpers/__init__.py b/elasticsearch/helpers/__init__.py index 66d17494..c1c6528a 100644 --- a/elasticsearch/helpers/__init__.py +++ b/elasticsearch/helpers/__init__.py @@ -95,6 +95,11 @@ def streaming_bulk(client, actions, chunk_size=500, raise_on_error=False, while True: chunk = islice(actions, chunk_size) + + # raise on exception means we might need to iterate on chunk twice + if not raise_on_exception: + chunk = list(chunk) + bulk_actions = [] for action, data in chunk: bulk_actions.append(action) diff --git a/test_elasticsearch/test_server/test_helpers.py b/test_elasticsearch/test_server/test_helpers.py index df0a3f70..ea141085 100644 --- a/test_elasticsearch/test_server/test_helpers.py +++ b/test_elasticsearch/test_server/test_helpers.py @@ -1,8 +1,20 @@ -from elasticsearch import helpers +from elasticsearch import helpers, TransportError from . import ElasticsearchTestCase from ..test_cases import SkipTest +class FailingBulkClient(object): + def __init__(self, client, fail_at=1): + self.client = client + self._called = -1 + self._fail_at = fail_at + + def bulk(self, *args, **kwargs): + self._called += 1 + if self._called == self._fail_at: + raise TransportError(599, "Error!", "INFO") + return self.client.bulk(*args, **kwargs) + class TestStreamingBulk(ElasticsearchTestCase): def test_actions_remain_unchanged(self): actions = [{'_id': 1}, {'_id': 2}] @@ -43,7 +55,7 @@ class TestStreamingBulk(ElasticsearchTestCase): docs = [ {'_index': 'i', '_type': 't', '_id': 47, 'f': 'v'}, {'_op_type': 'delete', '_index': 'i', '_type': 't', '_id': 45}, - {'_op_type': 'update', '_index': 'i', '_type': 't', '_id': 42, 'doc': {'answer': 42}}, + {'_op_type': 'update', '_index': 'i', '_type': 't', '_id': 42, 'doc': {'answer': 42}} ] for ok, item in helpers.streaming_bulk(self.client, docs): self.assertTrue(ok) @@ -52,6 +64,36 @@ class TestStreamingBulk(ElasticsearchTestCase): self.assertEquals({'answer': 42}, self.client.get(index='i', id=42)['_source']) self.assertEquals({'f': 'v'}, self.client.get(index='i', id=47)['_source']) + def test_transport_error_can_becaught(self): + failing_client = FailingBulkClient(self.client) + 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)) + self.assertEquals(3, len(results)) + self.assertEquals([True, False, True], [r[0] for r in results]) + + exc = results[1][1]['index'].pop('exception') + self.assertIsInstance(exc, TransportError) + self.assertEquals(599, exc.status_code) + self.assertEquals( + { + 'index': { + '_index': 'i', + '_type': 't', + '_id': 45, + + 'data': {'f': 'v'}, + 'error': "TransportError(599, 'Error!')", + 'status': 599 + } + }, + results[1][1] + ) + class TestBulk(ElasticsearchTestCase): def test_bulk_works_with_single_item(self):