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,6 +469,94 @@ def create_ami(instance_id, name, description):
 
     return image_id
 
+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))
+    print
+    for k, v in ec2_args.iteritems():
+        if k != 'user_data':
+            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)
+    inst = res.instances[0]
+    instance_id = inst.id
+
+    print "{:<40}".format("Waiting for running status:"),
+    status_start = time.time()
+    for _ in xrange(EC2_RUN_TIMEOUT):
+        res = ec2.get_all_instances(instance_ids=[instance_id])
+        if res[0].instances[0].state == 'running':
+            status_delta = time.time() - status_start
+            run_summary.append(('EC2 Launch', status_delta))
+            print "[ OK ] {:0>2.0f}:{:0>2.0f}".format(
+                status_delta / 60,
+                status_delta % 60)
+            break
+        else:
+            time.sleep(1)
+    else:
+        raise Exception("Timeout waiting for running status: {} ".format(
+            instance_id))
+
+    print "{:<40}".format("Waiting for system status:"),
+    system_start = time.time()
+    for _ in xrange(EC2_STATUS_TIMEOUT):
+        status = ec2.get_all_instance_status(inst.id)
+        if status[0].system_status.status == u'ok':
+            system_delta = time.time() - system_start
+            run_summary.append(('EC2 Status Checks', system_delta))
+            print "[ OK ] {:0>2.0f}:{:0>2.0f}".format(
+                system_delta / 60,
+                system_delta % 60)
+            break
+        else:
+            time.sleep(1)
+    else:
+        raise Exception("Timeout waiting for status checks: {} ".format(
+            instance_id))
+    user_start = time.time()
+
+    print
+    print "{:<40}".format(
+        "Waiting for user-data, polling sqs for Ansible events:")
+
+    (ansible_delta, task_report) = poll_sqs_ansible()
+    user_pre_ansible = time.time() - user_start - ansible_delta
+    run_summary.append(('Ansible run', ansible_delta))
+    print
+    print "{} longest Ansible tasks (seconds):".format(NUM_TASKS)
+    for task in sorted(
+            task_report, reverse=True,
+            key=lambda k: k['DELTA'])[:NUM_TASKS]:
+        print "{:0>3.0f} {}".format(task['DELTA'], task['TASK'])
+        print "  - {}".format(task['INVOCATION'])
+    print
+
+    print "{:<40}".format("Creating AMI:"),
+    ami_start = time.time()
+    ami = create_ami(instance_id, run_id, run_id)
+    ami_delta = time.time() - ami_start
+    print "[ OK ] {:0>2.0f}:{:0>2.0f}".format(
+        ami_delta / 60,
+        ami_delta % 60)
+    run_summary.append(('AMI Build', ami_delta))
+    total_time = time.time() - start_time
+    all_stages = sum(run[1] for run in run_summary)
+    if total_time - all_stages > 0:
+        run_summary.append(('Other', total_time - all_stages))
+    run_summary.append(('Total', total_time))
+
+    return run_summary, ami
 
 if __name__ == '__main__':
 
@@ -399,8 +569,10 @@ if __name__ == '__main__':
     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
@@ -420,6 +592,11 @@ if __name__ == '__main__':
         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
@@ -429,96 +606,31 @@ if __name__ == '__main__':
 
         ec2_args = create_instance_args()
 
-        print "{:<40}".format(
-            "Creating SQS queue and launching instance for {}:".format(run_id))
-        print
-        for k, v in ec2_args.iteritems():
-            if k != 'user_data':
-                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)
-        inst = res.instances[0]
-        instance_id = inst.id
-
-        print "{:<40}".format("Waiting for running status:"),
-        status_start = time.time()
-        for _ in xrange(EC2_RUN_TIMEOUT):
-            res = ec2.get_all_instances(instance_ids=[instance_id])
-            if res[0].instances[0].state == 'running':
-                status_delta = time.time() - status_start
-                run_summary.append(('EC2 Launch', status_delta))
-                print "[ OK ] {:0>2.0f}:{:0>2.0f}".format(
-                    status_delta / 60,
-                    status_delta % 60)
-                break
-            else:
-                time.sleep(1)
+        if args.noop:
+            print "Would have created sqs_queue with id: {}\nec2_args:".format(run_id)
+            pprint(ec2_args)
+            ami = "ami-00000"
         else:
-            raise Exception("Timeout waiting for running status: {} ".format(
-                instance_id))
-
-        print "{:<40}".format("Waiting for system status:"),
-        system_start = time.time()
-        for _ in xrange(EC2_STATUS_TIMEOUT):
-            status = ec2.get_all_instance_status(inst.id)
-            if status[0].system_status.status == u'ok':
-                system_delta = time.time() - system_start
-                run_summary.append(('EC2 Status Checks', system_delta))
-                print "[ OK ] {:0>2.0f}:{:0>2.0f}".format(
-                    system_delta / 60,
-                    system_delta % 60)
-                break
-            else:
-                time.sleep(1)
-        else:
-            raise Exception("Timeout waiting for status checks: {} ".format(
-                instance_id))
-        user_start = time.time()
-
-        print
-        print "{:<40}".format(
-            "Waiting for user-data, polling sqs for Ansible events:")
-
-        (ansible_delta, task_report) = poll_sqs_ansible()
-        user_pre_ansible = time.time() - user_start - ansible_delta
-        run_summary.append(('Ansible run', ansible_delta))
-        print
-        print "{} longest Ansible tasks (seconds):".format(NUM_TASKS)
-        for task in sorted(
-                task_report, reverse=True,
-                key=lambda k: k['DELTA'])[:NUM_TASKS]:
-            print "{:0>3.0f} {}".format(task['DELTA'], task['TASK'])
-            print "  - {}".format(task['INVOCATION'])
-        print
-
-        print "{:<40}".format("Creating AMI:"),
-        ami_start = time.time()
-        ami = create_ami(instance_id, run_id, run_id)
-        ami_delta = time.time() - ami_start
-        print "[ OK ] {:0>2.0f}:{:0>2.0f}".format(
-            ami_delta / 60,
-            ami_delta % 60)
-        run_summary.append(('AMI Build', ami_delta))
-        total_time = time.time() - start_time
-        all_stages = sum(run[1] for run in run_summary)
-        if total_time - all_stages > 0:
-            run_summary.append(('Other', total_time - all_stages))
-        run_summary.append(('Total', total_time))
-
-        print
-        print "Summary:\n"
-
-        for run in run_summary:
-            print "{:<30} {:0>2.0f}:{:0>5.2f}".format(
-                run[0], run[1] / 60, run[1] % 60)
-        print "AMI: {}".format(ami)
+            run_summary, ami = launch_and_configure(ec2_args)
+            print
+            print "Summary:\n"
+
+            for run in run_summary:
+                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)