1import common
2from autotest_lib.database import db_utils
3
4UP_SQL = """
5CREATE INDEX afe_drone_sets_drones_droneset_ibfk
6ON afe_drone_sets_drones (droneset_id);
7
8ALTER TABLE afe_drone_sets_drones
9DROP KEY afe_drone_sets_drones_unique;
10
11ALTER TABLE afe_drone_sets_drones
12ADD CONSTRAINT afe_drone_sets_drones_unique
13UNIQUE KEY (drone_id);
14"""
15
16# On first migration to 62, this key will be deleted automatically. However, if
17# you migrate to 62, then down to 61, then back to 62, this key will remain.
18DROP_KEY_SQL = """
19ALTER TABLE afe_drone_sets_drones
20DROP KEY afe_drone_sets_drones_drone_ibfk;
21"""
22
23DOWN_SQL = """
24CREATE INDEX afe_drone_sets_drones_drone_ibfk
25ON afe_drone_sets_drones (drone_id);
26
27ALTER TABLE afe_drone_sets_drones
28DROP KEY afe_drone_sets_drones_unique;
29
30ALTER TABLE afe_drone_sets_drones
31ADD CONSTRAINT afe_drone_sets_drones_unique
32UNIQUE KEY (droneset_id, drone_id);
33
34ALTER TABLE afe_drone_sets_drones
35DROP KEY afe_drone_sets_drones_droneset_ibfk;
36"""
37
38
39def migrate_up(manager):
40    query = ('SELECT * FROM afe_drone_sets_drones '
41             'GROUP BY drone_id HAVING COUNT(*) > 1')
42    rows = manager.execute(query)
43    if rows:
44        raise Exception('Some drones are associated with more than one drone '
45                        'set. Please remove all duplicates before running this '
46                        'migration.')
47    manager.execute_script(UP_SQL)
48
49    if db_utils.check_index_exists(manager, 'afe_drone_sets_drones',
50                                   'afe_drone_sets_drones_drone_ibfk'):
51        manager.execute(DROP_KEY_SQL)
52