Merge pull request #5901 from jctanner/ec2_tag_exact_count_2

Add exact_count and count_tag to the ec2 module.

cloud/ec2

@@ -198,6 +198,21 @@ options:
     required: false
     default: null
     aliases: []
+  exact_count:
+    version_added: "1.5"
+    description:
+      - An integer value which indicates how many instances that match the 'count_tag' parameter should be running. Instances are either created or terminated based on this value. 
+    required: false
+    default: null
+    aliases: []
+  count_tag:
+    version_added: "1.5"
+    description:
+      - Used with 'exact_count' to determine how many nodes based on a specific tag criteria should be running.  This can be expressed in multiple ways and is shown in the EXAMPLES section.  For instance, one can request 25 servers that are tagged with "class=webserver".  
+    required: false
+    default: null
+    aliases: []
+       
 
 requirements: [ "boto" ]
 author: Seth Vidal, Tim Gerla, Lester Wade
@@ -227,8 +242,9 @@ EXAMPLES = '''
     wait: yes
     wait_timeout: 500
     count: 5
-    instance_tags: '{"db":"postgres"}'
+    instance_tags: 
-    monitoring=yes
+       db: postgres
+    monitoring: yes
 
 # Single instance with additional IOPS volume from snapshot
 local_action:
@@ -245,7 +261,7 @@ local_action:
       device_type: io1
       iops: 1000
       volume_size: 100
-    monitoring=yes
+    monitoring: yes
 
 # Multiple groups example
 local_action:
@@ -257,8 +273,9 @@ local_action:
     wait: yes
     wait_timeout: 500
     count: 5
-    instance_tags: '{"db":"postgres"}'
+    instance_tags: 
-    monitoring=yes
+        db: postgres
+    monitoring: yes
 
 # Multiple instances with additional volume from snapshot
 local_action:
@@ -274,7 +291,7 @@ local_action:
     - device_name: /dev/sdb
       snapshot: snap-abcdef12
       volume_size: 10
-    monitoring=yes
+    monitoring: yes
 
 # VPC example
 - local_action:
@@ -372,10 +389,70 @@ local_action:
       region: '{{ region }}'
       state: stopped
       wait: True
+
+#
+# Enforce that 5 instances with a tag "foo" are running
+#
+
+- local_action:
+    module: ec2
+    keypair: mykey
+    instance_type: c1.medium
+    image: emi-40603AD1
+    wait: yes
+    group: webserver
+    instance_tags:
+        foo: bar
+    exact_count: 5
+    count_tag: foo
+
+#
+# Enforce that 5 running instances named "database" with a "dbtype" of "postgres"
+#
+
+- local_action:
+    module: ec2
+    keypair: mykey
+    instance_type: c1.medium
+    image: emi-40603AD1
+    wait: yes
+    group: webserver
+    instance_tags: 
+        Name: database
+        dbtype: postgres
+    exact_count: 5
+    count_tag: 
+        Name: database
+        dbtype: postgres
+
+#
+# count_tag complex argument examples
+#
+
+    # instances with tag foo
+    count_tag:
+        foo:
+
+    # instances with tag foo=bar
+    count_tag:
+        foo: bar
+
+    # instances with tags foo=bar & baz
+    count_tag:
+        foo: bar
+        baz:
+
+    # instances with tags foo & bar & baz=bang
+    count_tag:
+        - foo
+        - bar
+        - baz: bang
+
 '''
 
 import sys
 import time
+from ast import literal_eval
 
 try:
     import boto.ec2
@@ -385,6 +462,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 +612,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 +669,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 +686,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 +1014,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 +1041,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)