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 os 27 import socket 28 import sys 29 import tempfile 30 31 import common 32 # Installed via build_externals, must be after import common. 33 import MySQLdb 34 from autotest_lib.client.common_lib import global_config 35 from autotest_lib.client.common_lib import logging_config 36 from autotest_lib.server import frontend 37 from chromite.lib import metrics 38 from chromite.lib import ts_mon_config 39 40 41 _METRICS_PREFIX = 'chromeos/autotest/afe_db/admin/label_cleaner' 42 43 GLOBAL_AFE = global_config.global_config.get_config_value( 44 'SERVER', 'global_afe_hostname') 45 DB_SERVER = global_config.global_config.get_config_value('AUTOTEST_WEB', 'host') 46 USER = global_config.global_config.get_config_value('AUTOTEST_WEB', 'user') 47 PASSWD = global_config.global_config.get_config_value( 48 'AUTOTEST_WEB', 'password') 49 DATABASE = global_config.global_config.get_config_value( 50 'AUTOTEST_WEB', 'database') 51 RESPECT_STATIC_LABELS = global_config.global_config.get_config_value( 52 'SKYLAB', 'respect_static_labels', type=bool, default=False) 53 54 # Per-prefix metrics are generated only for the following prefixes. This 55 # whitelist is a second level defence against populating the 'label_prefix' 56 # field with arbitrary values provided on the commandline. 57 _LABEL_PREFIX_METRICS_WHITELIST = ( 58 'cros-version', 59 'fwro-version', 60 'fwrw-version', 61 'pool', 62 ) 63 64 SELECT_USED_LABELS_FORMAT = """ 65 SELECT DISTINCT(label_id) FROM afe_autotests_dependency_labels UNION 66 SELECT DISTINCT(label_id) FROM afe_hosts_labels UNION 67 SELECT DISTINCT(label_id) FROM afe_jobs_dependency_labels UNION 68 SELECT DISTINCT(label_id) FROM afe_shards_labels UNION 69 SELECT DISTINCT(label_id) FROM afe_parameterized_jobs UNION 70 SELECT DISTINCT(meta_host) FROM afe_host_queue_entries 71 """ 72 73 SELECT_REPLACED_LABELS = """ 74 SELECT label_id FROM afe_replaced_labels 75 """ 76 77 DELETE_LABELS_FORMAT = """ 78 DELETE FROM afe_labels WHERE id in (%s) 79 """ 80 81 82 def get_used_labels(conn): 83 """Get labels that are currently in use. 84 85 @param conn: MySQLdb Connection object. 86 87 @return: A list of label ids. 88 """ 89 cursor = conn.cursor() 90 sql = SELECT_USED_LABELS_FORMAT 91 logging.debug('Running: %r', sql) 92 cursor.execute(sql) 93 rows = cursor.fetchall() 94 return set(r[0] for r in rows) 95 96 97 def fetch_labels(conn, label=None, prefix=False): 98 """Fetch labels from database. 99 100 @param conn: MySQLdb Connection object. 101 @param label: (optional) Label name to fetch. 102 @param prefix: If True, use `label` as a prefix. Otherwise, fetch 103 labels whose name is exactly same as `label`. 104 105 @return: A list of label ids. 106 """ 107 cursor = conn.cursor() 108 if label is not None: 109 if prefix: 110 sql = 'SELECT id FROM afe_labels WHERE name LIKE "%s%%"' % label 111 else: 112 sql = 'SELECT id FROM afe_labels WHERE name = "%s"' % label 113 else: 114 sql = 'SELECT id FROM afe_labels' 115 logging.debug('Running: %r', sql) 116 cursor.execute(sql) 117 rows = cursor.fetchall() 118 # Don't delete labels whose replaced_by_static_label=True, since they're 119 # actually maintained by afe_static_labels, not afe_labels. 120 if not RESPECT_STATIC_LABELS: 121 return set(r[0] for r in rows) 122 else: 123 cursor.execute(SELECT_REPLACED_LABELS) 124 replaced_labels = cursor.fetchall() 125 replaced_label_ids = set([r[0] for r in replaced_labels]) 126 return set(r[0] for r in rows) - replaced_label_ids 127 128 129 def _delete_labels(conn, labels, dry_run): 130 """Helper function of `delete_labels`.""" 131 labels_str = ','.join([str(l) for l in labels]) 132 sql = DELETE_LABELS_FORMAT % labels_str 133 if dry_run: 134 logging.info('[DRY RUN] Would have run: %r', sql) 135 else: 136 logging.debug('Running: %r', sql) 137 conn.cursor().execute(sql) 138 conn.commit() 139 140 141 def delete_labels(conn, labels, max_delete, dry_run=False): 142 """Delete given labels from database. 143 144 @param conn: MySQLdb Connection object. 145 @param labels: iterable of labels to delete. 146 @param max_delete: Max number of records to delete in a query. 147 @param dry_run: (Boolean) Whether this is a dry run. 148 """ 149 while labels: 150 chunk = labels[:max_delete] 151 labels = labels[max_delete:] 152 _delete_labels(conn, chunk, dry_run) 153 154 155 def is_primary_server(): 156 """Check if this server's status is primary 157 158 @return: True if primary, False otherwise. 159 """ 160 server = frontend.AFE(server=GLOBAL_AFE).run( 161 'get_servers', hostname=socket.getfqdn()) 162 if server and server[0]['status'] == 'primary': 163 return True 164 return False 165 166 167 def clean_labels(options): 168 """Cleans unused labels from AFE database""" 169 msg = 'Label cleaner starts. Will delete ' 170 if options.prefix: 171 msg += 'all labels whose prefix is "%s".' 172 else: 173 msg += 'a label "%s".' 174 logging.info(msg, options.label) 175 logging.info('Target database: %s.', options.db_server) 176 if options.check_status and not is_primary_server(): 177 raise Exception('Cannot run in a non-primary server') 178 179 conn = MySQLdb.connect( 180 host=options.db_server, 181 user=options.db_user, 182 passwd=options.db_password, 183 db=DATABASE, 184 ) 185 186 all_labels = fetch_labels(conn) 187 logging.info('Found total %d labels', len(all_labels)) 188 metrics.Gauge(_METRICS_PREFIX + '/total_labels_count').set( 189 len(all_labels), 190 fields={ 191 'target_db': options.db_server, 192 'label_prefix': '', 193 }, 194 ) 195 196 labels = fetch_labels(conn, options.label, options.prefix) 197 logging.info('Found total %d labels matching %s', len(labels), 198 options.label) 199 if options.prefix and options.label in _LABEL_PREFIX_METRICS_WHITELIST: 200 metrics.Gauge(_METRICS_PREFIX + '/total_labels_count').set( 201 len(labels), 202 fields={ 203 'target_db': options.db_server, 204 'label_prefix': options.label, 205 }, 206 ) 207 208 used_labels = get_used_labels(conn) 209 logging.info('Found %d labels are used', len(used_labels)) 210 metrics.Gauge(_METRICS_PREFIX + '/used_labels_count').set( 211 len(used_labels), fields={'target_db': options.db_server}) 212 213 to_delete = list(labels - used_labels) 214 logging.info('Deleting %d unused labels', len(to_delete)) 215 delete_labels(conn, to_delete, options.max_delete, options.dry_run) 216 metrics.Counter(_METRICS_PREFIX + '/labels_deleted').increment_by( 217 len(to_delete), fields={'target_db': options.db_server}) 218 219 220 def main(): 221 """Cleans unused labels from AFE database""" 222 parser = argparse.ArgumentParser( 223 formatter_class=argparse.ArgumentDefaultsHelpFormatter) 224 parser.add_argument( 225 '--db', 226 dest='db_server', 227 help='Database server', 228 default=DB_SERVER, 229 ) 230 parser.add_argument( 231 '--db-user', 232 dest='db_user', 233 help='Database user', 234 default=USER, 235 ) 236 parser.add_argument( 237 '--db-password', 238 dest='db_password', 239 help='Database password', 240 default=PASSWD, 241 ) 242 parser.add_argument( 243 '-p', 244 dest='prefix', 245 action='store_true', 246 help=('Use argument <label> as a prefix for matching. ' 247 'For example, when the argument <label> is "cros-version" ' 248 'and this option is enabled, then labels whose name ' 249 'beginning with "cros-version" are matched. When this ' 250 'option is disabled, we match labels whose name is ' 251 'exactly same as the argument <label>.'), 252 ) 253 parser.add_argument( 254 '-n', 255 dest='max_delete', 256 type=int, 257 help='Max number of records to delete in each query.', 258 default=100, 259 ) 260 parser.add_argument( 261 '-s', 262 dest='check_status', 263 action='store_true', 264 help='Enforce to run only in a server that has primary status', 265 ) 266 parser.add_argument( 267 '--dry-run', 268 dest='dry_run', 269 action='store_true', 270 help='Dry run mode. Do not actually delete any labels.', 271 ) 272 parser.add_argument('label', help='Label name to delete') 273 options = parser.parse_args() 274 275 logging_config.LoggingConfig().configure_logging( 276 datefmt='%Y-%m-%d %H:%M:%S', 277 verbose=True) 278 279 if options.dry_run: 280 tfd, metrics_file=tempfile.mkstemp() 281 os.close(tfd) 282 ts_mon_context = ts_mon_config.SetupTsMonGlobalState( 283 'afe_label_cleaner', 284 auto_flush=False, 285 debug_file=metrics_file, 286 ) 287 else: 288 ts_mon_context = ts_mon_config.SetupTsMonGlobalState( 289 'afe_label_cleaner', 290 auto_flush=False, 291 ) 292 with ts_mon_context: 293 try: 294 clean_labels(options) 295 except: 296 metrics.Counter(_METRICS_PREFIX + '/tick').increment( 297 fields={'target_db': options.db_server, 298 'success': False}) 299 raise 300 else: 301 metrics.Counter(_METRICS_PREFIX + '/tick').increment( 302 fields={'target_db': options.db_server, 303 'success': True}) 304 finally: 305 metrics.Flush() 306 if options.dry_run: 307 logging.info('Dumped ts_mon metrics to %s', metrics_file) 308 309 310 if __name__ == '__main__': 311 sys.exit(main()) 312
