[7.x] Add async helpers

elasticsearch/_async/helpers.py

@@ -0,0 +1,430 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+# 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
+
+import asyncio
+
+from .client import AsyncElasticsearch  # noqa
+from ..exceptions import TransportError
+from ..compat import map
+
+from ..helpers.actions import (
+    _ActionChunker,
+    _process_bulk_chunk_error,
+    _process_bulk_chunk_success,
+    expand_action,
+)
+from ..helpers.errors import ScanError
+
+import logging
+
+
+logger = logging.getLogger("elasticsearch.helpers")
+
+
+async def _chunk_actions(actions, chunk_size, max_chunk_bytes, serializer):
+    """
+    Split actions into chunks by number or size, serialize them into strings in
+    the process.
+    """
+    chunker = _ActionChunker(
+        chunk_size=chunk_size, max_chunk_bytes=max_chunk_bytes, serializer=serializer
+    )
+    async for action, data in actions:
+        ret = chunker.feed(action, data)
+        if ret:
+            yield ret
+    ret = chunker.flush()
+    if ret:
+        yield ret
+
+
+async def _process_bulk_chunk(
+    client,
+    bulk_actions,
+    bulk_data,
+    raise_on_exception=True,
+    raise_on_error=True,
+    *args,
+    **kwargs
+):
+    """
+    Send a bulk request to elasticsearch and process the output.
+    """
+    try:
+        # send the actual request
+        resp = await client.bulk("\n".join(bulk_actions) + "\n", *args, **kwargs)
+    except TransportError as e:
+        gen = _process_bulk_chunk_error(
+            error=e,
+            bulk_data=bulk_data,
+            raise_on_exception=raise_on_exception,
+            raise_on_error=raise_on_error,
+        )
+    else:
+        gen = _process_bulk_chunk_success(
+            resp=resp, bulk_data=bulk_data, raise_on_error=raise_on_error
+        )
+    for item in gen:
+        yield item
+
+
+def aiter(x):
+    """Turns an async iterable or iterable into an async iterator"""
+    if hasattr(x, "__anext__"):
+        return x
+    elif hasattr(x, "__aiter__"):
+        return x.__aiter__()
+
+    async def f():
+        for item in x:
+            yield item
+
+    return f().__aiter__()
+
+
+async def azip(*iterables):
+    """Zips async iterables and iterables into an async iterator
+    with the same behavior as zip()
+    """
+    aiters = [aiter(x) for x in iterables]
+    try:
+        while True:
+            yield tuple([await x.__anext__() for x in aiters])
+    except StopAsyncIteration:
+        pass
+
+
+async def async_streaming_bulk(
+    client,
+    actions,
+    chunk_size=500,
+    max_chunk_bytes=100 * 1024 * 1024,
+    raise_on_error=True,
+    expand_action_callback=expand_action,
+    raise_on_exception=True,
+    max_retries=0,
+    initial_backoff=2,
+    max_backoff=600,
+    yield_ok=True,
+    *args,
+    **kwargs
+):
+
+    """
+    Streaming bulk consumes actions from the iterable passed in and yields
+    results per action. For non-streaming usecases use
+    :func:`~elasticsearch.helpers.async_bulk` which is a wrapper around streaming
+    bulk that returns summary information about the bulk operation once the
+    entire input is consumed and sent.
+
+    If you specify ``max_retries`` it will also retry any documents that were
+    rejected with a ``429`` status code. To do this it will wait (**by calling
+    asyncio.sleep**) for ``initial_backoff`` seconds and then,
+    every subsequent rejection for the same chunk, for double the time every
+    time up to ``max_backoff`` seconds.
+
+    :arg client: instance of :class:`~elasticsearch.AsyncElasticsearch` to use
+    :arg actions: iterable or async 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
+        call to ``bulk`` and just report the items that failed as failed.
+    :arg expand_action_callback: callback executed on each action passed in,
+        should return a tuple containing the action line and the data line
+        (`None` if data line should be omitted).
+    :arg max_retries: maximum number of times a document will be retried when
+        ``429`` is received, set to 0 (default) for no retries on ``429``
+    :arg initial_backoff: number of seconds we should wait before the first
+        retry. Any subsequent retries will be powers of ``initial_backoff *
+        2**retry_number``
+    :arg max_backoff: maximum number of seconds a retry will wait
+    :arg yield_ok: if set to False will skip successful documents in the output
+    """
+
+    async def map_actions():
+        async for item in aiter(actions):
+            yield expand_action_callback(item)
+
+    async for bulk_data, bulk_actions in _chunk_actions(
+        map_actions(), chunk_size, max_chunk_bytes, client.transport.serializer
+    ):
+
+        for attempt in range(max_retries + 1):
+            to_retry, to_retry_data = [], []
+            if attempt:
+                await asyncio.sleep(
+                    min(max_backoff, initial_backoff * 2 ** (attempt - 1))
+                )
+
+            try:
+                async for data, (ok, info) in azip(
+                    bulk_data,
+                    _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
+                        # in the last attempt
+                        if (
+                            max_retries
+                            and info["status"] == 429
+                            and (attempt + 1) <= max_retries
+                        ):
+                            # _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}
+                    elif yield_ok:
+                        yield ok, info
+
+            except TransportError as e:
+                # suppress 429 errors since we will retry them
+                if attempt == max_retries or e.status_code != 429:
+                    raise
+            else:
+                if not to_retry:
+                    break
+                # retry only subset of documents that didn't succeed
+                bulk_actions, bulk_data = to_retry, to_retry_data
+
+
+async def async_bulk(client, actions, stats_only=False, *args, **kwargs):
+    """
+    Helper for the :meth:`~elasticsearch.AsyncElasticsearch.bulk` api that provides
+    a more human friendly interface - it consumes an iterator of actions and
+    sends them to elasticsearch in chunks. It returns a tuple with summary
+    information - number of successfully executed actions and either list of
+    errors or number of errors if ``stats_only`` is set to ``True``. Note that
+    by default we raise a ``BulkIndexError`` when we encounter an error so
+    options like ``stats_only`` only+ apply when ``raise_on_error`` is set to
+    ``False``.
+
+    When errors are being collected original document data is included in the
+    error dictionary which can lead to an extra high memory usage. If you need
+    to process a lot of data and want to ignore/collect errors please consider
+    using the :func:`~elasticsearch.helpers.async_streaming_bulk` helper which will
+    just return the errors and not store them in memory.
+
+
+    :arg client: instance of :class:`~elasticsearch.AsyncElasticsearch` to use
+    :arg actions: iterator containing the actions
+    :arg stats_only: if `True` only report number of successful/failed
+        operations instead of just number of successful and a list of error responses
+
+    Any additional keyword arguments will be passed to
+    :func:`~elasticsearch.helpers.async_streaming_bulk` which is used to execute
+    the operation, see :func:`~elasticsearch.helpers.async_streaming_bulk` for more
+    accepted parameters.
+    """
+    success, failed = 0, 0
+
+    # list of errors to be collected is not stats_only
+    errors = []
+
+    # make streaming_bulk yield successful results so we can count them
+    kwargs["yield_ok"] = True
+    async for ok, item in async_streaming_bulk(client, actions, *args, **kwargs):
+        # go through request-response pairs and detect failures
+        if not ok:
+            if not stats_only:
+                errors.append(item)
+            failed += 1
+        else:
+            success += 1
+
+    return success, failed if stats_only else errors
+
+
+async def async_scan(
+    client,
+    query=None,
+    scroll="5m",
+    raise_on_error=True,
+    preserve_order=False,
+    size=1000,
+    request_timeout=None,
+    clear_scroll=True,
+    scroll_kwargs=None,
+    **kwargs
+):
+    """
+    Simple abstraction on top of the
+    :meth:`~elasticsearch.AsyncElasticsearch.scroll` api - a simple iterator that
+    yields all hits as returned by underlining scroll requests.
+
+    By default scan does not return results in any pre-determined order. To
+    have a standard order in the returned documents (either by score or
+    explicit sort definition) when scrolling, use ``preserve_order=True``. This
+    may be an expensive operation and will negate the performance benefits of
+    using ``scan``.
+
+    :arg client: instance of :class:`~elasticsearch.AsyncElasticsearch` to use
+    :arg query: body for the :meth:`~elasticsearch.AsyncElasticsearch.search` api
+    :arg scroll: Specify how long a consistent view of the index should be
+        maintained for scrolled search
+    :arg raise_on_error: raises an exception (``ScanError``) if an error is
+        encountered (some shards fail to execute). By default we raise.
+    :arg preserve_order: don't set the ``search_type`` to ``scan`` - this will
+        cause the scroll to paginate with preserving the order. Note that this
+        can be an extremely expensive operation and can easily lead to
+        unpredictable results, use with caution.
+    :arg size: size (per shard) of the batch send at each iteration.
+    :arg request_timeout: explicit timeout for each call to ``scan``
+    :arg clear_scroll: explicitly calls delete on the scroll id via the clear
+        scroll API at the end of the method on completion or error, defaults
+        to true.
+    :arg scroll_kwargs: additional kwargs to be passed to
+        :meth:`~elasticsearch.AsyncElasticsearch.scroll`
+
+    Any additional keyword arguments will be passed to the initial
+    :meth:`~elasticsearch.AsyncElasticsearch.search` call::
+
+        async_scan(es,
+            query={"query": {"match": {"title": "python"}}},
+            index="orders-*",
+            doc_type="books"
+        )
+
+    """
+    scroll_kwargs = scroll_kwargs or {}
+
+    if not preserve_order:
+        query = query.copy() if query else {}
+        query["sort"] = "_doc"
+
+    # initial search
+    resp = await client.search(
+        body=query, scroll=scroll, size=size, request_timeout=request_timeout, **kwargs
+    )
+    scroll_id = resp.get("_scroll_id")
+
+    try:
+        while scroll_id and resp["hits"]["hits"]:
+            for hit in resp["hits"]["hits"]:
+                yield hit
+
+            # check if we have any errors
+            if (resp["_shards"]["successful"] + resp["_shards"]["skipped"]) < resp[
+                "_shards"
+            ]["total"]:
+                logger.warning(
+                    "Scroll request has only succeeded on %d (+%d skipped) shards out of %d.",
+                    resp["_shards"]["successful"],
+                    resp["_shards"]["skipped"],
+                    resp["_shards"]["total"],
+                )
+                if raise_on_error:
+                    raise ScanError(
+                        scroll_id,
+                        "Scroll request has only succeeded on %d (+%d skiped) shards out of %d."
+                        % (
+                            resp["_shards"]["successful"],
+                            resp["_shards"]["skipped"],
+                            resp["_shards"]["total"],
+                        ),
+                    )
+            resp = await client.scroll(
+                body={"scroll_id": scroll_id, "scroll": scroll}, **scroll_kwargs
+            )
+            scroll_id = resp.get("_scroll_id")
+
+    finally:
+        if scroll_id and clear_scroll:
+            await client.clear_scroll(body={"scroll_id": [scroll_id]}, ignore=(404,))
+
+
+async def async_reindex(
+    client,
+    source_index,
+    target_index,
+    query=None,
+    target_client=None,
+    chunk_size=500,
+    scroll="5m",
+    scan_kwargs={},
+    bulk_kwargs={},
+):
+
+    """
+    Reindex all documents from one index that satisfy a given query
+    to another, potentially (if `target_client` is specified) on a different cluster.
+    If you don't specify the query you will reindex all the documents.
+
+    Since ``2.3`` a :meth:`~elasticsearch.AsyncElasticsearch.reindex` api is
+    available as part of elasticsearch itself. It is recommended to use the api
+    instead of this helper wherever possible. The helper is here mostly for
+    backwards compatibility and for situations where more flexibility is
+    needed.
+
+    .. note::
+
+        This helper doesn't transfer mappings, just the data.
+
+    :arg client: instance of :class:`~elasticsearch.AsyncElasticsearch` to use (for
+        read if `target_client` is specified as well)
+    :arg source_index: index (or list of indices) to read documents from
+    :arg target_index: name of the index in the target cluster to populate
+    :arg query: body for the :meth:`~elasticsearch.AsyncElasticsearch.search` api
+    :arg target_client: optional, is specified will be used for writing (thus
+        enabling reindex between clusters)
+    :arg chunk_size: number of docs in one chunk sent to es (default: 500)
+    :arg scroll: Specify how long a consistent view of the index should be
+        maintained for scrolled search
+    :arg scan_kwargs: additional kwargs to be passed to
+        :func:`~elasticsearch.helpers.async_scan`
+    :arg bulk_kwargs: additional kwargs to be passed to
+        :func:`~elasticsearch.helpers.async_bulk`
+    """
+    target_client = client if target_client is None else target_client
+    docs = async_scan(
+        client, query=query, index=source_index, scroll=scroll, **scan_kwargs
+    )
+
+    async def _change_doc_index(hits, index):
+        async for h in hits:
+            h["_index"] = index
+            if "fields" in h:
+                h.update(h.pop("fields"))
+            yield h
+
+    kwargs = {"stats_only": True}
+    kwargs.update(bulk_kwargs)
+    return await async_bulk(
+        target_client,
+        _change_doc_index(docs, target_index),
+        chunk_size=chunk_size,
+        **kwargs
+    )

elasticsearch/helpers/__init__.py

@@ -15,6 +15,7 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
+import sys
 from .errors import BulkIndexError, ScanError
 from .actions import expand_action, streaming_bulk, bulk, parallel_bulk
 from .actions import scan, reindex
@@ -32,3 +33,20 @@ __all__ = [
     "_chunk_actions",
     "_process_bulk_chunk",
 ]
+
+
+try:
+    # Asyncio only supported on Python 3.6+
+    if sys.version_info < (3, 6):
+        raise ImportError
+
+    from .._async.helpers import (
+        async_scan,
+        async_bulk,
+        async_reindex,
+        async_streaming_bulk,
+    )
+
+    __all__ += ["async_scan", "async_bulk", "async_reindex", "async_streaming_bulk"]
+except (ImportError, SyntaxError):
+    pass

elasticsearch/helpers/actions.py

@@ -80,91 +80,77 @@ def expand_action(data):
     return action, data.get("_source", data)
 
 
+class _ActionChunker:
+    def __init__(self, chunk_size, max_chunk_bytes, serializer):
+        self.chunk_size = chunk_size
+        self.max_chunk_bytes = max_chunk_bytes
+        self.serializer = serializer
+
+        self.size = 0
+        self.action_count = 0
+        self.bulk_actions = []
+        self.bulk_data = []
+
+    def feed(self, action, data):
+        ret = None
+        raw_data, raw_action = data, action
+        action = self.serializer.dumps(action)
+        # +1 to account for the trailing new line character
+        cur_size = len(action.encode("utf-8")) + 1
+
+        if data is not None:
+            data = self.serializer.dumps(data)
+            cur_size += len(data.encode("utf-8")) + 1
+
+        # full chunk, send it and start a new one
+        if self.bulk_actions and (
+            self.size + cur_size > self.max_chunk_bytes
+            or self.action_count == self.chunk_size
+        ):
+            ret = (self.bulk_data, self.bulk_actions)
+            self.bulk_actions, self.bulk_data = [], []
+            self.size, self.action_count = 0, 0
+
+        self.bulk_actions.append(action)
+        if data is not None:
+            self.bulk_actions.append(data)
+            self.bulk_data.append((raw_action, raw_data))
+        else:
+            self.bulk_data.append((raw_action,))
+
+        self.size += cur_size
+        self.action_count += 1
+        return ret
+
+    def flush(self):
+        ret = None
+        if self.bulk_actions:
+            ret = (self.bulk_data, self.bulk_actions)
+            self.bulk_actions, self.bulk_data = [], []
+        return ret
+
+
 def _chunk_actions(actions, chunk_size, max_chunk_bytes, serializer):
     """
     Split actions into chunks by number or size, serialize them into strings in
     the process.
     """
-    bulk_actions, bulk_data = [], []
+    chunker = _ActionChunker(
-    size, action_count = 0, 0
+        chunk_size=chunk_size, max_chunk_bytes=max_chunk_bytes, serializer=serializer
+    )
     for action, data in actions:
-        raw_data, raw_action = data, action
+        ret = chunker.feed(action, data)
-        action = serializer.dumps(action)
+        if ret:
-        # +1 to account for the trailing new line character
+            yield ret
-        cur_size = len(action.encode("utf-8")) + 1
+    ret = chunker.flush()
-
+    if ret:
-        if data is not None:
+        yield ret
-            data = serializer.dumps(data)
-            cur_size += len(data.encode("utf-8")) + 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_data, bulk_actions
-            bulk_actions, bulk_data = [], []
-            size, action_count = 0, 0
-
-        bulk_actions.append(action)
-        if data is not None:
-            bulk_actions.append(data)
-            bulk_data.append((raw_action, raw_data))
-        else:
-            bulk_data.append((raw_action,))
-
-        size += cur_size
-        action_count += 1
-
-    if bulk_actions:
-        yield bulk_data, bulk_actions
 
 
-def _process_bulk_chunk(
+def _process_bulk_chunk_success(resp, bulk_data, raise_on_error=True):
-    client,
-    bulk_actions,
-    bulk_data,
-    raise_on_exception=True,
-    raise_on_error=True,
-    *args,
-    **kwargs
-):
-    """
-    Send a bulk request to elasticsearch and process the output.
-    """
     # if raise on error is set, we need to collect errors per chunk before raising them
     errors = []
 
-    try:
-        # send the actual request
-        resp = client.bulk("\n".join(bulk_actions) + "\n", *args, **kwargs)
-    except TransportError as e:
-        # default behavior - just propagate exception
-        if raise_on_exception:
-            raise e
-
-        # if we are not propagating, mark all actions in current chunk as failed
-        err_message = str(e)
-        exc_errors = []
-
-        for data in bulk_data:
-            # collect all the information about failed actions
-            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]
-            info.update(action)
-            exc_errors.append({op_type: info})
-
-        # emulate standard behavior for failed actions
-        if raise_on_error:
-            raise BulkIndexError(
-                "%i document(s) failed to index." % len(exc_errors), exc_errors
-            )
-        else:
-            for err in exc_errors:
-                yield False, err
-            return
-
     # go through request-response pairs and detect failures
     for data, (op_type, item) in zip(
         bulk_data, map(methodcaller("popitem"), resp["items"])
@@ -185,6 +171,66 @@ def _process_bulk_chunk(
         raise BulkIndexError("%i document(s) failed to index." % len(errors), errors)
 
 
+def _process_bulk_chunk_error(
+    error, bulk_data, raise_on_exception=True, raise_on_error=True
+):
+    # default behavior - just propagate exception
+    if raise_on_exception:
+        raise error
+
+    # if we are not propagating, mark all actions in current chunk as failed
+    err_message = str(error)
+    exc_errors = []
+
+    for data in bulk_data:
+        # collect all the information about failed actions
+        op_type, action = data[0].copy().popitem()
+        info = {"error": err_message, "status": error.status_code, "exception": error}
+        if op_type != "delete":
+            info["data"] = data[1]
+        info.update(action)
+        exc_errors.append({op_type: info})
+
+    # emulate standard behavior for failed actions
+    if raise_on_error:
+        raise BulkIndexError(
+            "%i document(s) failed to index." % len(exc_errors), exc_errors
+        )
+    else:
+        for err in exc_errors:
+            yield False, err
+
+
+def _process_bulk_chunk(
+    client,
+    bulk_actions,
+    bulk_data,
+    raise_on_exception=True,
+    raise_on_error=True,
+    *args,
+    **kwargs
+):
+    """
+    Send a bulk request to elasticsearch and process the output.
+    """
+    try:
+        # send the actual request
+        resp = client.bulk("\n".join(bulk_actions) + "\n", *args, **kwargs)
+    except TransportError as e:
+        gen = _process_bulk_chunk_error(
+            error=e,
+            bulk_data=bulk_data,
+            raise_on_exception=raise_on_exception,
+            raise_on_error=raise_on_error,
+        )
+    else:
+        gen = _process_bulk_chunk_success(
+            resp=resp, bulk_data=bulk_data, raise_on_error=raise_on_error
+        )
+    for item in gen:
+        yield item
+
+
 def streaming_bulk(
     client,
     actions,

test_elasticsearch/test_async/test_server/test_helpers.py

@@ -0,0 +1,751 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+# 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
+
+import pytest
+import asyncio
+from mock import patch, MagicMock
+
+from elasticsearch import helpers, TransportError
+from elasticsearch.helpers import ScanError
+
+
+pytestmark = pytest.mark.asyncio
+
+
+class AsyncMock(MagicMock):
+    async def __call__(self, *args, **kwargs):
+        return super(AsyncMock, self).__call__(*args, **kwargs)
+
+
+class FailingBulkClient(object):
+    def __init__(
+        self, client, fail_at=(2,), fail_with=TransportError(599, "Error!", {})
+    ):
+        self.client = client
+        self._called = 0
+        self._fail_at = fail_at
+        self.transport = client.transport
+        self._fail_with = fail_with
+
+    async def bulk(self, *args, **kwargs):
+        self._called += 1
+        if self._called in self._fail_at:
+            raise self._fail_with
+        return await self.client.bulk(*args, **kwargs)
+
+
+class TestStreamingBulk(object):
+    async def test_actions_remain_unchanged(self, async_client):
+        actions = [{"_id": 1}, {"_id": 2}]
+        async for ok, item in helpers.async_streaming_bulk(
+            async_client, actions, index="test-index"
+        ):
+            assert ok
+        assert [{"_id": 1}, {"_id": 2}] == actions
+
+    async def test_all_documents_get_inserted(self, async_client):
+        docs = [{"answer": x, "_id": x} for x in range(100)]
+        async for ok, item in helpers.async_streaming_bulk(
+            async_client, docs, index="test-index", refresh=True
+        ):
+            assert ok
+
+        assert 100 == (await async_client.count(index="test-index"))["count"]
+        assert {"answer": 42} == (await async_client.get(index="test-index", id=42))[
+            "_source"
+        ]
+
+    async def test_documents_data_types(self, async_client):
+        async def async_gen():
+            for x in range(100):
+                await asyncio.sleep(0)
+                yield {"answer": x, "_id": x}
+
+        def sync_gen():
+            for x in range(100):
+                yield {"answer": x, "_id": x}
+
+        async for ok, item in helpers.async_streaming_bulk(
+            async_client, async_gen(), index="test-index", refresh=True
+        ):
+            assert ok
+
+        assert 100 == (await async_client.count(index="test-index"))["count"]
+        assert {"answer": 42} == (await async_client.get(index="test-index", id=42))[
+            "_source"
+        ]
+
+        await async_client.delete_by_query(
+            index="test-index", body={"query": {"match_all": {}}}
+        )
+
+        async for ok, item in helpers.async_streaming_bulk(
+            async_client, sync_gen(), index="test-index", refresh=True
+        ):
+            assert ok
+
+        assert 100 == (await async_client.count(index="test-index"))["count"]
+        assert {"answer": 42} == (await async_client.get(index="test-index", id=42))[
+            "_source"
+        ]
+
+    async def test_all_errors_from_chunk_are_raised_on_failure(self, async_client):
+        await async_client.indices.create(
+            "i",
+            {
+                "mappings": {"properties": {"a": {"type": "integer"}}},
+                "settings": {"number_of_shards": 1, "number_of_replicas": 0},
+            },
+        )
+        await async_client.cluster.health(wait_for_status="yellow")
+
+        try:
+            async for ok, item in helpers.async_streaming_bulk(
+                async_client, [{"a": "b"}, {"a": "c"}], index="i", raise_on_error=True
+            ):
+                assert ok
+        except helpers.BulkIndexError as e:
+            assert 2 == len(e.errors)
+        else:
+            assert False, "exception should have been raised"
+
+    async def test_different_op_types(self, async_client):
+        await async_client.index(index="i", id=45, body={})
+        await async_client.index(index="i", id=42, body={})
+        docs = [
+            {"_index": "i", "_id": 47, "f": "v"},
+            {"_op_type": "delete", "_index": "i", "_id": 45},
+            {"_op_type": "update", "_index": "i", "_id": 42, "doc": {"answer": 42}},
+        ]
+        async for ok, item in helpers.async_streaming_bulk(async_client, docs):
+            assert ok
+
+        assert not await async_client.exists(index="i", id=45)
+        assert {"answer": 42} == (await async_client.get(index="i", id=42))["_source"]
+        assert {"f": "v"} == (await async_client.get(index="i", id=47))["_source"]
+
+    async def test_transport_error_can_becaught(self, async_client):
+        failing_client = FailingBulkClient(async_client)
+        docs = [
+            {"_index": "i", "_id": 47, "f": "v"},
+            {"_index": "i", "_id": 45, "f": "v"},
+            {"_index": "i", "_id": 42, "f": "v"},
+        ]
+
+        results = [
+            x
+            async for x in helpers.async_streaming_bulk(
+                failing_client,
+                docs,
+                raise_on_exception=False,
+                raise_on_error=False,
+                chunk_size=1,
+            )
+        ]
+        assert 3 == len(results)
+        assert [True, False, True] == [r[0] for r in results]
+
+        exc = results[1][1]["index"].pop("exception")
+        assert isinstance(exc, TransportError)
+        assert 599 == exc.status_code
+        assert {
+            "index": {
+                "_index": "i",
+                "_id": 45,
+                "data": {"f": "v"},
+                "error": "TransportError(599, 'Error!')",
+                "status": 599,
+            }
+        } == results[1][1]
+
+    async def test_rejected_documents_are_retried(self, async_client):
+        failing_client = FailingBulkClient(
+            async_client, fail_with=TransportError(429, "Rejected!", {})
+        )
+        docs = [
+            {"_index": "i", "_id": 47, "f": "v"},
+            {"_index": "i", "_id": 45, "f": "v"},
+            {"_index": "i", "_id": 42, "f": "v"},
+        ]
+        results = [
+            x
+            async for x in helpers.async_streaming_bulk(
+                failing_client,
+                docs,
+                raise_on_exception=False,
+                raise_on_error=False,
+                chunk_size=1,
+                max_retries=1,
+                initial_backoff=0,
+            )
+        ]
+        assert 3 == len(results)
+        assert [True, True, True] == [r[0] for r in results]
+        await async_client.indices.refresh(index="i")
+        res = await async_client.search(index="i")
+        assert {"value": 3, "relation": "eq"} == res["hits"]["total"]
+        assert 4 == failing_client._called
+
+    async def test_rejected_documents_are_retried_at_most_max_retries_times(
+        self, async_client
+    ):
+        failing_client = FailingBulkClient(
+            async_client, fail_at=(1, 2), fail_with=TransportError(429, "Rejected!", {})
+        )
+
+        docs = [
+            {"_index": "i", "_id": 47, "f": "v"},
+            {"_index": "i", "_id": 45, "f": "v"},
+            {"_index": "i", "_id": 42, "f": "v"},
+        ]
+        results = [
+            x
+            async for x in helpers.async_streaming_bulk(
+                failing_client,
+                docs,
+                raise_on_exception=False,
+                raise_on_error=False,
+                chunk_size=1,
+                max_retries=1,
+                initial_backoff=0,
+            )
+        ]
+        assert 3 == len(results)
+        assert [False, True, True] == [r[0] for r in results]
+        await async_client.indices.refresh(index="i")
+        res = await async_client.search(index="i")
+        assert {"value": 2, "relation": "eq"} == res["hits"]["total"]
+        assert 4 == failing_client._called
+
+    async def test_transport_error_is_raised_with_max_retries(self, async_client):
+        failing_client = FailingBulkClient(
+            async_client,
+            fail_at=(1, 2, 3, 4),
+            fail_with=TransportError(429, "Rejected!", {}),
+        )
+
+        async def streaming_bulk():
+            results = [
+                x
+                async for x in helpers.async_streaming_bulk(
+                    failing_client,
+                    [{"a": 42}, {"a": 39}],
+                    raise_on_exception=True,
+                    max_retries=3,
+                    initial_backoff=0,
+                )
+            ]
+            return results
+
+        with pytest.raises(TransportError):
+            await streaming_bulk()
+        assert 4 == failing_client._called
+
+
+class TestBulk(object):
+    async def test_bulk_works_with_single_item(self, async_client):
+        docs = [{"answer": 42, "_id": 1}]
+        success, failed = await helpers.async_bulk(
+            async_client, docs, index="test-index", refresh=True
+        )
+
+        assert 1 == success
+        assert not failed
+        assert 1 == (await async_client.count(index="test-index"))["count"]
+        assert {"answer": 42} == (await async_client.get(index="test-index", id=1))[
+            "_source"
+        ]
+
+    async def test_all_documents_get_inserted(self, async_client):
+        docs = [{"answer": x, "_id": x} for x in range(100)]
+        success, failed = await helpers.async_bulk(
+            async_client, docs, index="test-index", refresh=True
+        )
+
+        assert 100 == success
+        assert not failed
+        assert 100 == (await async_client.count(index="test-index"))["count"]
+        assert {"answer": 42} == (await async_client.get(index="test-index", id=42))[
+            "_source"
+        ]
+
+    async def test_stats_only_reports_numbers(self, async_client):
+        docs = [{"answer": x} for x in range(100)]
+        success, failed = await helpers.async_bulk(
+            async_client, docs, index="test-index", refresh=True, stats_only=True
+        )
+
+        assert 100 == success
+        assert 0 == failed
+        assert 100 == (await async_client.count(index="test-index"))["count"]
+
+    async def test_errors_are_reported_correctly(self, async_client):
+        await async_client.indices.create(
+            "i",
+            {
+                "mappings": {"properties": {"a": {"type": "integer"}}},
+                "settings": {"number_of_shards": 1, "number_of_replicas": 0},
+            },
+        )
+        await async_client.cluster.health(wait_for_status="yellow")
+
+        success, failed = await helpers.async_bulk(
+            async_client,
+            [{"a": 42}, {"a": "c", "_id": 42}],
+            index="i",
+            raise_on_error=False,
+        )
+        assert 1 == success
+        assert 1 == len(failed)
+        error = failed[0]
+        assert "42" == error["index"]["_id"]
+        assert "i" == error["index"]["_index"]
+        print(error["index"]["error"])
+        assert "MapperParsingException" in repr(
+            error["index"]["error"]
+        ) or "mapper_parsing_exception" in repr(error["index"]["error"])
+
+    async def test_error_is_raised(self, async_client):
+        await async_client.indices.create(
+            "i",
+            {
+                "mappings": {"properties": {"a": {"type": "integer"}}},
+                "settings": {"number_of_shards": 1, "number_of_replicas": 0},
+            },
+        )
+        await async_client.cluster.health(wait_for_status="yellow")
+
+        with pytest.raises(helpers.BulkIndexError):
+            await helpers.async_bulk(async_client, [{"a": 42}, {"a": "c"}], index="i")
+
+    async def test_errors_are_collected_properly(self, async_client):
+        await async_client.indices.create(
+            "i",
+            {
+                "mappings": {"properties": {"a": {"type": "integer"}}},
+                "settings": {"number_of_shards": 1, "number_of_replicas": 0},
+            },
+        )
+        await async_client.cluster.health(wait_for_status="yellow")
+
+        success, failed = await helpers.async_bulk(
+            async_client,
+            [{"a": 42}, {"a": "c"}],
+            index="i",
+            stats_only=True,
+            raise_on_error=False,
+        )
+        assert 1 == success
+        assert 1 == failed
+
+
+class MockScroll:
+    def __init__(self):
+        self.calls = []
+
+    async def __call__(self, *args, **kwargs):
+        self.calls.append((args, kwargs))
+        if len(self.calls) == 1:
+            return {
+                "_scroll_id": "dummy_id",
+                "_shards": {"successful": 4, "total": 5, "skipped": 0},
+                "hits": {"hits": [{"scroll_data": 42}]},
+            }
+        elif len(self.calls) == 2:
+            return {
+                "_scroll_id": "dummy_id",
+                "_shards": {"successful": 4, "total": 5, "skipped": 0},
+                "hits": {"hits": []},
+            }
+        else:
+            raise Exception("no more responses")
+
+
+class MockResponse:
+    def __init__(self, resp):
+        self.resp = resp
+
+    async def __call__(self, *args, **kwargs):
+        return self.resp
+
+
+class TestScan(object):
+    async def teardown_method(self, async_client):
+        await async_client.transport.perform_request("DELETE", "/_search/scroll/_all")
+
+    async def test_order_can_be_preserved(self, async_client):
+        bulk = []
+        for x in range(100):
+            bulk.append({"index": {"_index": "test_index", "_id": x}})
+            bulk.append({"answer": x, "correct": x == 42})
+        await async_client.bulk(bulk, refresh=True)
+
+        docs = [
+            doc
+            async for doc in helpers.async_scan(
+                async_client,
+                index="test_index",
+                query={"sort": "answer"},
+                preserve_order=True,
+            )
+        ]
+
+        assert 100 == len(docs)
+        assert list(map(str, range(100))) == list(d["_id"] for d in docs)
+        assert list(range(100)) == list(d["_source"]["answer"] for d in docs)
+
+    async def test_all_documents_are_read(self, async_client):
+        bulk = []
+        for x in range(100):
+            bulk.append({"index": {"_index": "test_index", "_id": x}})
+            bulk.append({"answer": x, "correct": x == 42})
+        await async_client.bulk(bulk, refresh=True)
+
+        docs = [
+            x
+            async for x in helpers.async_scan(async_client, index="test_index", size=2)
+        ]
+
+        assert 100 == len(docs)
+        assert set(map(str, range(100))) == set(d["_id"] for d in docs)
+        assert set(range(100)) == set(d["_source"]["answer"] for d in docs)
+
+    async def test_scroll_error(self, async_client):
+        bulk = []
+        for x in range(4):
+            bulk.append({"index": {"_index": "test_index"}})
+            bulk.append({"value": x})
+        await async_client.bulk(bulk, refresh=True)
+
+        with patch.object(async_client, "scroll", MockScroll()):
+            data = [
+                x
+                async for x in helpers.async_scan(
+                    async_client,
+                    index="test_index",
+                    size=2,
+                    raise_on_error=False,
+                    clear_scroll=False,
+                )
+            ]
+            assert len(data) == 3
+            assert data[-1] == {"scroll_data": 42}
+
+        with patch.object(async_client, "scroll", MockScroll()):
+            with pytest.raises(ScanError):
+                data = [
+                    x
+                    async for x in helpers.async_scan(
+                        async_client,
+                        index="test_index",
+                        size=2,
+                        raise_on_error=True,
+                        clear_scroll=False,
+                    )
+                ]
+            assert len(data) == 3
+            assert data[-1] == {"scroll_data": 42}
+
+    async def test_initial_search_error(self, async_client):
+        with patch.object(async_client, "clear_scroll", new_callable=AsyncMock):
+            with patch.object(
+                async_client,
+                "search",
+                MockResponse(
+                    {
+                        "_scroll_id": "dummy_id",
+                        "_shards": {"successful": 4, "total": 5, "skipped": 0},
+                        "hits": {"hits": [{"search_data": 1}]},
+                    }
+                ),
+            ):
+                with patch.object(async_client, "scroll", MockScroll()):
+
+                    data = [
+                        x
+                        async for x in helpers.async_scan(
+                            async_client,
+                            index="test_index",
+                            size=2,
+                            raise_on_error=False,
+                        )
+                    ]
+                    assert data == [{"search_data": 1}, {"scroll_data": 42}]
+
+            with patch.object(
+                async_client,
+                "search",
+                MockResponse(
+                    {
+                        "_scroll_id": "dummy_id",
+                        "_shards": {"successful": 4, "total": 5, "skipped": 0},
+                        "hits": {"hits": [{"search_data": 1}]},
+                    }
+                ),
+            ):
+                with patch.object(async_client, "scroll", MockScroll()) as mock_scroll:
+
+                    with pytest.raises(ScanError):
+                        data = [
+                            x
+                            async for x in helpers.async_scan(
+                                async_client,
+                                index="test_index",
+                                size=2,
+                                raise_on_error=True,
+                            )
+                        ]
+                        assert data == [{"search_data": 1}]
+                        assert mock_scroll.calls == []
+
+    async def test_no_scroll_id_fast_route(self, async_client):
+        with patch.object(async_client, "search", MockResponse({"no": "_scroll_id"})):
+            with patch.object(async_client, "scroll") as scroll_mock:
+                with patch.object(async_client, "clear_scroll") as clear_mock:
+                    data = [
+                        x
+                        async for x in helpers.async_scan(
+                            async_client, index="test_index"
+                        )
+                    ]
+
+                    assert data == []
+                    scroll_mock.assert_not_called()
+                    clear_mock.assert_not_called()
+
+    @patch("elasticsearch._async.helpers.logger")
+    async def test_logger(self, logger_mock, async_client):
+        bulk = []
+        for x in range(4):
+            bulk.append({"index": {"_index": "test_index"}})
+            bulk.append({"value": x})
+        await async_client.bulk(bulk, refresh=True)
+
+        with patch.object(async_client, "scroll", MockScroll()):
+            _ = [
+                x
+                async for x in helpers.async_scan(
+                    async_client,
+                    index="test_index",
+                    size=2,
+                    raise_on_error=False,
+                    clear_scroll=False,
+                )
+            ]
+            logger_mock.warning.assert_called()
+
+        with patch.object(async_client, "scroll", MockScroll()):
+            try:
+                _ = [
+                    x
+                    async for x in helpers.async_scan(
+                        async_client,
+                        index="test_index",
+                        size=2,
+                        raise_on_error=True,
+                        clear_scroll=False,
+                    )
+                ]
+            except ScanError:
+                pass
+            logger_mock.warning.assert_called_with(
+                "Scroll request has only succeeded on %d (+%d skipped) shards out of %d.",
+                4,
+                0,
+                5,
+            )
+
+    async def test_clear_scroll(self, async_client):
+        bulk = []
+        for x in range(4):
+            bulk.append({"index": {"_index": "test_index"}})
+            bulk.append({"value": x})
+        await async_client.bulk(bulk, refresh=True)
+
+        with patch.object(
+            async_client, "clear_scroll", wraps=async_client.clear_scroll
+        ) as spy:
+            _ = [
+                x
+                async for x in helpers.async_scan(
+                    async_client, index="test_index", size=2
+                )
+            ]
+            spy.assert_called_once()
+
+            spy.reset_mock()
+            _ = [
+                x
+                async for x in helpers.async_scan(
+                    async_client, index="test_index", size=2, clear_scroll=True
+                )
+            ]
+            spy.assert_called_once()
+
+            spy.reset_mock()
+            _ = [
+                x
+                async for x in helpers.async_scan(
+                    async_client, index="test_index", size=2, clear_scroll=False
+                )
+            ]
+            spy.assert_not_called()
+
+
+@pytest.fixture(scope="function")
+async def reindex_setup(async_client):
+    bulk = []
+    for x in range(100):
+        bulk.append({"index": {"_index": "test_index", "_id": x}})
+        bulk.append(
+            {
+                "answer": x,
+                "correct": x == 42,
+                "type": "answers" if x % 2 == 0 else "questions",
+            }
+        )
+    await async_client.bulk(bulk, refresh=True)
+    yield
+
+
+class TestReindex(object):
+    async def test_reindex_passes_kwargs_to_scan_and_bulk(
+        self, async_client, reindex_setup
+    ):
+        await helpers.async_reindex(
+            async_client,
+            "test_index",
+            "prod_index",
+            scan_kwargs={"q": "type:answers"},
+            bulk_kwargs={"refresh": True},
+        )
+
+        assert await async_client.indices.exists("prod_index")
+        assert (
+            50
+            == (await async_client.count(index="prod_index", q="type:answers"))["count"]
+        )
+
+        assert {"answer": 42, "correct": True, "type": "answers"} == (
+            await async_client.get(index="prod_index", id=42)
+        )["_source"]
+
+    async def test_reindex_accepts_a_query(self, async_client, reindex_setup):
+        await helpers.async_reindex(
+            async_client,
+            "test_index",
+            "prod_index",
+            query={"query": {"bool": {"filter": {"term": {"type": "answers"}}}}},
+        )
+        await async_client.indices.refresh()
+
+        assert await async_client.indices.exists("prod_index")
+        assert (
+            50
+            == (await async_client.count(index="prod_index", q="type:answers"))["count"]
+        )
+
+        assert {"answer": 42, "correct": True, "type": "answers"} == (
+            await async_client.get(index="prod_index", id=42)
+        )["_source"]
+
+    async def test_all_documents_get_moved(self, async_client, reindex_setup):
+        await helpers.async_reindex(async_client, "test_index", "prod_index")
+        await async_client.indices.refresh()
+
+        assert await async_client.indices.exists("prod_index")
+        assert (
+            50
+            == (await async_client.count(index="prod_index", q="type:questions"))[
+                "count"
+            ]
+        )
+        assert (
+            50
+            == (await async_client.count(index="prod_index", q="type:answers"))["count"]
+        )
+
+        assert {"answer": 42, "correct": True, "type": "answers"} == (
+            await async_client.get(index="prod_index", id=42)
+        )["_source"]
+
+
+@pytest.fixture(scope="function")
+async def parent_reindex_setup(async_client):
+    body = {
+        "settings": {"number_of_shards": 1, "number_of_replicas": 0},
+        "mappings": {
+            "properties": {
+                "question_answer": {
+                    "type": "join",
+                    "relations": {"question": "answer"},
+                }
+            }
+        },
+    }
+    await async_client.indices.create(index="test-index", body=body)
+    await async_client.indices.create(index="real-index", body=body)
+
+    await async_client.index(
+        index="test-index", id=42, body={"question_answer": "question"}
+    )
+    await async_client.index(
+        index="test-index",
+        id=47,
+        routing=42,
+        body={"some": "data", "question_answer": {"name": "answer", "parent": 42}},
+    )
+    await async_client.indices.refresh(index="test-index")
+
+
+class TestParentChildReindex:
+    async def test_children_are_reindexed_correctly(
+        self, async_client, parent_reindex_setup
+    ):
+        await helpers.async_reindex(async_client, "test-index", "real-index")
+
+        q = await async_client.get(index="real-index", id=42)
+        assert {
+            "_id": "42",
+            "_index": "real-index",
+            "_primary_term": 1,
+            "_seq_no": 0,
+            "_source": {"question_answer": "question"},
+            "_type": "_doc",
+            "_version": 1,
+            "found": True,
+        } == q
+
+        q = await async_client.get(index="test-index", id=47, routing=42)
+        assert {
+            "_routing": "42",
+            "_id": "47",
+            "_index": "test-index",
+            "_primary_term": 1,
+            "_seq_no": 1,
+            "_source": {
+                "some": "data",
+                "question_answer": {"name": "answer", "parent": 42},
+            },
+            "_type": "_doc",
+            "_version": 1,
+            "found": True,
+        } == q

utils/build_dists.py

@@ -100,6 +100,11 @@ def test_dist(dist):
         # Install aiohttp and see that async is now available
         run(venv_python, "-m", "pip", "install", "aiohttp")
         run(venv_python, "-c", f"from {dist_name} import AsyncElasticsearch")
+        run(
+            venv_python,
+            "-c",
+            f"from {dist_name}.helpers import async_scan, async_bulk, async_streaming_bulk, async_reindex",
+        )
 
         # Ensure that the namespaces are correct for the dist
         for suffix in ("", "1", "2", "5", "6", "7", "8", "9", "10"):