site_utils/label_cleaner.py - platform/external/autotest - Gitiles

 #!/usr/bin/python
 #
 # 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
 # whitelist is a second level defence against populating the 'label_prefix'
 # field with arbitrary values provided on the commandline.
 _LABEL_PREFIX_METRICS_WHITELIST = (
         '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_WHITELIST:
         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())
	#!/usr/bin/python
	#
	# 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
	# whitelist is a second level defence against populating the 'label_prefix'
	# field with arbitrary values provided on the commandline.
	_LABEL_PREFIX_METRICS_WHITELIST = (
	'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_WHITELIST:
	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())