1#!/usr/bin/python2 2# Copyright 2017 The Chromium Authors. All rights reserved. 3# Use of this source code is governed by a BSD-style license that can be 4# found in the LICENSE file. 5 6"""Utility to cleanup TKO database by removing old records. 7""" 8 9import argparse 10import logging 11import os 12import time 13 14import common 15from autotest_lib.client.bin import utils 16from autotest_lib.client.common_lib import global_config 17from autotest_lib.client.common_lib import logging_config 18 19from chromite.lib import metrics 20from chromite.lib import ts_mon_config 21 22 23CONFIG = global_config.global_config 24 25# SQL command to remove old test results in TKO database. 26CLEANUP_TKO_CMD = 'call remove_old_tests_sp()' 27CLEANUP_METRIC = 'chromeos/autotest/tko/cleanup_duration' 28RECREATE_TEST_ATTRIBUTES_METRIC = ( 29 'chromeos/autotest/tko/recreate_test_attributes') 30RECREATE_TABLE = 'tko_test_attributes' 31 32 33def parse_options(): 34 """Parse command line inputs. 35 36 @return: Options to run the script. 37 """ 38 parser = argparse.ArgumentParser() 39 parser.add_argument('--recreate_test_attributes', 40 action="store_true", 41 default=False, 42 help=('Delete and recreate table tko_test_attributes.' 43 'Please use it MANUALLY with CAREFULNESS & make' 44 'sure the table is properly created back.')) 45 parser.add_argument('-l', '--logfile', type=str, 46 default=None, 47 help='Path to the log file to save logs.') 48 return parser.parse_args() 49 50 51def _recreate_test_attributes(server, user, password, database): 52 """Drop & recreate the table tko_test_attributes.""" 53 table_schema = utils.run_sql_cmd( 54 server, user, password, 55 'SHOW CREATE TABLE %s\G' % RECREATE_TABLE, database) 56 logging.info(table_schema) 57 # Format executable command for creating table. 58 create_table_cmd = table_schema.split('Create Table: ')[1] 59 create_table_cmd = create_table_cmd.replace('`', '').replace('\n', '') 60 utils.run_sql_cmd(server, user, password, 61 'DROP TABLE IF EXISTS %s' % RECREATE_TABLE, database) 62 utils.run_sql_cmd(server, user, password, create_table_cmd, database) 63 64 65def main(): 66 """Main script.""" 67 options = parse_options() 68 log_config = logging_config.LoggingConfig() 69 if options.logfile: 70 log_config.add_file_handler( 71 file_path=os.path.abspath(options.logfile), level=logging.DEBUG) 72 73 with ts_mon_config.SetupTsMonGlobalState(service_name='cleanup_tko_db', 74 indirect=True): 75 server = CONFIG.get_config_value( 76 'AUTOTEST_WEB', 'global_db_host', 77 default=CONFIG.get_config_value('AUTOTEST_WEB', 'host')) 78 user = CONFIG.get_config_value( 79 'AUTOTEST_WEB', 'global_db_user', 80 default=CONFIG.get_config_value('AUTOTEST_WEB', 'user')) 81 password = CONFIG.get_config_value( 82 'AUTOTEST_WEB', 'global_db_password', 83 default=CONFIG.get_config_value('AUTOTEST_WEB', 'password')) 84 database = CONFIG.get_config_value( 85 'AUTOTEST_WEB', 'global_db_database', 86 default=CONFIG.get_config_value('AUTOTEST_WEB', 'database')) 87 88 logging.info('Starting cleaning up old records in TKO database %s on ' 89 'server %s.', database, server) 90 91 start_time = time.time() 92 try: 93 if options.recreate_test_attributes: 94 with metrics.SecondsTimer(RECREATE_TEST_ATTRIBUTES_METRIC, 95 fields={'success': False}) as fields: 96 _recreate_test_attributes(server, user, password, database) 97 fields['success'] = True 98 else: 99 with metrics.SecondsTimer(CLEANUP_METRIC, 100 fields={'success': False}) as fields: 101 utils.run_sql_cmd(server, user, password, CLEANUP_TKO_CMD, 102 database) 103 fields['success'] = True 104 except: 105 logging.exception('Cleanup failed with exception.') 106 finally: 107 duration = time.time() - start_time 108 logging.info('Cleanup attempt finished in %s seconds.', duration) 109 110 111if __name__ == '__main__': 112 main() 113