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 10 Delete given labels from database when they are not in use. 11 Labels that match the query `SELECT_USED_LABELS_FORMAT` are considered in use. 12 When given labels are not in the used labels, those labels are deleted. 13 14 For 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 19 If '-p' option is not given, we delete labels whose name is exactly 20 'cros-version' and are not in use. 21 """ 22 23 24 import argparse 25 import logging 26 import MySQLdb 27 import socket 28 import sys 29 import traceback 30 31 import common 32 from autotest_lib.client.common_lib import global_config 33 from autotest_lib.client.common_lib import logging_config 34 from autotest_lib.server import frontend 35 36 37 GLOBAL_AFE = global_config.global_config.get_config_value( 38 'SERVER', 'global_afe_hostname') 39 DB_SERVER = global_config.global_config.get_config_value('AUTOTEST_WEB', 'host') 40 USER = global_config.global_config.get_config_value('AUTOTEST_WEB', 'user') 41 PASSWD = global_config.global_config.get_config_value( 42 'AUTOTEST_WEB', 'password') 43 DATABASE = global_config.global_config.get_config_value( 44 'AUTOTEST_WEB', 'database') 45 46 SELECT_USED_LABELS_FORMAT = """ 47 SELECT DISTINCT(label_id) FROM afe_autotests_dependency_labels UNION 48 SELECT DISTINCT(label_id) FROM afe_hosts_labels UNION 49 SELECT DISTINCT(label_id) FROM afe_jobs_dependency_labels UNION 50 SELECT DISTINCT(label_id) FROM afe_shards_labels UNION 51 SELECT DISTINCT(label_id) FROM afe_parameterized_jobs UNION 52 SELECT DISTINCT(meta_host) FROM afe_host_queue_entries 53 """ 54 55 SELECT_LABELS_FORMAT = """ 56 SELECT id FROM afe_labels WHERE name %s 57 """ 58 59 DELETE_LABELS_FORMAT = """ 60 DELETE FROM afe_labels WHERE id in (%s) 61 """ 62 63 64 def 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 82 def 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 106 def _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 120 def 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 137 def 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 149 def 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 195 if __name__ == '__main__': 196 sys.exit(main()) 197
