From 680ab852754ee9da74e910dcd9d98be1a65ed8e8 Mon Sep 17 00:00:00 2001 From: Seth Michael Larson Date: Thu, 20 Aug 2020 07:18:23 -0500 Subject: [PATCH] [7.x] Add async helpers --- elasticsearch/_async/helpers.py | 430 ++++++++++ elasticsearch/helpers/__init__.py | 18 + elasticsearch/helpers/actions.py | 194 +++-- .../test_async/test_server/test_helpers.py | 751 ++++++++++++++++++ utils/build_dists.py | 5 + 5 files changed, 1324 insertions(+), 74 deletions(-) create mode 100644 elasticsearch/_async/helpers.py create mode 100644 test_elasticsearch/test_async/test_server/test_helpers.py diff --git a/elasticsearch/_async/helpers.py b/elasticsearch/_async/helpers.py new file mode 100644 index 00000000..34b927a0 --- /dev/null +++ b/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 + ) diff --git a/elasticsearch/helpers/__init__.py b/elasticsearch/helpers/__init__.py index 32fbbde3..08871a96 100644 --- a/elasticsearch/helpers/__init__.py +++ b/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 diff --git a/elasticsearch/helpers/actions.py b/elasticsearch/helpers/actions.py index 25f8bae7..87fcc305 100644 --- a/elasticsearch/helpers/actions.py +++ b/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 = [], [] - size, action_count = 0, 0 + chunker = _ActionChunker( + chunk_size=chunk_size, max_chunk_bytes=max_chunk_bytes, serializer=serializer + ) for action, data in actions: - raw_data, raw_action = data, action - action = 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 = 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 + ret = chunker.feed(action, data) + if ret: + yield ret + ret = chunker.flush() + if ret: + yield ret -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. - """ +def _process_bulk_chunk_success(resp, bulk_data, raise_on_error=True): # 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, diff --git a/test_elasticsearch/test_async/test_server/test_helpers.py b/test_elasticsearch/test_async/test_server/test_helpers.py new file mode 100644 index 00000000..e9d12ae6 --- /dev/null +++ b/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 diff --git a/utils/build_dists.py b/utils/build_dists.py index 0a021c15..a9602561 100644 --- a/utils/build_dists.py +++ b/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"):