Add an option to use scan helper without search_type=scan, use with caution
Fixes #160
This commit is contained in:
@@ -156,7 +156,7 @@ def bulk(client, actions, stats_only=False, **kwargs):
|
|||||||
# preserve the name for backwards compatibility
|
# preserve the name for backwards compatibility
|
||||||
bulk_index = bulk
|
bulk_index = bulk
|
||||||
|
|
||||||
def scan(client, query=None, scroll='5m', **kwargs):
|
def scan(client, query=None, scroll='5m', preserve_order=False, **kwargs):
|
||||||
"""
|
"""
|
||||||
Simple abstraction on top of the
|
Simple abstraction on top of the
|
||||||
:meth:`~elasticsearch.Elasticsearch.scroll` api - a simple iterator that
|
:meth:`~elasticsearch.Elasticsearch.scroll` api - a simple iterator that
|
||||||
@@ -166,19 +166,30 @@ def scan(client, query=None, scroll='5m', **kwargs):
|
|||||||
:arg query: body for the :meth:`~elasticsearch.Elasticsearch.search` api
|
:arg query: body for the :meth:`~elasticsearch.Elasticsearch.search` api
|
||||||
:arg scroll: Specify how long a consistent view of the index should be
|
:arg scroll: Specify how long a consistent view of the index should be
|
||||||
maintained for scrolled search
|
maintained for scrolled search
|
||||||
|
: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.
|
||||||
|
|
||||||
Any additional keyword arguments will be passed to the initial
|
Any additional keyword arguments will be passed to the initial
|
||||||
:meth:`~elasticsearch.Elasticsearch.search` call.
|
:meth:`~elasticsearch.Elasticsearch.search` call.
|
||||||
"""
|
"""
|
||||||
# initial search to
|
if not preserve_order:
|
||||||
resp = client.search(body=query, search_type='scan', scroll=scroll, **kwargs)
|
kwargs['search_type'] = 'scan'
|
||||||
|
# initial search
|
||||||
|
resp = client.search(body=query, scroll=scroll, **kwargs)
|
||||||
|
|
||||||
scroll_id = resp.get('_scroll_id')
|
scroll_id = resp.get('_scroll_id')
|
||||||
if scroll_id is None:
|
if scroll_id is None:
|
||||||
return
|
return
|
||||||
|
|
||||||
|
first_run = True
|
||||||
while True:
|
while True:
|
||||||
resp = client.scroll(scroll_id, scroll=scroll)
|
# if we didn't set search_type to scan initial search contains data
|
||||||
|
if preserve_order and first_run:
|
||||||
|
first_run = False
|
||||||
|
else:
|
||||||
|
resp = client.scroll(scroll_id, scroll=scroll)
|
||||||
if not resp['hits']['hits']:
|
if not resp['hits']['hits']:
|
||||||
break
|
break
|
||||||
for hit in resp['hits']['hits']:
|
for hit in resp['hits']['hits']:
|
||||||
|
|||||||
@@ -138,6 +138,19 @@ class TestBulk(ElasticsearchTestCase):
|
|||||||
|
|
||||||
|
|
||||||
class TestScan(ElasticsearchTestCase):
|
class TestScan(ElasticsearchTestCase):
|
||||||
|
def test_order_can_be_preserved(self):
|
||||||
|
bulk = []
|
||||||
|
for x in range(100):
|
||||||
|
bulk.append({"index": {"_index": "test_index", "_type": "answers", "_id": x}})
|
||||||
|
bulk.append({"answer": x, "correct": x == 42})
|
||||||
|
self.client.bulk(bulk, refresh=True)
|
||||||
|
|
||||||
|
docs = list(helpers.scan(self.client, index="test_index", doc_type="answers", size=2, query={"sort": ["answer"]}, preserve_order=True))
|
||||||
|
|
||||||
|
self.assertEquals(100, len(docs))
|
||||||
|
self.assertEquals(list(map(str, range(100))), list(d['_id'] for d in docs))
|
||||||
|
self.assertEquals(list(range(100)), list(d['_source']['answer'] for d in docs))
|
||||||
|
|
||||||
def test_all_documents_are_read(self):
|
def test_all_documents_are_read(self):
|
||||||
bulk = []
|
bulk = []
|
||||||
for x in range(100):
|
for x in range(100):
|
||||||
|
|||||||
Reference in New Issue
Block a user