diff --git a/elasticsearch/helpers.py b/elasticsearch/helpers.py index 0a9a1720..505c8cdb 100644 --- a/elasticsearch/helpers.py +++ b/elasticsearch/helpers.py @@ -26,3 +26,18 @@ def bulk_index(client, docs, chunk_size=500, **kwargs): success.append(item) else: 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'] + diff --git a/test_elasticsearch/test_server/test_helpers.py b/test_elasticsearch/test_server/test_helpers.py index fb3fa516..951045e9 100644 --- a/test_elasticsearch/test_server/test_helpers.py +++ b/test_elasticsearch/test_server/test_helpers.py @@ -11,3 +11,16 @@ class TestBulkIndex(ElasticTestCase): self.assertFalse(failed) 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))