Add keys flag (#68)

* Add keys flag

elasticsearch_loader/__init__.py

@@ -86,6 +86,7 @@ def log(sevirity, msg):
 @click.option('--index-settings-file', type=click.File('rb'), help='Specify path to json file containing index mapping and settings, creates index if missing')
 @click.option('--timeout', type=float, help='Specify request timeout in seconds for Elasticsearch client', default=10)
 @click.option('--encoding', type=str, help='Specify content encoding for input files', default='utf-8')
+@click.option('--keys', type=str, help='Comma separated keys to pick from each document', default='', callback=lambda c, p, v: [x for x in v.split(',') if x])
 @click.pass_context
 def cli(ctx, **opts):
     ctx.obj = opts

elasticsearch_loader/iter.py

@@ -14,9 +14,13 @@ def grouper(iterable, n, fillvalue=None):
 
 def bulk_builder(bulk, config):
     for item in filter(None, bulk):
+        source = item
+        if config['keys']:
+            source = {x: y for x, y in item.items() if x in config['keys']}
+
         body = {'_index': config['index'],
                 '_type': config['type'],
-                '_source': item}
+                '_source': source}
 
         if config['id_field']:
             body['_id'] = item[config['id_field']]
@@ -28,7 +32,7 @@ def bulk_builder(bulk, config):
         if config['update']:
             # default _op_type is 'index', which will overwrites existing doc
             body['_op_type'] = 'update'
-            body['doc'] = item
+            body['doc'] = source
             del body['_source']
 
         yield body

test/test_json.py

@@ -1,6 +1,7 @@
+import mock
 from click.testing import CliRunner
+
 from elasticsearch_loader import cli
-import mock
 
 
 def invoke(*args, **kwargs):