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