adding mongo updates to the abbey script

util/vpc-tools/abbey.py

-#!/usr/bin/env python -u
 
+#!/usr/bin/env python -u
 import sys
 from argparse import ArgumentParser
 import time
@@ -14,11 +14,73 @@ except ImportError:
     print "boto required for script"
     sys.exit(1)
 
+from pymongo import MongoClient
+from pymongo.errors import ConnectionFailure
+from pprint import pprint
+
 AMI_TIMEOUT = 600  # time to wait for AMIs to complete
 EC2_RUN_TIMEOUT = 180  # time to wait for ec2 state transition
 EC2_STATUS_TIMEOUT = 300  # time to wait for ec2 system status checks
 NUM_TASKS = 5  # number of tasks for time summary report
 
+class MongoConnection:
+
+    def __init__(self):
+        try:
+            mongo = MongoClient(host=args.mongo_host, port=args.mongo_port)
+            mongo_db = getattr(mongo, args.mongo_db)
+            if args.mongo_ami_collection not in mongo_db.collection_names():
+                mongo_db.create_collection(args.mongo_ami_collection)
+            if args.mongo_deployment_collection not in mongo_db.collection_names():
+                mongo_db.create_collection(args.mongo_deployment_collection)
+            self.mongo_ami = getattr(mongo_db, args.mongo_ami_collection)
+            self.mongo_deployment = getattr(mongo_db, args.mongo_deployment_collection)
+
+            self.query = {
+                'play': args.play,
+                'env': args.environment,
+                'deployment': args.deployment,
+                'configuration_ref': args.configuration_version,
+                'configuration_secure_ref': args.configuration_secure_version,
+                'vars': extra_vars,
+            }
+
+        except ConnectionFailure:
+            print "Unable to connect to the mongo database specified"
+            sys.exit(1)
+
+    def update_ami(self, status, ami=None):
+        """
+        Creates a new document in the AMI
+        or updates an existing one with a status
+        """
+
+
+        update = self.query.copy()
+        update['status'] = status
+        if ami:
+            update['ami'] = ami
+        self.mongo_ami.update(self.query, update, True)
+
+    def update_deployment(self, ami):
+        """
+        Adds the built AMI to the deployment
+        collection
+        """
+        query = {
+            '_id': args.jenkins_build,
+            'plays': {
+                args.play: {
+                    'amis': {},
+                },
+            },
+        }
+        update = query.copy()
+        pprint(update)
+        update['plays'][args.play]['amis'][args.environment] = ami
+        self.mongo_deployment.update(query, update, True)
+
+
 
 class Unbuffered:
     """
@@ -101,6 +163,25 @@ def parse_args():
                         default=5,
                         help="How long to delay message display from sqs "
                              "to ensure ordering")
+    parser.add_argument("--mongo-host", required=False,
+                        default=None,
+                        help="Mongo host that contains the AMI collection")
+    parser.add_argument("--mongo-pass", required=False,
+                        default=None,
+                        help="Mongo password")
+    parser.add_argument("--mongo-port", required=False,
+                        default=27017,
+                        help="Mongo port")
+    parser.add_argument("--mongo-db", required=False,
+                        default="test",
+                        help="Mongo database")
+    parser.add_argument("--mongo-ami-collection", required=False,
+                        default="ami",
+                        help="Mongo ami collection")
+    parser.add_argument("--mongo-deployment-collection", required=False,
+                        default="deployment",
+                        help="Mongo deployment collection")
+
     return parser.parse_args()
 
 
@@ -254,6 +335,7 @@ rm -rf $base_dir
         'instance_type': args.instance_type,
         'instance_profile_name': args.role_name,
         'user_data': user_data,
+
     }
 
     return ec2_args
@@ -387,47 +469,12 @@ def create_ami(instance_id, name, description):
 
     return image_id
 
-
-if __name__ == '__main__':
-
-    args = parse_args()
-
-    run_summary = []
-
-    start_time = time.time()
-
-    if args.vars:
-        with open(args.vars) as f:
-            extra_vars_yml = f.read()
-    else:
-        extra_vars_yml = "---\n"
-
-    if args.secure_vars:
-        secure_vars = args.secure_vars
-    else:
-        secure_vars = "../../../configuration-secure/" \
-                      "ansible/vars/{}/{}.yml".format(
-                      args.deployment, args.environment)
-    if args.stack_name:
-        stack_name = args.stack_name
-    else:
-        stack_name = "{}-{}".format(args.environment, args.deployment)
-
-    try:
-        sqs = boto.sqs.connect_to_region(args.region)
-        ec2 = boto.ec2.connect_to_region(args.region)
-    except NoAuthHandlerFound:
-        print 'You must be able to connect to sqs and ec2 to use this script'
-        sys.exit(1)
-
-    try:
-        sqs_queue = None
-        instance_id = None
-
-        run_id = "abbey-{}-{}-{}".format(
-            args.environment, args.deployment, int(time.time() * 100))
-
-        ec2_args = create_instance_args()
+def launch_and_configure(ec2_args):
+    """
+    Creates an sqs queue, launches an ec2 instance,
+    configures it and creates an AMI. Polls
+    SQS for updates
+    """
 
     print "{:<40}".format(
         "Creating SQS queue and launching instance for {}:".format(run_id))
@@ -437,6 +484,7 @@ if __name__ == '__main__':
             print "    {:<25}{}".format(k, v)
     print
 
+
     sqs_queue = sqs.create_queue(run_id)
     sqs_queue.set_message_class(RawMessage)
     res = ec2.run_instances(**ec2_args)
@@ -508,6 +556,62 @@ if __name__ == '__main__':
         run_summary.append(('Other', total_time - all_stages))
     run_summary.append(('Total', total_time))
 
+    return run_summary, ami
+
+if __name__ == '__main__':
+
+    args = parse_args()
+
+    run_summary = []
+
+    start_time = time.time()
+
+    if args.vars:
+        with open(args.vars) as f:
+            extra_vars_yml = f.read()
+            extra_vars = yaml.load(f.read)
+    else:
+        extra_vars_yml = "---\n"
+        extra_vars = {}
+
+    if args.secure_vars:
+        secure_vars = args.secure_vars
+    else:
+        secure_vars = "../../../configuration-secure/" \
+                      "ansible/vars/{}/{}.yml".format(
+                      args.deployment, args.environment)
+    if args.stack_name:
+        stack_name = args.stack_name
+    else:
+        stack_name = "{}-{}".format(args.environment, args.deployment)
+
+    try:
+        sqs = boto.sqs.connect_to_region(args.region)
+        ec2 = boto.ec2.connect_to_region(args.region)
+    except NoAuthHandlerFound:
+        print 'You must be able to connect to sqs and ec2 to use this script'
+        sys.exit(1)
+
+    if args.mongo_host:
+        mongo_con = MongoConnection()
+        mongo_con.update_ami(status='Generating')
+
+
+    try:
+        sqs_queue = None
+        instance_id = None
+
+        run_id = "abbey-{}-{}-{}".format(
+            args.environment, args.deployment, int(time.time() * 100))
+
+        ec2_args = create_instance_args()
+
+        if args.noop:
+            print "Would have created sqs_queue with id: {}\nec2_args:".format(run_id)
+            pprint(ec2_args)
+            ami = "ami-00000"
+        else:
+            run_summary, ami = launch_and_configure(ec2_args)
             print
             print "Summary:\n"
 
@@ -515,10 +619,18 @@ if __name__ == '__main__':
                 print "{:<30} {:0>2.0f}:{:0>5.2f}".format(
                     run[0], run[1] / 60, run[1] % 60)
             print "AMI: {}".format(ami)
+        if args.mongo_host:
+            mongo_con.update_ami(status='Completed', ami=ami)
+            mongo_con.update_deployment(ami)
+
+    except Exception as e:
+        if args.mongo_host:
+            mongo_con.update_ami(status='Error: {}'.format(e.message))
+        raise
 
     finally:
         print
-        if not args.no_cleanup:
+        if not args.no_cleanup and not args.noop:
             if sqs_queue:
                 print "Cleaning up - Removing SQS queue - {}".format(run_id)
                 sqs.delete_queue(sqs_queue)