From c19936fba153410dd421694be966149ec6ab57c5 Mon Sep 17 00:00:00 2001 From: James Tanner Date: Thu, 6 Feb 2014 20:22:45 -0500 Subject: [PATCH] Add exact_count and count_tag to the ec2 module. --- cloud/ec2 | 118 ++++++++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 115 insertions(+), 3 deletions(-) diff --git a/cloud/ec2 b/cloud/ec2 index c28c4b51576..f8e48948c03 100644 --- a/cloud/ec2 +++ b/cloud/ec2 @@ -376,6 +376,7 @@ local_action: import sys import time +from ast import literal_eval try: import boto.ec2 @@ -385,6 +386,68 @@ except ImportError: print "failed=True msg='boto required for this module'" sys.exit(1) +def find_running_instances_by_count_tag(module, ec2, count_tag): + + # get reservations for instances that match tag(s) and are running + reservations = get_reservations(module, ec2, tags=count_tag, state="running") + + instances = [] + for res in reservations: + if hasattr(res, 'instances'): + for inst in res.instances: + instances.append(inst) + + return reservations, instances + + +def _set_none_to_blank(dictionary): + result = dictionary + for k in result.iterkeys(): + if type(result[k]) == dict: + result[k] = _set_non_to_blank(result[k]) + elif not result[k]: + result[k] = "" + return result + + +def get_reservations(module, ec2, tags=None, state=None): + + # TODO: filters do not work with tags that have underscores + filters = dict() + + if tags is not None: + + if type(tags) is str: + try: + tags = literal_eval(tags) + except: + pass + + # if string, we only care that a tag of that name exists + if type(tags) is str: + filters.update({"tag-key": tags}) + + # if list, append each item to filters + if type(tags) is list: + for x in tags: + if type(x) is dict: + x = _set_none_to_blank(x) + filters.update(dict(("tag:"+tn, tv) for (tn,tv) in x.iteritems())) + else: + filters.update({"tag-key": x}) + + # if dict, add the key and value to the filter + if type(tags) is dict: + tags = _set_none_to_blank(tags) + filters.update(dict(("tag:"+tn, tv) for (tn,tv) in tags.iteritems())) + + if state: + # http://stackoverflow.com/questions/437511/what-are-the-valid-instancestates-for-the-amazon-ec2-api + filters.update({'instance-state-name': state}) + + results = ec2.get_all_instances(filters=filters) + + return results def get_instance_info(inst): """ @@ -473,7 +536,45 @@ def create_block_device(module, ec2, volume): delete_on_termination=volume.get('delete_on_termination', False), iops=volume.get('iops')) -def create_instances(module, ec2): + +def enforce_count(module, ec2): + + exact_count = module.params.get('exact_count') + count_tag = module.params.get('count_tag') + + reservations, instances = find_running_instances_by_count_tag(module, ec2, count_tag) + + changed = None + checkmode = False + instance_dict_array = None + changed_instance_ids = None + + if len(instances) == exact_count: + changed = False + elif len(instances) < exact_count: + changed = True + to_create = exact_count - len(instances) + if not checkmode: + (instance_dict_array, changed_instance_ids, changed) \ + = create_instances(module, ec2, override_count=to_create) + elif len(instances) > exact_count: + changed = True + to_remove = len(instances) - exact_count + if not checkmode: + all_instance_ids = sorted([ x.id for x in instances ]) + remove_ids = all_instance_ids[0:to_remove] + (changed, instance_dict_array, changed_instance_ids) \ + = terminate_instances(module, ec2, remove_ids) + terminated_list = [] + for inst in instance_dict_array: + inst['state'] = "terminated" + terminated_list.append(inst) + instance_dict_array = terminated_list + + return (instance_dict_array, changed_instance_ids, changed) + + +def create_instances(module, ec2, override_count=None): """ Creates new instances @@ -492,7 +593,10 @@ def create_instances(module, ec2): zone = module.params.get('zone') instance_type = module.params.get('instance_type') image = module.params.get('image') - count = module.params.get('count') + if override_count: + count = override_count + else: + count = module.params.get('count') monitoring = module.params.get('monitoring') kernel = module.params.get('kernel') ramdisk = module.params.get('ramdisk') @@ -506,6 +610,8 @@ def create_instances(module, ec2): private_ip = module.params.get('private_ip') instance_profile_name = module.params.get('instance_profile_name') volumes = module.params.get('volumes') + exact_count = module.params.get('exact_count') + count_tag = module.params.get('count_tag') # group_id and group_name are exclusive of each other if group_id and group_name: @@ -832,6 +938,8 @@ def main(): instance_profile_name = dict(), instance_ids = dict(type='list'), state = dict(default='present'), + exact_count = dict(type='int'), + count_tag = dict(), volumes = dict(type='list'), ) ) @@ -857,7 +965,11 @@ def main(): # Changed is always set to true when provisioning new instances if not module.params.get('image'): module.fail_json(msg='image parameter is required for new instance') - (instance_dict_array, new_instance_ids, changed) = create_instances(module, ec2) + + if module.params.get('exact_count'): + (instance_dict_array, new_instance_ids, changed) = enforce_count(module, ec2) + else: + (instance_dict_array, new_instance_ids, changed) = create_instances(module, ec2) module.exit_json(changed=changed, instance_ids=new_instance_ids, instances=instance_dict_array)