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
10Delete given labels from database when they are not in use.
11Labels that match the query `SELECT_USED_LABELS_FORMAT` are considered in use.
12When given labels are not in the used labels, those labels are deleted.
13
14For 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
19If '-p' option is not given, we delete labels whose name is exactly
20'cros-version' and are not in use.
21"""
22
23
24import argparse
25import logging
26import os
27import socket
28import sys
29import tempfile
30
31import common
32# Installed via build_externals, must be after import common.
33import MySQLdb
34from autotest_lib.client.common_lib import global_config
35from autotest_lib.client.common_lib import logging_config
36from autotest_lib.server import frontend
37from chromite.lib import metrics
38from chromite.lib import ts_mon_config
39
40
41_METRICS_PREFIX = 'chromeos/autotest/afe_db/admin/label_cleaner'
42
43GLOBAL_AFE = global_config.global_config.get_config_value(
44        'SERVER', 'global_afe_hostname')
45DB_SERVER = global_config.global_config.get_config_value('AUTOTEST_WEB', 'host')
46USER = global_config.global_config.get_config_value('AUTOTEST_WEB', 'user')
47PASSWD = global_config.global_config.get_config_value(
48        'AUTOTEST_WEB', 'password')
49DATABASE = global_config.global_config.get_config_value(
50        'AUTOTEST_WEB', 'database')
51RESPECT_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
64SELECT_USED_LABELS_FORMAT = """
65SELECT DISTINCT(label_id) FROM afe_autotests_dependency_labels UNION
66SELECT DISTINCT(label_id) FROM afe_hosts_labels UNION
67SELECT DISTINCT(label_id) FROM afe_jobs_dependency_labels UNION
68SELECT DISTINCT(label_id) FROM afe_shards_labels UNION
69SELECT DISTINCT(label_id) FROM afe_parameterized_jobs UNION
70SELECT DISTINCT(meta_host) FROM afe_host_queue_entries
71"""
72
73SELECT_REPLACED_LABELS = """
74SELECT label_id FROM afe_replaced_labels
75"""
76
77DELETE_LABELS_FORMAT = """
78DELETE FROM afe_labels WHERE id in (%s)
79"""
80
81
82def 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
97def 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
129def _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
141def 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
155def 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
167def 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
220def 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
310if __name__ == '__main__':
311    sys.exit(main())
312