Home | History | Annotate | Download | only in site_utils 
      1 #!/usr/bin/python
      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 
      9 import argparse
     10 import logging
     11 import os
     12 import time
     13 
     14 import common
     15 from autotest_lib.client.bin import utils
     16 from autotest_lib.client.common_lib import global_config
     17 from autotest_lib.client.common_lib import logging_config
     18 
     19 from chromite.lib import metrics
     20 from chromite.lib import ts_mon_config
     21 
     22 
     23 CONFIG = global_config.global_config
     24 
     25 # SQL command to remove old test results in TKO database.
     26 CLEANUP_TKO_CMD = 'call remove_old_tests_sp()'
     27 CLEANUP_METRIC = 'chromeos/autotest/tko/cleanup_duration'
     28 RECREATE_TEST_ATTRIBUTES_METRIC = (
     29     'chromeos/autotest/tko/recreate_test_attributes')
     30 RECREATE_TABLE = 'tko_test_attributes'
     31 
     32 
     33 def 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 
     51 def _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 
     65 def 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 
    111 if __name__ == '__main__':
    112     main()
    113