From 5b1a2fd6505808f2ae1919230ea24cd758326047 Mon Sep 17 00:00:00 2001 From: Honza Kral Date: Thu, 1 Aug 2013 14:48:39 +0200 Subject: [PATCH] bulk_index helper to easily load docs into ES --- elasticsearch/helpers.py | 28 +++++++++++++++++++ .../test_server/test_helpers.py | 13 +++++++++ 2 files changed, 41 insertions(+) create mode 100644 elasticsearch/helpers.py create mode 100644 test_elasticsearch/test_server/test_helpers.py diff --git a/elasticsearch/helpers.py b/elasticsearch/helpers.py new file mode 100644 index 00000000..0a9a1720 --- /dev/null +++ b/elasticsearch/helpers.py @@ -0,0 +1,28 @@ +from itertools import islice + +def bulk_index(client, docs, chunk_size=500, **kwargs): + success, failed = [], [] + docs = iter(docs) + while True: + chunk = islice(docs, chunk_size) + bulk_actions = [] + for d in chunk: + action = {'index': {}} + for key in ('_index', '_parent', '_percolate', '_routing', + '_timestamp', '_ttl', '_type', '_version',): + if key in d: + action['index'][key] = d.pop(key) + + bulk_actions.append(action) + bulk_actions.append(d.get('_source', d)) + + if not bulk_actions: + return success, failed + + resp = client.bulk(bulk_actions, **kwargs) + + for item in resp['items']: + if item['create']['ok']: + success.append(item) + else: + failed.append(item) diff --git a/test_elasticsearch/test_server/test_helpers.py b/test_elasticsearch/test_server/test_helpers.py new file mode 100644 index 00000000..fb3fa516 --- /dev/null +++ b/test_elasticsearch/test_server/test_helpers.py @@ -0,0 +1,13 @@ +from elasticsearch import helpers + +from . import ElasticTestCase + +class TestBulkIndex(ElasticTestCase): + def test_all_documents_get_inserted(self): + docs = [{"answer": x} for x in range(100)] + success, failed = helpers.bulk_index(self.client, docs, index='test-index', doc_type='answers', refresh=True) + + self.assertEquals(len(success), len(docs)) + self.assertFalse(failed) + self.assertEquals(len(docs), self.client.count(index='test-index', doc_type='answers')['count']) +