From 3400179153cc13b6ae2c26734337202569bdfd80 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Honza=20Kr=C3=A1l?= <honza.kral@gmail.com>
Date: Wed, 30 Sep 2015 21:06:24 +0200
Subject: [PATCH] Allos the size of the bulk request to be defined in bytes

Fixes #199
---
 elasticsearch/helpers/__init__.py             | 41 ++++++++++++-------
 .../test_server/test_helpers.py               | 18 ++++++--
 2 files changed, 42 insertions(+), 17 deletions(-)

diff --git a/elasticsearch/helpers/__init__.py b/elasticsearch/helpers/__init__.py
index d1528c12..c1fc81bd 100644
--- a/elasticsearch/helpers/__init__.py
+++ b/elasticsearch/helpers/__init__.py
@@ -1,7 +1,6 @@
 from __future__ import unicode_literals
 
 import logging
-from itertools import islice
 from operator import methodcaller
 
 from ..exceptions import ElasticsearchException, TransportError
@@ -40,22 +39,35 @@ def expand_action(data):
 
     return action, data.get('_source', data)
 
-def _chunk_actions(actions, chunk_size, serializer):
-    while True:
-        bulk_actions = []
-        for action, data in islice(actions, chunk_size):
-            bulk_actions.append(serializer.dumps(action))
-            if data is not None:
-                bulk_actions.append(serializer.dumps(data))
+def _chunk_actions(actions, chunk_size, max_chunk_bytes, serializer):
+    bulk_actions = []
+    size, action_count = 0, 0
+    for action, data in actions:
+        action = serializer.dumps(action)
+        cur_size = len(action) + 1
 
-        if not bulk_actions:
-            return
+        if data is not None:
+            data = serializer.dumps(data)
+            cur_size += len(data) + 1
 
+        # full chunk, send it and start a new one
+        if bulk_actions and (size + cur_size > max_chunk_bytes or action_count == chunk_size):
+            yield bulk_actions
+            bulk_actions = []
+            size, action_count = 0, 0
+
+        bulk_actions.append(action)
+        if data is not None:
+            bulk_actions.append(data)
+        size += cur_size
+        action_count += 1
+
+    if bulk_actions:
         yield bulk_actions
 
-def streaming_bulk(client, actions, chunk_size=500, raise_on_error=True,
-        expand_action_callback=expand_action, raise_on_exception=True, 
-        **kwargs):
+def streaming_bulk(client, actions, chunk_size=500, max_chunk_bytes=100 * 1014 * 1024,
+        raise_on_error=True, expand_action_callback=expand_action,
+        raise_on_exception=True, **kwargs):
     """
     Streaming bulk consumes actions from the iterable passed in and yields
     results per action. For non-streaming usecases use
@@ -101,6 +113,7 @@ def streaming_bulk(client, actions, chunk_size=500, raise_on_error=True,
     :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 raise_on_error: raise ``BulkIndexError`` containing errors (as `.errors`)
         from the execution of the last chunk when some occur. By default we raise.
     :arg raise_on_exception: if ``False`` then don't propagate exceptions from
@@ -115,7 +128,7 @@ def streaming_bulk(client, actions, chunk_size=500, raise_on_error=True,
     # if raise on error is set, we need to collect errors per chunk before raising them
     errors = []
 
-    for bulk_actions in _chunk_actions(actions, chunk_size, serializer):
+    for bulk_actions in _chunk_actions(actions, chunk_size, max_chunk_bytes, serializer):
         try:
             # send the actual request
             resp = client.bulk('\n'.join(bulk_actions) + '\n', **kwargs)
diff --git a/test_elasticsearch/test_server/test_helpers.py b/test_elasticsearch/test_server/test_helpers.py
index 2b457240..292815f9 100644
--- a/test_elasticsearch/test_server/test_helpers.py
+++ b/test_elasticsearch/test_server/test_helpers.py
@@ -1,9 +1,21 @@
-from datetime import datetime
-
 from elasticsearch import helpers, TransportError
+from elasticsearch.serializer import JSONSerializer
 
 from . import ElasticsearchTestCase
-from ..test_cases import SkipTest
+from ..test_cases import SkipTest, TestCase
+
+
+class TestChunkActions(TestCase):
+    def setUp(self):
+        super(TestChunkActions, self).setUp()
+        self.actions = [({'index': {}}, {'some': 'data', 'i': i}) for i in range(100)]
+
+    def test_chunks_are_chopped_by_byte_size(self):
+        self.assertEquals(100, len(list(helpers._chunk_actions(self.actions, 100000, 1, JSONSerializer()))))
+
+    def test_chunks_are_chopped_by_chunk_size(self):
+        self.assertEquals(10, len(list(helpers._chunk_actions(self.actions, 10, 99999999, JSONSerializer()))))
+
 
 class FailingBulkClient(object):
     def __init__(self, client, fail_at=1):