[7.x] Add async helpers
This commit is contained in:
Seth Michael Larson
committed by
Seth Michael Larson
Seth Michael Larson
parent
62e2b16cdf
commit
680ab85275
@@ -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
|
||||
)
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
@@ -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"):
|
||||
|
||||
Reference in New Issue
Block a user