blob: 2113fc1da3e601f0cf65ae3abd70fddc0ca65453 [file] [log] [blame]
MK Ryu40557042015-08-18 10:59:47 -07001#!/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
MK Ryua34e3b12015-08-21 16:20:47 -07008"""Tool for cleaning up labels that are not in use.
MK Ryu40557042015-08-18 10:59:47 -07009
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
Prathmesh Prabhu2f02ab82018-01-11 11:15:54 -080026import os
MK Ryua34e3b12015-08-21 16:20:47 -070027import socket
MK Ryu40557042015-08-18 10:59:47 -070028import sys
Prathmesh Prabhu2f02ab82018-01-11 11:15:54 -080029import tempfile
MK Ryu40557042015-08-18 10:59:47 -070030
31import common
Prathmesh Prabhu194b8702017-08-07 15:55:50 -070032# Installed via build_externals, must be after import common.
33import MySQLdb
MK Ryu40557042015-08-18 10:59:47 -070034from autotest_lib.client.common_lib import global_config
35from autotest_lib.client.common_lib import logging_config
MK Ryua34e3b12015-08-21 16:20:47 -070036from autotest_lib.server import frontend
Prathmesh Prabhu82684492017-08-07 16:28:37 -070037from chromite.lib import metrics
38from chromite.lib import ts_mon_config
MK Ryu40557042015-08-18 10:59:47 -070039
40
Prathmesh Prabhu82684492017-08-07 16:28:37 -070041_METRICS_PREFIX = 'chromeos/autotest/afe_db/admin/label_cleaner'
42
MK Ryu37325b42015-08-25 15:41:39 -070043GLOBAL_AFE = global_config.global_config.get_config_value(
44 'SERVER', 'global_afe_hostname')
MK Ryu40557042015-08-18 10:59:47 -070045DB_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')
Xixuan Wu211ae692017-12-18 15:32:42 -080051RESPECT_STATIC_LABELS = global_config.global_config.get_config_value(
52 'SKYLAB', 'respect_static_labels', type=bool, default=False)
MK Ryu40557042015-08-18 10:59:47 -070053
Prathmesh Prabhu2f02ab82018-01-11 11:15:54 -080054# 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
MK Ryu40557042015-08-18 10:59:47 -070064SELECT_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
Xixuan Wu211ae692017-12-18 15:32:42 -080073SELECT_REPLACED_LABELS = """
74SELECT label_id FROM afe_replaced_labels
75"""
MK Ryu40557042015-08-18 10:59:47 -070076
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
Prathmesh Prabhu82684492017-08-07 16:28:37 -070091 logging.debug('Running: %r', sql)
92 cursor.execute(sql)
93 rows = cursor.fetchall()
MK Ryu40557042015-08-18 10:59:47 -070094 return set(r[0] for r in rows)
95
96
Prathmesh Prabhu2f02ab82018-01-11 11:15:54 -080097def fetch_labels(conn, label=None, prefix=False):
MK Ryu40557042015-08-18 10:59:47 -070098 """Fetch labels from database.
99
100 @param conn: MySQLdb Connection object.
Prathmesh Prabhu2f02ab82018-01-11 11:15:54 -0800101 @param label: (optional) Label name to fetch.
MK Ryu40557042015-08-18 10:59:47 -0700102 @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()
Prathmesh Prabhu2f02ab82018-01-11 11:15:54 -0800108 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
MK Ryu40557042015-08-18 10:59:47 -0700113 else:
Prathmesh Prabhu2f02ab82018-01-11 11:15:54 -0800114 sql = 'SELECT id FROM afe_labels'
Prathmesh Prabhu82684492017-08-07 16:28:37 -0700115 logging.debug('Running: %r', sql)
116 cursor.execute(sql)
117 rows = cursor.fetchall()
Xixuan Wu211ae692017-12-18 15:32:42 -0800118 # 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
MK Ryu40557042015-08-18 10:59:47 -0700127
128
Prathmesh Prabhu2f02ab82018-01-11 11:15:54 -0800129def _delete_labels(conn, labels, dry_run):
MK Ryu40557042015-08-18 10:59:47 -0700130 """Helper function of `delete_labels`."""
131 labels_str = ','.join([str(l) for l in labels])
MK Ryu40557042015-08-18 10:59:47 -0700132 sql = DELETE_LABELS_FORMAT % labels_str
Prathmesh Prabhu2f02ab82018-01-11 11:15:54 -0800133 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()
MK Ryu40557042015-08-18 10:59:47 -0700139
140
Prathmesh Prabhu2f02ab82018-01-11 11:15:54 -0800141def delete_labels(conn, labels, max_delete, dry_run=False):
MK Ryu40557042015-08-18 10:59:47 -0700142 """Delete given labels from database.
143
144 @param conn: MySQLdb Connection object.
Prathmesh Prabhu82684492017-08-07 16:28:37 -0700145 @param labels: iterable of labels to delete.
MK Ryu40557042015-08-18 10:59:47 -0700146 @param max_delete: Max number of records to delete in a query.
Prathmesh Prabhu2f02ab82018-01-11 11:15:54 -0800147 @param dry_run: (Boolean) Whether this is a dry run.
MK Ryu40557042015-08-18 10:59:47 -0700148 """
MK Ryu40557042015-08-18 10:59:47 -0700149 while labels:
Prathmesh Prabhu82684492017-08-07 16:28:37 -0700150 chunk = labels[:max_delete]
151 labels = labels[max_delete:]
Prathmesh Prabhu2f02ab82018-01-11 11:15:54 -0800152 _delete_labels(conn, chunk, dry_run)
MK Ryu40557042015-08-18 10:59:47 -0700153
154
MK Ryua34e3b12015-08-21 16:20:47 -0700155def is_primary_server():
156 """Check if this server's status is primary
157
158 @return: True if primary, False otherwise.
159 """
MK Ryu37325b42015-08-25 15:41:39 -0700160 server = frontend.AFE(server=GLOBAL_AFE).run(
161 'get_servers', hostname=socket.getfqdn())
162 if server and server[0]['status'] == 'primary':
MK Ryua34e3b12015-08-21 16:20:47 -0700163 return True
164 return False
165
166
Prathmesh Prabhu82684492017-08-07 16:28:37 -0700167def 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
Prathmesh Prabhu2f02ab82018-01-11 11:15:54 -0800179 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 )
Prathmesh Prabhu82684492017-08-07 16:28:37 -0700195
196 labels = fetch_labels(conn, options.label, options.prefix)
Prathmesh Prabhu2f02ab82018-01-11 11:15:54 -0800197 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 )
Prathmesh Prabhu82684492017-08-07 16:28:37 -0700207
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))
Prathmesh Prabhu2f02ab82018-01-11 11:15:54 -0800215 delete_labels(conn, to_delete, options.max_delete, options.dry_run)
Prathmesh Prabhu82684492017-08-07 16:28:37 -0700216 metrics.Counter(_METRICS_PREFIX + '/labels_deleted').increment_by(
217 len(to_delete), fields={'target_db': options.db_server})
218
219
MK Ryu40557042015-08-18 10:59:47 -0700220def main():
Prathmesh Prabhu82684492017-08-07 16:28:37 -0700221 """Cleans unused labels from AFE database"""
MK Ryu40557042015-08-18 10:59:47 -0700222 parser = argparse.ArgumentParser(
223 formatter_class=argparse.ArgumentDefaultsHelpFormatter)
Prathmesh Prabhu2f02ab82018-01-11 11:15:54 -0800224 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',
MK Ryu40557042015-08-18 10:59:47 -0700246 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 '
Prathmesh Prabhu2f02ab82018-01-11 11:15:54 -0800251 '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 )
MK Ryu40557042015-08-18 10:59:47 -0700272 parser.add_argument('label', help='Label name to delete')
273 options = parser.parse_args()
274
275 logging_config.LoggingConfig().configure_logging(
Prathmesh Prabhu82684492017-08-07 16:28:37 -0700276 datefmt='%Y-%m-%d %H:%M:%S',
277 verbose=True)
MK Ryu40557042015-08-18 10:59:47 -0700278
Prathmesh Prabhu2f02ab82018-01-11 11:15:54 -0800279 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:
Prathmesh Prabhu82684492017-08-07 16:28:37 -0700293 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
MK Ryu105e1c22015-09-01 15:58:43 -0700300 else:
Prathmesh Prabhu82684492017-08-07 16:28:37 -0700301 metrics.Counter(_METRICS_PREFIX + '/tick').increment(
302 fields={'target_db': options.db_server,
303 'success': True})
304 finally:
305 metrics.Flush()
Prathmesh Prabhu2f02ab82018-01-11 11:15:54 -0800306 if options.dry_run:
307 logging.info('Dumped ts_mon metrics to %s', metrics_file)
MK Ryu40557042015-08-18 10:59:47 -0700308
309
310if __name__ == '__main__':
311 sys.exit(main())