Fixed bug in async_streaming_bulk()

Author: sethmlarson

elasticsearch/_async/compat.py

@@ -26,22 +26,17 @@ def get_sleep():
 
     async def sleep(duration):
         await asyncio.sleep(duration, loop=loop)
+
     return sleep
 
 
 def azip(*iterables):
-    print("AZIP", iterables)
     iterators = [aiter(x) for x in iterables]
-    print("AZIPTOR", iterators)
 
     async def generator():
         while True:
             try:
-                tuple_items = []
-                for iterator in iterators:
-                    tuple_items.append(await iterator.__anext__())
-                print("azip tuple", tuple_items)
-                yield tuple(tuple_items)
+                yield tuple([await i.__anext__() for i in iterators])
             except StopAsyncIteration:
                 break
 

elasticsearch/_async/helpers/__init__.py

@@ -0,0 +1,4 @@
+# Licensed to Elasticsearch B.V under one or more agreements.
+# Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
+# See the LICENSE file in the project root for more information
+

elasticsearch/_async/helpers/actions.py

@@ -115,7 +115,6 @@ async def _process_bulk_chunk(
     """
     Send a bulk request to elasticsearch and process the output.
     """
-    print("BULK DATA", bulk_data)
     # if raise on error is set, we need to collect errors per chunk before raising them
     errors = []
 
@@ -226,7 +225,10 @@ async def actions_generator():
 
     async def generator():
         async for bulk_data, bulk_actions in _chunk_actions(
-            actions_generator(), chunk_size, max_chunk_bytes, client.transport.serializer
+            aiter(actions_generator()),
+            chunk_size,
+            max_chunk_bytes,
+            client.transport.serializer,
         ):
 
             for attempt in range(max_retries + 1):
@@ -235,17 +237,18 @@ async def generator():
                     await sleep(min(max_backoff, initial_backoff * 2 ** (attempt - 1)))
 
                 try:
-                    print("before zip", bulk_actions, bulk_data)
-                    async for data, (ok, info) in azip(bulk_actions, _process_bulk_chunk(
-                        client,
-                        bulk_actions,
+                    async for data, (ok, info) in azip(
                         bulk_data,
-                        raise_on_exception,
-                        raise_on_error,
-                        *args,
-                        **kwargs
-                    )):
-                        print("zipped", data, ok, info)
+                        _process_bulk_chunk(
+                            client,
+                            bulk_actions,
+                            bulk_data,
+                            raise_on_exception,
+                            raise_on_error,
+                            *args,
+                            **kwargs
+                        ),
+                    ):
                         if not ok:
                             action, info = info.popitem()
                             # retry if retries enabled, we get 429, and we are not
@@ -257,7 +260,6 @@ async def generator():
                             ):
                                 # _process_bulk_chunk expects strings so we need to
                                 # re-serialize the data
-                                print("RETRY", data)
                                 to_retry.extend(
                                     map(client.transport.serializer.dumps, data)
                                 )
@@ -442,6 +444,7 @@ def scan(
         )
 
     """
+
     async def generator(query, scroll_kwargs):
         scroll_kwargs = scroll_kwargs or {}
 
@@ -451,7 +454,11 @@ async def generator(query, scroll_kwargs):
 
         # initial search
         resp = await client.search(
-            body=query, scroll=scroll, size=size, request_timeout=request_timeout, **kwargs
+            body=query,
+            scroll=scroll,
+            size=size,
+            request_timeout=request_timeout,
+            **kwargs
         )
         scroll_id = resp.get("_scroll_id")
 
@@ -487,7 +494,9 @@ async def generator(query, scroll_kwargs):
 
         finally:
             if scroll_id and clear_scroll:
-                await client.clear_scroll(body={"scroll_id": [scroll_id]}, ignore=(404,))
+                await client.clear_scroll(
+                    body={"scroll_id": [scroll_id]}, ignore=(404,)
+                )
 
     return aiter(generator(query, scroll_kwargs))
 

elasticsearch/client/__init__.py

@@ -229,6 +229,9 @@ def __enter__(self):
         return self
 
     def __exit__(self, *_):
+        self.close()
+
+    def close(self):
         self.transport.close()
 
     # AUTO-GENERATED-API-DEFINITIONS #

elasticsearch/compat.py

@@ -17,6 +17,7 @@
     string_types = str, bytes
     from urllib.parse import quote, quote_plus, urlencode, urlparse, unquote
 
+    zip = zip
     map = map
     from queue import Queue
 
@@ -25,25 +26,6 @@ def get_sleep():
     return time.sleep
 
 
-def zip(*iterables):
-    print("ZIP", iterables)
-    iterators = [iter(x) for x in iterables]
-    print("ZIPTOR", iterators)
-
-    def generator():
-        while True:
-            try:
-                tuple_items = []
-                for iterator in iterators:
-                    tuple_items.append(iterator.__next__())
-                print("zip tuple", tuple_items)
-                yield tuple(tuple_items)
-            except StopIteration:
-                break
-
-    return generator().__iter__()
-
-
 # These match against 'anext' and 'aiter'
 next = next
 iter = iter