Enforce that a drone may appear in at most one drone set. Also hide drones that
are already used from the drone set filter_horizontal widget.

Signed-off-by: James Ren <jamesren@google.com>



git-svn-id: http://test.kernel.org/svn/autotest/trunk@4580 592f7852-d20e-0410-864c-8624ca9c26a4
diff --git a/frontend/afe/admin.py b/frontend/afe/admin.py
index 0b92670..b6ba64a 100644
--- a/frontend/afe/admin.py
+++ b/frontend/afe/admin.py
@@ -136,8 +136,21 @@
 admin.site.register(models.AclGroup, AclGroupAdmin)
 
 
+class DroneSetForm(forms.ModelForm):
+    def __init__(self, *args, **kwargs):
+        super(DroneSetForm, self).__init__(*args, **kwargs)
+        drone_ids_used = set()
+        for drone_set in models.DroneSet.objects.exclude(id=self.instance.id):
+            drone_ids_used.update(drone_set.drones.values_list('id', flat=True))
+        available_drones = models.Drone.objects.exclude(id__in=drone_ids_used)
+
+        self.fields['drones'].widget.choices = [(drone.id, drone.hostname)
+                                                for drone in available_drones]
+
+
 class DroneSetAdmin(SiteAdmin):
     filter_horizontal = ('drones',)
+    form = DroneSetForm
 
 admin.site.register(models.DroneSet, DroneSetAdmin)
 
diff --git a/frontend/migrations/062_drone_sets_unique.py b/frontend/migrations/062_drone_sets_unique.py
new file mode 100644
index 0000000..b1253eb
--- /dev/null
+++ b/frontend/migrations/062_drone_sets_unique.py
@@ -0,0 +1,40 @@
+UP_SQL = """
+CREATE INDEX afe_drone_sets_drones_droneset_ibfk
+ON afe_drone_sets_drones (droneset_id);
+
+ALTER TABLE afe_drone_sets_drones
+DROP KEY afe_drone_sets_drones_unique;
+
+ALTER TABLE afe_drone_sets_drones
+ADD CONSTRAINT afe_drone_sets_drones_unique
+UNIQUE KEY (drone_id);
+
+ALTER TABLE afe_drone_sets_drones
+DROP KEY afe_drone_sets_drones_drone_ibfk;
+"""
+
+DOWN_SQL = """
+CREATE INDEX afe_drone_sets_drones_drone_ibfk
+ON afe_drone_sets_drones (drone_id);
+
+ALTER TABLE afe_drone_sets_drones
+DROP KEY afe_drone_sets_drones_unique;
+
+ALTER TABLE afe_drone_sets_drones
+ADD CONSTRAINT afe_drone_sets_drones_unique
+UNIQUE KEY (droneset_id, drone_id);
+
+ALTER TABLE afe_drone_sets_drones
+DROP KEY afe_drone_sets_drones_droneset_ibfk;
+"""
+
+
+def migrate_up(manager):
+    query = ('SELECT * FROM afe_drone_sets_drones '
+             'GROUP BY drone_id HAVING COUNT(*) > 1')
+    rows = manager.execute(query)
+    if rows:
+      raise Exception('Some drones are associated with more than one drone '
+                      'set. Please remove all duplicates before running this '
+                      'migration.')
+    manager.execute_script(UP_SQL)