more cmd line options and replace print with logging

Author: oliver006

src/index_emails.py

@@ -7,20 +7,20 @@
 import email.utils
 import mailbox
 import email
-
+import logging
 
 http_client = HTTPClient()
 
 DEFAULT_BATCH_SIZE = 500
-ES_URL = "http://localhost:9200/gmail"
-
+DEFAULT_ES_URL = "http://localhost:9200"
+DEFAULT_INDEX_NAME = "gmail"
 
 def delete_index():
     try:
-        body = {"refresh": True}
-        request = HTTPRequest(ES_URL, method="DELETE", body=json.dumps(body), request_timeout=240)
+        url = "%s/%s?refresh=true" % (tornado.options.options.es_url, tornado.options.options.index_name)
+        request = HTTPRequest(url, method="DELETE", request_timeout=240)
         response = http_client.fetch(request)
-        print 'Delete index done   %s' % response.body
+        logging.info('Delete index done   %s' % response.body)
     except:
         pass
 
@@ -50,28 +50,31 @@ def create_index():
     }
 
     body = json.dumps(schema)
-    request = HTTPRequest(ES_URL, method="PUT", body=body, request_timeout=240)
-    response = http_client.fetch(request)
-    print 'Create index done   %s' % response.body
-
-total_uploaded = 0
+    url = "%s/%s" % (tornado.options.options.es_url, tornado.options.options.index_name)
+    try:
+        request = HTTPRequest(url, method="PUT", body=body, request_timeout=240)
+        response = http_client.fetch(request)
+        logging.info('Create index done   %s' % response.body)
+    except:
+        pass
 
 
+total_uploaded = 0
 def upload_batch(upload_data):
     upload_data_txt = ""
     for item in upload_data:
         cmd = {'index': {'_index': 'gmail', '_type': 'email', '_id': item['message-id']}}
         upload_data_txt += json.dumps(cmd) + "\n"
         upload_data_txt += json.dumps(item) + "\n"
 
-    request = HTTPRequest("http://localhost:9200/_bulk", method="POST", body=upload_data_txt, request_timeout=240)
+    request = HTTPRequest(tornado.options.options.es_url + "/_bulk", method="POST", body=upload_data_txt, request_timeout=240)
     response = http_client.fetch(request)
     result = json.loads(response.body)
 
     global total_uploaded
     total_uploaded += len(upload_data)
     res_txt = "OK" if not result['errors'] else "FAILED"
-    print "Upload: %s - upload took: %4dms, total messages uploaded: %6d" % (res_txt, result['took'], total_uploaded)
+    logging.info("Upload: %s - upload took: %4dms, total messages uploaded: %6d" % (res_txt, result['took'], total_uploaded))
 
 
 def normalize_email(email_in):
@@ -122,13 +125,15 @@ def load_from_file():
 
     if tornado.options.options.init:
         delete_index()
-        create_index()
+    create_index()
+
 
     if tornado.options.options.skip:
-        print "Skipping first %d messages from mbox file" % tornado.options.options.skip
+        logging.info("Skipping first %d messages from mbox file" % tornado.options.options.skip)
 
     count = 0
     upload_data = list()
+    logging.info("Starting import from file %s" % tornado.options.options.infile)
     mbox = mailbox.UnixMailbox(open(tornado.options.options.infile, 'rb'), email.message_from_file)
     for msg in mbox:
         count += 1
@@ -145,34 +150,28 @@ def load_from_file():
     if upload_data:
         upload_batch(upload_data)
 
-    print "Done - total count %d" % count
+    logging.info("Import done - total count %d" % count)
 
 
 if __name__ == '__main__':
 
-    tornado.options.define(
-        "infile",
-        type=str,
-        default=None,
-        help="The mbox input file")
-
-    tornado.options.define(
-        "init",
-        type=bool,
-        default=False,
-        help="Delete and re-initialize the Elasticsearch index")
-
-    tornado.options.define(
-        "batch_size",
-        type=int,
-        default=DEFAULT_BATCH_SIZE,
-        help="Elasticsearch bulk index batch size")
-
-    tornado.options.define(
-        "skip",
-        type=int,
-        default=0,
-        help="Number of messages to skip from the mbox file")
+    tornado.options.define("es_url", type=str, default=DEFAULT_ES_URL,
+                           help="URL of your Elasticsearch node")
+
+    tornado.options.define("index_name", type=str, default=DEFAULT_INDEX_NAME,
+                           help="Name of the index to store your messages")
+
+    tornado.options.define("infile", type=str, default=None,
+                           help="The mbox input file")
+
+    tornado.options.define("init", type=bool, default=False,
+                           help="Force deleting and re-initializing the Elasticsearch index")
+
+    tornado.options.define("batch_size", type=int, default=DEFAULT_BATCH_SIZE,
+                           help="Elasticsearch bulk index batch size")
+
+    tornado.options.define("skip", type=int, default=0,
+                           help="Number of messages to skip from the mbox file")
 
     tornado.options.parse_command_line()