1def execute_safely(manager, statement):
2    try:
3        manager.execute(statement)
4    except Exception:
5        print 'Statement %r failed (this is not fatal)' % statement
6
7
8def delete_duplicates(manager, table, first_id, second_id):
9    rows = manager.execute(
10        'SELECT %s, %s, COUNT(1) AS count FROM %s '
11        'GROUP BY %s, %s HAVING count > 1' %
12        (first_id, second_id, table, first_id, second_id))
13    for first_id_value, second_id_value, count_unused in rows:
14        manager.execute('DELETE FROM %s '
15                        'WHERE %s = %%s AND %s = %%s LIMIT 1' %
16                        (table, first_id, second_id),
17                        first_id_value, second_id_value)
18    if rows:
19        print 'Deleted %s duplicate rows from %s' % (len(rows), table)
20
21
22def delete_invalid_foriegn_keys(manager, pivot_table, foreign_key_field,
23                                destination_table):
24    manager.execute(
25        'DELETE %(table)s.* FROM %(table)s '
26        'LEFT JOIN %(destination_table)s '
27        'ON %(table)s.%(field)s = %(destination_table)s.id '
28        'WHERE %(destination_table)s.id IS NULL' %
29        dict(table=pivot_table, field=foreign_key_field,
30             destination_table=destination_table))
31    deleted_count = manager._database.rowcount
32    if deleted_count:
33        print ('Deleted %s invalid foreign key references from %s (%s)' %
34               (deleted_count, pivot_table, foreign_key_field))
35
36
37def unique_index_name(table):
38    return table + '_both_ids'
39
40
41def basic_index_name(table, field):
42    if field == 'aclgroup_id':
43        field = 'acl_group_id'
44    return table + '_' + field
45
46
47def create_unique_index(manager, pivot_table, first_field, second_field):
48    index_name = unique_index_name(pivot_table)
49    manager.execute('CREATE UNIQUE INDEX %s ON %s (%s, %s)' %
50                    (index_name, pivot_table, first_field, second_field))
51
52    # these indices are in the migrations but may not exist for historical
53    # reasons
54    old_index_name = basic_index_name(pivot_table, first_field)
55    execute_safely(manager, 'DROP INDEX %s ON %s' %
56                   (old_index_name, pivot_table))
57
58
59def drop_unique_index(manager, pivot_table, first_field):
60    index_name = unique_index_name(pivot_table)
61    manager.execute('DROP INDEX %s ON %s' % (index_name, pivot_table))
62
63    old_index_name = basic_index_name(pivot_table, first_field)
64    manager.execute('CREATE INDEX %s ON %s (%s)' %
65                    (old_index_name, pivot_table, first_field))
66
67
68def foreign_key_name(table, field):
69    return '_'.join([table, field, 'fk'])
70
71
72def create_foreign_key_constraint(manager, table, field, destination_table):
73    key_name = foreign_key_name(table, field)
74    manager.execute('ALTER TABLE %s ADD CONSTRAINT %s FOREIGN KEY (%s) '
75                    'REFERENCES %s (id) ON DELETE NO ACTION' %
76                    (table, key_name, field, destination_table))
77
78
79def drop_foreign_key_constraint(manager, table, field):
80    key_name = foreign_key_name(table, field)
81    manager.execute('ALTER TABLE %s DROP FOREIGN KEY %s' % (table, key_name))
82
83
84def cleanup_m2m_pivot(manager, pivot_table, first_field, first_table,
85                      second_field, second_table, create_unique):
86    delete_duplicates(manager, pivot_table, first_field, second_field)
87    delete_invalid_foriegn_keys(manager, pivot_table, first_field, first_table)
88    delete_invalid_foriegn_keys(manager, pivot_table, second_field,
89                                second_table)
90
91    if create_unique:
92        # first field is the more commonly used one, so we'll replace the
93        # less-commonly-used index with the larger unique index
94        create_unique_index(manager, pivot_table, second_field, first_field)
95
96    create_foreign_key_constraint(manager, pivot_table, first_field,
97                                  first_table)
98    create_foreign_key_constraint(manager, pivot_table, second_field,
99                                  second_table)
100
101
102def reverse_cleanup_m2m_pivot(manager, pivot_table, first_field, second_field,
103                              drop_unique):
104    drop_foreign_key_constraint(manager, pivot_table, second_field)
105    drop_foreign_key_constraint(manager, pivot_table, first_field)
106    if drop_unique:
107        drop_unique_index(manager, pivot_table, second_field)
108
109
110TABLES = (
111        ('hosts_labels', 'host_id', 'hosts', 'label_id', 'labels', True),
112        ('acl_groups_hosts', 'host_id', 'hosts', 'aclgroup_id', 'acl_groups',
113         True),
114        ('acl_groups_users', 'user_id', 'users', 'aclgroup_id', 'acl_groups',
115         True),
116        ('autotests_dependency_labels', 'test_id', 'autotests', 'label_id',
117         'labels', False),
118        ('jobs_dependency_labels', 'job_id', 'jobs', 'label_id', 'labels',
119         False),
120        ('ineligible_host_queues', 'job_id', 'jobs', 'host_id', 'hosts', True),
121    )
122
123
124def migrate_up(manager):
125    for (table, first_field, first_table, second_field, second_table,
126         create_unique) in TABLES:
127        cleanup_m2m_pivot(manager, table, first_field, first_table,
128                          second_field, second_table, create_unique)
129
130
131def migrate_down(manager):
132    for (table, first_field, first_table, second_field, second_table,
133         drop_unique) in reversed(TABLES):
134        reverse_cleanup_m2m_pivot(manager, table, first_field, second_field,
135                                  drop_unique)
136