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