append registered hostname to policy list

* The Instance Group list of instances was getting over-written with
every call to the register_instance management command.
* This changeset appends --hostnames to the Instance Group policy list.

awx/main/management/commands/register_queue.py

@@ -1,6 +1,7 @@
 # Copyright (c) 2017 Ansible Tower by Red Hat
 # All Rights Reserved.
 import sys
+import six
 
 from awx.main.utils.pglock import advisory_lock
 from awx.main.models import Instance, InstanceGroup
@@ -8,6 +9,13 @@ from awx.main.models import Instance, InstanceGroup
 from django.core.management.base import BaseCommand, CommandError
 
 
+class InstanceNotFound(Exception):
+    def __init__(self, message, changed, *args, **kwargs):
+        self.message = message
+        self.changed = changed
+        super(InstanceNotFound, self).__init__(*args, **kwargs)
+
+
 class Command(BaseCommand):
 
     def add_arguments(self, parser):
@@ -22,51 +30,95 @@ class Command(BaseCommand):
         parser.add_argument('--instance_minimum', dest='instance_minimum', type=int, default=0,
                             help='The minimum number of instance that will be retained for this group from available instances')
 
+
+    def get_create_update_instance_group(self, queuename, instance_percent, instance_min):
+        ig = InstanceGroup.objects.filter(name=queuename)
+        created = False
+        changed = False
+
+        (ig, created) = InstanceGroup.objects.get_or_create(name=queuename)
+        if ig.policy_instance_percentage != instance_percent:
+            ig.policy_instance_percentage = instance_percent
+            changed = True
+        if ig.policy_instance_minimum != instance_min:
+            ig.policy_instance_minimum = instance_min
+            changed = True
+
+        return (ig, created, changed)
+
+    def update_instance_group_controller(self, ig, controller):
+        changed = False
+        control_ig = None
+
+        if controller:
+            control_ig = InstanceGroup.objects.filter(name=controller).first()
+
+        if control_ig and ig.controller_id != control_ig.pk:
+            ig.controller = control_ig
+            ig.save()
+            changed = True
+
+        return (control_ig, changed)
+
+    def add_instances_to_group(self, ig, hostname_list):
+        changed = False
+
+        instance_list_unique = set([x.strip() for x in hostname_list if x])
+        instances = []
+        for inst_name in instance_list_unique:
+            instance = Instance.objects.filter(hostname=inst_name)
+            if instance.exists():
+                instances.append(instance[0])
+            else:
+                raise InstanceNotFound(six.text_type("Instance does not exist: {}").format(inst_name), changed)
+
+        ig.instances = instances
+
+        instance_list_before = set(ig.policy_instance_list)
+        instance_list_after = set(instance_list_unique)
+        if len(instance_list_before) != len(instance_list_after) or \
+                len(set(instance_list_before) - set(instance_list_after)) != 0:
+            changed = True
+
+        ig.policy_instance_list = list(instance_list_unique)
+        ig.save()
+        return (instances, changed)
+
     def handle(self, **options):
+        instance_not_found_err = None
         queuename = options.get('queuename')
         if not queuename:
             raise CommandError("Specify `--queuename` to use this command.")
-        changed = False
+        ctrl = options.get('controller')
+        inst_per = options.get('instance_percent')
+        inst_min = options.get('instance_minimum')
+        hostname_list = []
+        if options.get('hostnames'):
+            hostname_list = options.get('hostnames').split(",")
+
         with advisory_lock('instance_group_registration_%s' % queuename):
-            ig = InstanceGroup.objects.filter(name=queuename)
+            (ig, created, changed) = self.get_create_update_instance_group(queuename, inst_per, inst_min)
-            control_ig = None
+            if created:
-            if options.get('controller'):
+                print(six.text_type("Creating instance group {}".format(ig.name)))
-                control_ig = InstanceGroup.objects.filter(name=options.get('controller')).first()
+            elif not created:
-            if ig.exists():
+                print(six.text_type("Instance Group already registered {}").format(ig.name))
-                print("Instance Group already registered {}".format(ig[0].name))
+
-                ig = ig[0]
+            if ctrl:
-                if control_ig and ig.controller_id != control_ig.pk:
+                (ig_ctrl, changed) = self.update_instance_group_controller(ig, ctrl)
-                    ig.controller = control_ig
+                if changed:
-                    ig.save()
+                    print(six.text_type("Set controller group {} on {}.").format(ctrl, queuename))
-                    print("Set controller group {} on {}.".format(control_ig.name, ig.name))
+
-                    changed = True
+            try:
-            else:
+                (instances, changed) = self.add_instances_to_group(ig, hostname_list)
-                print("Creating instance group {}".format(queuename))
+                for i in instances:
-                ig = InstanceGroup(name=queuename,
+                    print(six.text_type("Added instance {} to {}").format(i.hostname, ig.name))
-                                   policy_instance_percentage=options.get('instance_percent'),
+            except InstanceNotFound as e:
-                                   policy_instance_minimum=options.get('instance_minimum'))
+                instance_not_found_err = e
-                if control_ig:
+
-                    ig.controller = control_ig
+        if changed:
-                ig.save()
+            print('(changed: True)')
-                changed = True
+
-            hostname_list = []
+        if instance_not_found_err:
-            if options.get('hostnames'):
+            print(instance_not_found_err.message)
-                hostname_list = options.get('hostnames').split(",")
+            sys.exit(1)
-            instance_list = [x.strip() for x in hostname_list if x]
+
-            for inst_name in instance_list:
-                instance = Instance.objects.filter(hostname=inst_name)
-                if instance.exists() and instance[0] not in ig.instances.all():
-                    ig.instances.add(instance[0])
-                    print("Added instance {} to {}".format(instance[0].hostname, ig.name))
-                    changed = True
-                elif not instance.exists():
-                    print("Instance does not exist: {}".format(inst_name))
-                    if changed:
-                        print('(changed: True)')
-                    sys.exit(1)
-                else:
-                    print("Instance already registered {}".format(instance[0].hostname))
-            ig.policy_instance_list = instance_list
-            ig.save()
-            if changed:
-                print('(changed: True)')