| #!/usr/bin/python2 |
| # |
| # Copyright (c) 2015 The Chromium OS Authors. All rights reserved. |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| |
| |
| """Tool for cleaning up labels that are not in use. |
| |
| Delete given labels from database when they are not in use. |
| Labels that match the query `SELECT_USED_LABELS_FORMAT` are considered in use. |
| When given labels are not in the used labels, those labels are deleted. |
| |
| For example, following command deletes all labels whose name begins with |
| 'cros-version' and are not in use. |
| |
| ./label_cleaner.py -p cros-version |
| |
| If '-p' option is not given, we delete labels whose name is exactly |
| 'cros-version' and are not in use. |
| """ |
| |
| |
| import argparse |
| import logging |
| import os |
| import socket |
| import sys |
| import tempfile |
| |
| import common |
| # Installed via build_externals, must be after import common. |
| import MySQLdb |
| from autotest_lib.client.common_lib import global_config |
| from autotest_lib.client.common_lib import logging_config |
| from autotest_lib.server import frontend |
| from chromite.lib import metrics |
| from chromite.lib import ts_mon_config |
| |
| |
| _METRICS_PREFIX = 'chromeos/autotest/afe_db/admin/label_cleaner' |
| |
| GLOBAL_AFE = global_config.global_config.get_config_value( |
| 'SERVER', 'global_afe_hostname') |
| DB_SERVER = global_config.global_config.get_config_value('AUTOTEST_WEB', 'host') |
| USER = global_config.global_config.get_config_value('AUTOTEST_WEB', 'user') |
| PASSWD = global_config.global_config.get_config_value( |
| 'AUTOTEST_WEB', 'password') |
| DATABASE = global_config.global_config.get_config_value( |
| 'AUTOTEST_WEB', 'database') |
| RESPECT_STATIC_LABELS = global_config.global_config.get_config_value( |
| 'SKYLAB', 'respect_static_labels', type=bool, default=False) |
| |
| # Per-prefix metrics are generated only for the following prefixes. This |
| # allowlist is a second level defence against populating the 'label_prefix' |
| # field with arbitrary values provided on the commandline. |
| _LABEL_PREFIX_METRICS_ALLOWLIST = ( |
| 'cros-version', |
| 'fwro-version', |
| 'fwrw-version', |
| 'pool', |
| ) |
| |
| SELECT_USED_LABELS_FORMAT = """ |
| SELECT DISTINCT(label_id) FROM afe_autotests_dependency_labels UNION |
| SELECT DISTINCT(label_id) FROM afe_hosts_labels UNION |
| SELECT DISTINCT(label_id) FROM afe_jobs_dependency_labels UNION |
| SELECT DISTINCT(label_id) FROM afe_shards_labels UNION |
| SELECT DISTINCT(label_id) FROM afe_parameterized_jobs UNION |
| SELECT DISTINCT(meta_host) FROM afe_host_queue_entries |
| """ |
| |
| SELECT_REPLACED_LABELS = """ |
| SELECT label_id FROM afe_replaced_labels |
| """ |
| |
| DELETE_LABELS_FORMAT = """ |
| DELETE FROM afe_labels WHERE id in (%s) |
| """ |
| |
| |
| def get_used_labels(conn): |
| """Get labels that are currently in use. |
| |
| @param conn: MySQLdb Connection object. |
| |
| @return: A list of label ids. |
| """ |
| cursor = conn.cursor() |
| sql = SELECT_USED_LABELS_FORMAT |
| logging.debug('Running: %r', sql) |
| cursor.execute(sql) |
| rows = cursor.fetchall() |
| return set(r[0] for r in rows) |
| |
| |
| def fetch_labels(conn, label=None, prefix=False): |
| """Fetch labels from database. |
| |
| @param conn: MySQLdb Connection object. |
| @param label: (optional) Label name to fetch. |
| @param prefix: If True, use `label` as a prefix. Otherwise, fetch |
| labels whose name is exactly same as `label`. |
| |
| @return: A list of label ids. |
| """ |
| cursor = conn.cursor() |
| if label is not None: |
| if prefix: |
| sql = 'SELECT id FROM afe_labels WHERE name LIKE "%s%%"' % label |
| else: |
| sql = 'SELECT id FROM afe_labels WHERE name = "%s"' % label |
| else: |
| sql = 'SELECT id FROM afe_labels' |
| logging.debug('Running: %r', sql) |
| cursor.execute(sql) |
| rows = cursor.fetchall() |
| # Don't delete labels whose replaced_by_static_label=True, since they're |
| # actually maintained by afe_static_labels, not afe_labels. |
| if not RESPECT_STATIC_LABELS: |
| return set(r[0] for r in rows) |
| else: |
| cursor.execute(SELECT_REPLACED_LABELS) |
| replaced_labels = cursor.fetchall() |
| replaced_label_ids = set([r[0] for r in replaced_labels]) |
| return set(r[0] for r in rows) - replaced_label_ids |
| |
| |
| def _delete_labels(conn, labels, dry_run): |
| """Helper function of `delete_labels`.""" |
| labels_str = ','.join([str(l) for l in labels]) |
| sql = DELETE_LABELS_FORMAT % labels_str |
| if dry_run: |
| logging.info('[DRY RUN] Would have run: %r', sql) |
| else: |
| logging.debug('Running: %r', sql) |
| conn.cursor().execute(sql) |
| conn.commit() |
| |
| |
| def delete_labels(conn, labels, max_delete, dry_run=False): |
| """Delete given labels from database. |
| |
| @param conn: MySQLdb Connection object. |
| @param labels: iterable of labels to delete. |
| @param max_delete: Max number of records to delete in a query. |
| @param dry_run: (Boolean) Whether this is a dry run. |
| """ |
| while labels: |
| chunk = labels[:max_delete] |
| labels = labels[max_delete:] |
| _delete_labels(conn, chunk, dry_run) |
| |
| |
| def is_primary_server(): |
| """Check if this server's status is primary |
| |
| @return: True if primary, False otherwise. |
| """ |
| server = frontend.AFE(server=GLOBAL_AFE).run( |
| 'get_servers', hostname=socket.getfqdn()) |
| if server and server[0]['status'] == 'primary': |
| return True |
| return False |
| |
| |
| def clean_labels(options): |
| """Cleans unused labels from AFE database""" |
| msg = 'Label cleaner starts. Will delete ' |
| if options.prefix: |
| msg += 'all labels whose prefix is "%s".' |
| else: |
| msg += 'a label "%s".' |
| logging.info(msg, options.label) |
| logging.info('Target database: %s.', options.db_server) |
| if options.check_status and not is_primary_server(): |
| raise Exception('Cannot run in a non-primary server') |
| |
| conn = MySQLdb.connect( |
| host=options.db_server, |
| user=options.db_user, |
| passwd=options.db_password, |
| db=DATABASE, |
| ) |
| |
| all_labels = fetch_labels(conn) |
| logging.info('Found total %d labels', len(all_labels)) |
| metrics.Gauge(_METRICS_PREFIX + '/total_labels_count').set( |
| len(all_labels), |
| fields={ |
| 'target_db': options.db_server, |
| 'label_prefix': '', |
| }, |
| ) |
| |
| labels = fetch_labels(conn, options.label, options.prefix) |
| logging.info('Found total %d labels matching %s', len(labels), |
| options.label) |
| if options.prefix and options.label in _LABEL_PREFIX_METRICS_ALLOWLIST: |
| metrics.Gauge(_METRICS_PREFIX + '/total_labels_count').set( |
| len(labels), |
| fields={ |
| 'target_db': options.db_server, |
| 'label_prefix': options.label, |
| }, |
| ) |
| |
| used_labels = get_used_labels(conn) |
| logging.info('Found %d labels are used', len(used_labels)) |
| metrics.Gauge(_METRICS_PREFIX + '/used_labels_count').set( |
| len(used_labels), fields={'target_db': options.db_server}) |
| |
| to_delete = list(labels - used_labels) |
| logging.info('Deleting %d unused labels', len(to_delete)) |
| delete_labels(conn, to_delete, options.max_delete, options.dry_run) |
| metrics.Counter(_METRICS_PREFIX + '/labels_deleted').increment_by( |
| len(to_delete), fields={'target_db': options.db_server}) |
| |
| |
| def main(): |
| """Cleans unused labels from AFE database""" |
| parser = argparse.ArgumentParser( |
| formatter_class=argparse.ArgumentDefaultsHelpFormatter) |
| parser.add_argument( |
| '--db', |
| dest='db_server', |
| help='Database server', |
| default=DB_SERVER, |
| ) |
| parser.add_argument( |
| '--db-user', |
| dest='db_user', |
| help='Database user', |
| default=USER, |
| ) |
| parser.add_argument( |
| '--db-password', |
| dest='db_password', |
| help='Database password', |
| default=PASSWD, |
| ) |
| parser.add_argument( |
| '-p', |
| dest='prefix', |
| action='store_true', |
| help=('Use argument <label> as a prefix for matching. ' |
| 'For example, when the argument <label> is "cros-version" ' |
| 'and this option is enabled, then labels whose name ' |
| 'beginning with "cros-version" are matched. When this ' |
| 'option is disabled, we match labels whose name is ' |
| 'exactly same as the argument <label>.'), |
| ) |
| parser.add_argument( |
| '-n', |
| dest='max_delete', |
| type=int, |
| help='Max number of records to delete in each query.', |
| default=100, |
| ) |
| parser.add_argument( |
| '-s', |
| dest='check_status', |
| action='store_true', |
| help='Enforce to run only in a server that has primary status', |
| ) |
| parser.add_argument( |
| '--dry-run', |
| dest='dry_run', |
| action='store_true', |
| help='Dry run mode. Do not actually delete any labels.', |
| ) |
| parser.add_argument('label', help='Label name to delete') |
| options = parser.parse_args() |
| |
| logging_config.LoggingConfig().configure_logging( |
| datefmt='%Y-%m-%d %H:%M:%S', |
| verbose=True) |
| |
| if options.dry_run: |
| tfd, metrics_file=tempfile.mkstemp() |
| os.close(tfd) |
| ts_mon_context = ts_mon_config.SetupTsMonGlobalState( |
| 'afe_label_cleaner', |
| auto_flush=False, |
| debug_file=metrics_file, |
| ) |
| else: |
| ts_mon_context = ts_mon_config.SetupTsMonGlobalState( |
| 'afe_label_cleaner', |
| auto_flush=False, |
| ) |
| with ts_mon_context: |
| try: |
| clean_labels(options) |
| except: |
| metrics.Counter(_METRICS_PREFIX + '/tick').increment( |
| fields={'target_db': options.db_server, |
| 'success': False}) |
| raise |
| else: |
| metrics.Counter(_METRICS_PREFIX + '/tick').increment( |
| fields={'target_db': options.db_server, |
| 'success': True}) |
| finally: |
| metrics.Flush() |
| if options.dry_run: |
| logging.info('Dumped ts_mon metrics to %s', metrics_file) |
| |
| |
| if __name__ == '__main__': |
| sys.exit(main()) |