Added a convenient scan heper API to iterate over docs
This commit is contained in:
@@ -26,3 +26,18 @@ def bulk_index(client, docs, chunk_size=500, **kwargs):
|
|||||||
success.append(item)
|
success.append(item)
|
||||||
else:
|
else:
|
||||||
failed.append(item)
|
failed.append(item)
|
||||||
|
|
||||||
|
def scan(client, query=None, scroll='5m', **kwargs):
|
||||||
|
# initial search to
|
||||||
|
resp = client.search(body=query, search_type='scan', scroll=scroll, **kwargs)
|
||||||
|
|
||||||
|
scroll_id = resp['_scroll_id']
|
||||||
|
|
||||||
|
while True:
|
||||||
|
resp = client.scroll(scroll_id, scroll=scroll)
|
||||||
|
if not resp['hits']['hits']:
|
||||||
|
break
|
||||||
|
for hit in resp['hits']['hits']:
|
||||||
|
yield hit
|
||||||
|
scroll_id = resp['_scroll_id']
|
||||||
|
|
||||||
|
|||||||
@@ -11,3 +11,16 @@ class TestBulkIndex(ElasticTestCase):
|
|||||||
self.assertFalse(failed)
|
self.assertFalse(failed)
|
||||||
self.assertEquals(len(docs), self.client.count(index='test-index', doc_type='answers')['count'])
|
self.assertEquals(len(docs), self.client.count(index='test-index', doc_type='answers')['count'])
|
||||||
|
|
||||||
|
class TestScan(ElasticTestCase):
|
||||||
|
def test_all_documents_are_read(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))
|
||||||
|
|
||||||
|
self.assertEquals(100, len(docs))
|
||||||
|
self.assertEquals(set(map(str, range(100))), set(d['_id'] for d in docs))
|
||||||
|
self.assertEquals(set(range(100)), set(d['_source']['answer'] for d in docs))
|
||||||
|
|||||||
Reference in New Issue
Block a user