site_utils/label_cleaner.py

#!/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.


"""Script to clean 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 MySQLdb
import sys
import traceback

import common
from autotest_lib.client.common_lib import global_config
from autotest_lib.client.common_lib import logging_config


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')

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_LABELS_FORMAT = """
SELECT id FROM afe_labels WHERE name %s
"""

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
    try:
        cursor.execute(sql)
        rows = cursor.fetchall()
    except:
        logging.error("Query failed: %s", sql)
        raise
    return set(r[0] for r in rows)


def fetch_labels(conn, label, prefix):
    """Fetch labels from database.

    @param conn: MySQLdb Connection object.
    @param label: 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 prefix:
        sql = SELECT_LABELS_FORMAT % ('LIKE "%s%%"' % label)
    else:
        sql = SELECT_LABELS_FORMAT % ('= "%s"' % label)
    try:
        cursor.execute(sql)
        rows = cursor.fetchall()
    except:
        logging.error("Query failed: %s", sql)
        raise
    return set(r[0] for r in rows)


def _delete_labels(conn, labels):
    """Helper function of `delete_labels`."""
    labels_str = ','.join([str(l) for l in labels])
    logging.info("Deleting following labels: %s ..", labels_str)
    sql = DELETE_LABELS_FORMAT % labels_str
    try:
        conn.cursor().execute(sql)
        conn.commit()
    except:
        logging.error("Query failed: %s", sql)
        raise
    logging.info("Done.")


def delete_labels(conn, labels, max_delete):
    """Delete given labels from database.

    @param conn: MySQLdb Connection object.
    @param labels: Labels to delete. Set type.
    @param max_delete: Max number of records to delete in a query.
    """
    if not labels:
        logging.warn("No label to delete.")
        return
    while labels:
        labels_to_del = set()
        for i in xrange(min(len(labels), max_delete)):
            labels_to_del.add(labels.pop())
        _delete_labels(conn, labels_to_del)


def main():
    parser = argparse.ArgumentParser(
            formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--db', dest='db_server',
                        help='Database server', default=DB_SERVER)
    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('label', help='Label name to delete')
    options = parser.parse_args()

    logging_config.LoggingConfig().configure_logging(
            datefmt='%Y-%m-%d %H:%M:%S')

    try:
        conn = MySQLdb.connect(host=options.db_server, user=USER,
                               passwd=PASSWD, db=DATABASE)
        used_labels = get_used_labels(conn)
        labels = fetch_labels(conn, options.label, options.prefix)
        delete_labels(conn, labels - used_labels, options.max_delete)
    except:
        logging.error(traceback.format_exc())
        return 1


if __name__ == '__main__':
    sys.exit(main())