1#!/usr/bin/python
2#
3# Copyright (c) 2015 The Chromium OS Authors. All rights reserved.
4# Use of this source code is governed by a BSD-style license that can be
5# found in the LICENSE file.
6
7
8"""Tool for cleaning up labels that are not in use.
9
10Delete given labels from database when they are not in use.
11Labels that match the query `SELECT_USED_LABELS_FORMAT` are considered in use.
12When given labels are not in the used labels, those labels are deleted.
13
14For example, following command deletes all labels whose name begins with
15'cros-version' and are not in use.
16
17./label_cleaner.py -p cros-version
18
19If '-p' option is not given, we delete labels whose name is exactly
20'cros-version' and are not in use.
21"""
22
23
24import argparse
25import logging
26import MySQLdb
27import socket
28import sys
29import traceback
30
31import common
32from autotest_lib.client.common_lib import global_config
33from autotest_lib.client.common_lib import logging_config
34from autotest_lib.server import frontend
35
36
37GLOBAL_AFE = global_config.global_config.get_config_value(
38        'SERVER', 'global_afe_hostname')
39DB_SERVER = global_config.global_config.get_config_value('AUTOTEST_WEB', 'host')
40USER = global_config.global_config.get_config_value('AUTOTEST_WEB', 'user')
41PASSWD = global_config.global_config.get_config_value(
42        'AUTOTEST_WEB', 'password')
43DATABASE = global_config.global_config.get_config_value(
44        'AUTOTEST_WEB', 'database')
45
46SELECT_USED_LABELS_FORMAT = """
47SELECT DISTINCT(label_id) FROM afe_autotests_dependency_labels UNION
48SELECT DISTINCT(label_id) FROM afe_hosts_labels UNION
49SELECT DISTINCT(label_id) FROM afe_jobs_dependency_labels UNION
50SELECT DISTINCT(label_id) FROM afe_shards_labels UNION
51SELECT DISTINCT(label_id) FROM afe_parameterized_jobs UNION
52SELECT DISTINCT(meta_host) FROM afe_host_queue_entries
53"""
54
55SELECT_LABELS_FORMAT = """
56SELECT id FROM afe_labels WHERE name %s
57"""
58
59DELETE_LABELS_FORMAT = """
60DELETE FROM afe_labels WHERE id in (%s)
61"""
62
63
64def get_used_labels(conn):
65    """Get labels that are currently in use.
66
67    @param conn: MySQLdb Connection object.
68
69    @return: A list of label ids.
70    """
71    cursor = conn.cursor()
72    sql = SELECT_USED_LABELS_FORMAT
73    try:
74        cursor.execute(sql)
75        rows = cursor.fetchall()
76    except:
77        logging.error("Query failed: %s", sql)
78        raise
79    return set(r[0] for r in rows)
80
81
82def fetch_labels(conn, label, prefix):
83    """Fetch labels from database.
84
85    @param conn: MySQLdb Connection object.
86    @param label: Label name to fetch.
87    @param prefix: If True, use `label` as a prefix. Otherwise, fetch
88                   labels whose name is exactly same as `label`.
89
90    @return: A list of label ids.
91    """
92    cursor = conn.cursor()
93    if prefix:
94        sql = SELECT_LABELS_FORMAT % ('LIKE "%s%%"' % label)
95    else:
96        sql = SELECT_LABELS_FORMAT % ('= "%s"' % label)
97    try:
98        cursor.execute(sql)
99        rows = cursor.fetchall()
100    except:
101        logging.error("Query failed: %s", sql)
102        raise
103    return set(r[0] for r in rows)
104
105
106def _delete_labels(conn, labels):
107    """Helper function of `delete_labels`."""
108    labels_str = ','.join([str(l) for l in labels])
109    logging.info("Deleting following labels: %s ..", labels_str)
110    sql = DELETE_LABELS_FORMAT % labels_str
111    try:
112        conn.cursor().execute(sql)
113        conn.commit()
114    except:
115        logging.error("Query failed: %s", sql)
116        raise
117    logging.info("Done.")
118
119
120def delete_labels(conn, labels, max_delete):
121    """Delete given labels from database.
122
123    @param conn: MySQLdb Connection object.
124    @param labels: Labels to delete. Set type.
125    @param max_delete: Max number of records to delete in a query.
126    """
127    if not labels:
128        logging.warn("No label to delete.")
129        return
130    while labels:
131        labels_to_del = set()
132        for i in xrange(min(len(labels), max_delete)):
133            labels_to_del.add(labels.pop())
134        _delete_labels(conn, labels_to_del)
135
136
137def is_primary_server():
138    """Check if this server's status is primary
139
140    @return: True if primary, False otherwise.
141    """
142    server = frontend.AFE(server=GLOBAL_AFE).run(
143            'get_servers', hostname=socket.getfqdn())
144    if server and server[0]['status'] == 'primary':
145        return True
146    return False
147
148
149def main():
150    parser = argparse.ArgumentParser(
151            formatter_class=argparse.ArgumentDefaultsHelpFormatter)
152    parser.add_argument('--db', dest='db_server',
153                        help='Database server', default=DB_SERVER)
154    parser.add_argument('-p', dest='prefix', action='store_true',
155            help=('Use argument <label> as a prefix for matching. '
156                  'For example, when the argument <label> is "cros-version" '
157                  'and this option is enabled, then labels whose name '
158                  'beginning with "cros-version" are matched. When this '
159                  'option is disabled, we match labels whose name is '
160                  'exactly same as the argument <label>.'))
161    parser.add_argument('-n', dest='max_delete', type=int,
162           help=('Max number of records to delete in each query.'),
163           default=100)
164    parser.add_argument('-s', dest='check_status', action='store_true',
165           help=('Enforce to run only in a server that has primary status'))
166    parser.add_argument('label', help='Label name to delete')
167    options = parser.parse_args()
168
169    logging_config.LoggingConfig().configure_logging(
170            datefmt='%Y-%m-%d %H:%M:%S')
171
172    try:
173        msg = 'Label cleaner starts. Will delete '
174        if options.prefix:
175            msg += 'all labels whose prefix is "%s".'
176        else:
177            msg += 'a label "%s".'
178        logging.info(msg, options.label)
179        logging.info('Target database: %s.', options.db_server)
180        if options.check_status and not is_primary_server():
181            logging.error('Cannot run in a non-primary server.')
182            return 1
183
184        conn = MySQLdb.connect(host=options.db_server, user=USER,
185                               passwd=PASSWD, db=DATABASE)
186        used_labels = get_used_labels(conn)
187        labels = fetch_labels(conn, options.label, options.prefix)
188        delete_labels(conn, labels - used_labels, options.max_delete)
189        logging.info('Done.')
190    except:
191        logging.error(traceback.format_exc())
192        return 1
193
194
195if __name__ == '__main__':
196    sys.exit(main())
197