Add JSON export + BigQuery compatibility

es2csv.py

@@ -1,12 +1,12 @@
+import codecs
+import json
 import os
 import time
-import json
-import codecs
+from functools import wraps
+
 import elasticsearch
 import progressbar
 from backports import csv
-from functools import wraps
-
 
 FLUSH_BUFFER = 1000  # Chunk of docs to flush in temp file
 CONNECTION_TIMEOUT = 120
@@ -52,6 +52,8 @@ def __init__(self, opts):
         self.scroll_time = '30m'
 
         self.csv_headers = list(META_FIELDS) if self.opts.meta_fields else []
+        self.header_delimiter = self.opts.header_delimiter or '.'
+        self.big_query_compat = self.opts.big_query
         self.tmp_file = '{}.tmp'.format(opts.output_file)
 
     @retry(elasticsearch.exceptions.ConnectionError, tries=TIMES_TO_TRY)
@@ -171,7 +173,9 @@ def next_scroll(scroll_id):
             bar.finish()
 
     def flush_to_file(self, hit_list):
-        def to_keyvalue_pairs(source, ancestors=[], header_delimeter='.'):
+        header_delimiter = self.header_delimiter
+        big_query_compat = self.big_query_compat
+        def to_keyvalue_pairs(source, ancestors=[]):
             def is_list(arg):
                 return type(arg) is list
 
@@ -188,13 +192,19 @@ def is_dict(arg):
                 else:
                     [to_keyvalue_pairs(item, ancestors + [str(index)]) for index, item in enumerate(source)]
             else:
-                header = header_delimeter.join(ancestors)
+                header = header_delimiter.join(ancestors)
                 if header not in self.csv_headers:
                     self.csv_headers.append(header)
-                try:
-                    out[header] = '{}{}{}'.format(out[header], self.opts.delimiter, source)
-                except:
-                    out[header] = source
+                if big_query_compat == False:
+                    try:
+                        out[header] = '{}{}{}'.format(out[header], self.opts.delimiter, source)
+                    except:
+                        out[header] = source
+                else:
+                    try:
+                        out[header.replace('@', '_')] = '{}{}{}'.format(out[header], self.opts.delimiter, source)
+                    except:
+                        out[header.replace('@', '_')] = source
 
         with codecs.open(self.tmp_file, mode='a', encoding='utf-8') as tmp_file:
             for hit in hit_list:
@@ -232,6 +242,10 @@ def write_to_csv(self):
                 print('There is no docs with selected field(s): {}.'.format(','.join(self.opts.fields)))
             os.remove(self.tmp_file)
 
+    def write_to_json(self):
+        if self.num_results > 0:
+            os.rename(self.tmp_file, self.opts.output_file)
+
     def clean_scroll_ids(self):
         try:
             self.es_conn.clear_scroll(body=','.join(self.scroll_ids))

es2csv_cli.py

@@ -11,8 +11,9 @@
                  es2csv -q '*' -t dev prod -u http://login:[email protected]:6666/es/ -o ~/file.csv
                  es2csv -q '{"query": {"match_all": {}}, "filter":{"term": {"tags": "dev"}}}' -r -u http://login:[email protected]:6666/es/ -o ~/file.csv
 """
-import sys
 import argparse
+import sys
+
 import es2csv
 
 __version__ = '5.5.2'
@@ -35,6 +36,9 @@ def main():
     p.add_argument('-k', '--kibana-nested', dest='kibana_nested', action='store_true', help='Format nested fields in Kibana style.')
     p.add_argument('-r', '--raw-query', dest='raw_query', action='store_true', help='Switch query format in the Query DSL.')
     p.add_argument('-e', '--meta-fields', dest='meta_fields', action='store_true', help='Add meta-fields in output.')
+    p.add_argument('-j', '--json', dest='json', action='store_true', help='Output as line-separated JSON instead of CSV')
+    p.add_argument('-l', '--header-delimiter', dest='header_delimiter', type=str, help='Delimiter to use with JSON nested fields')
+    p.add_argument('-b', '--big-query', dest='big_query', action="store_true", help='Export with BigQuery compatibility')
     p.add_argument('--verify-certs', dest='verify_certs', action='store_true', help='Verify SSL certificates. Default is %(default)s.')
     p.add_argument('--ca-certs', dest='ca_certs', default=None, type=str, help='Location of CA bundle.')
     p.add_argument('--client-cert', dest='client_cert', default=None, type=str, help='Location of Client Auth cert.')
@@ -51,7 +55,10 @@ def main():
     es.create_connection()
     es.check_indexes()
     es.search_query()
-    es.write_to_csv()
+    if (opts.json != True):
+        es.write_to_csv()
+    else:
+        es.write_to_json()
     es.clean_scroll_ids()