Tests for retry logic in streaming_bulk

Author: honzakral

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}

test_elasticsearch/test_server/test_helpers.py

@@ -5,16 +5,17 @@
 
 
 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 @@ def test_transport_error_can_becaught(self):
  '_index': 'i',
  '_type': 't',
  '_id': 45,
-
  'data': {'f': 'v'},
  'error': "TransportError(599, 'Error!')",
  'status': 599
@@ -96,6 +96,46 @@ def test_transport_error_can_becaught(self):
  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):