contrib/db_cleanup.py

#!/usr/bin/python
# Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import argparse
import datetime
import os
import re
import sys
import logging

os.environ['DJANGO_SETTINGS_MODULE'] = 'frontend.settings'

import common
from autotest_lib.server import utils
from django.db import connections, transaction


# Format Appears as: [Date] [Time] - [Msg Level] - [Message]
LOGGING_FORMAT = '%(asctime)s - %(levelname)s - %(message)s'
# This regex makes sure the input is in the format of YYYY-MM-DD (2012-02-01)
DATE_FORMAT_REGEX = ('^(19|20)\d\d[- /.](0[1-9]|1[012])[- /.](0[1-9]|[12][0-9]'
                     '|3[01])$')
SELECT_CMD_FORMAT = """
SELECT %(table)s.%(primary_key)s FROM %(table)s
WHERE %(table)s.%(time_column)s <= "%(date)s"
"""
SELECT_JOIN_CMD_FORMAT = """
SELECT %(table)s.%(primary_key)s FROM %(table)s
INNER JOIN %(related_table)s
  ON %(table)s.%(foreign_key)s=%(related_table)s.%(related_primary_key)s
WHERE %(related_table)s.%(time_column)s <= "%(date)s"
"""
SELECT_WITH_INDIRECTION_FORMAT = """
SELECT %(table)s.%(primary_key)s FROM %(table)s
INNER JOIN %(indirection_table)s
  ON %(table)s.%(foreign_key)s =
     %(indirection_table)s.%(indirection_primary_key)s
INNER JOIN %(related_table)s
  ON %(indirection_table)s.%(indirection_foreign_key)s =
  %(related_table)s.%(related_primary_key)s
WHERE %(related_table)s.%(time_column)s <= "%(date)s"
"""
DELETE_ROWS_FORMAT = """
DELETE FROM %(table)s
WHERE %(table)s.%(primary_key)s IN (%(rows)s)
"""


AFE_JOB_ID = 'afe_job_id'
JOB_ID = 'job_id'
JOB_IDX = 'job_idx'
TEST_IDX = 'test_idx'

# CAUTION: Make sure only the 'default' connection is used. Otherwise
# db_cleanup may delete stuff from the global database, which is generally not
# intended.
cursor = connections['default'].cursor()

STEP_SIZE = None  # Threading this through properly is disgusting.

class ProgressBar(object):
    TEXT = "{:<40s} [{:<20s}] ({:>9d}/{:>9d})"

    def __init__(self, name, amount):
        self._name = name
        self._amount = amount
        self._cur = 0

    def __enter__(self):
        return self

    def __exit__(self, a, b, c):
        sys.stdout.write('\n')
        sys.stdout.flush()

    def update(self, x):
        """
        Advance the counter by `x`.

        @param x: An integer of how many more elements were processed.
        """
        self._cur += x

    def show(self):
        """
        Display the progress bar on the current line.  Repeated invocations
        "update" the display.
        """
        if self._amount == 0:
            barlen = 20
        else:
            barlen = int(20 * self._cur / float(self._amount))
        if barlen:
            bartext = '=' * (barlen-1) + '>'
        else:
            bartext = ''
        text = self.TEXT.format(self._name, bartext, self._cur, self._amount)
        sys.stdout.write('\r')
        sys.stdout.write(text)
        sys.stdout.flush()


def grouper(iterable, n):
    """
    Group the elements of `iterable` into groups of maximum size `n`.

    @param iterable: An iterable.
    @param n: Max size of returned groups.
    @returns: Yields iterables of size <= n.

    >>> grouper('ABCDEFG', 3)
    [['A', 'B', C'], ['D', 'E', 'F'], ['G']]
    """
    args = [iter(iterable)] * n
    while True:
        lst = []
        try:
            for itr in args:
                lst.append(next(itr))
            yield lst
        except StopIteration:
            if lst:
                yield lst
            break


def _delete_table_data_before_date(table_to_delete_from, primary_key,
                                   related_table, related_primary_key,
                                   date, foreign_key=None,
                                   time_column="started_time",
                                   indirection_table=None,
                                   indirection_primary_key=None,
                                   indirection_foreign_key=None):
    """
    We want a delete statement that will only delete from one table while
    using a related table to find the rows to delete.

    An example mysql command:
    DELETE FROM tko_iteration_result USING tko_iteration_result INNER JOIN
    tko_tests WHERE tko_iteration_result.test_idx=tko_tests.test_idx AND
    tko_tests.started_time <= '2012-02-01';

    There are also tables that require 2 joins to determine which rows we want
    to delete and we determine these rows by joining the table we want to
    delete from with an indirection table to the actual jobs table.

    @param table_to_delete_from: Table whose rows we want to delete.
    @param related_table: Table with the date information we are selecting by.
    @param foreign_key: Foreign key used in table_to_delete_from to reference
                        the related table. If None, the primary_key is used.
    @param primary_key: Primary key in the related table.
    @param date: End date of the information we are trying to delete.
    @param time_column: Column that we want to use to compare the date to.
    @param indirection_table: Table we use to link the data we are trying to
                              delete with the table with the date information.
    @param indirection_primary_key: Key we use to connect the indirection table
                                    to the table we are trying to delete rows
                                    from.
    @param indirection_foreign_key: Key we use to connect the indirection table
                                    to the table with the date information.
    """
    if not foreign_key:
        foreign_key = primary_key

    if not related_table:
        # Deleting from a table directly.
        variables = dict(table=table_to_delete_from, primary_key=primary_key,
                         time_column=time_column, date=date)
        sql = SELECT_CMD_FORMAT % variables
    elif not indirection_table:
        # Deleting using a single JOIN to get the date information.
        variables = dict(primary_key=primary_key, table=table_to_delete_from,
                         foreign_key=foreign_key, related_table=related_table,
                         related_primary_key=related_primary_key,
                         time_column=time_column, date=date)
        sql = SELECT_JOIN_CMD_FORMAT % variables
    else:
        # There are cases where we need to JOIN 3 TABLES to determine the rows
        # we want to delete.
        variables = dict(primary_key=primary_key, table=table_to_delete_from,
                         indirection_table=indirection_table,
                         foreign_key=foreign_key,
                         indirection_primary_key=indirection_primary_key,
                         related_table=related_table,
                         related_primary_key=related_primary_key,
                         indirection_foreign_key=indirection_foreign_key,
                         time_column=time_column, date=date)
        sql = SELECT_WITH_INDIRECTION_FORMAT % variables

    logging.debug('SQL: %s', sql)
    cursor.execute(sql, [])
    rows = [x[0] for x in cursor.fetchall()]
    logging.debug(rows)

    if not rows or rows == [None]:
        with ProgressBar(table_to_delete_from, 0) as pb:
            pb.show()
        logging.debug('Noting to delete for %s', table_to_delete_from)
        return

    with ProgressBar(table_to_delete_from, len(rows)) as pb:
        for row_keys in grouper(rows, STEP_SIZE):
            variables['rows'] = ','.join([str(x) for x in row_keys])
            sql = DELETE_ROWS_FORMAT % variables
            logging.debug('SQL: %s', sql)
            cursor.execute(sql, [])
            transaction.commit_unless_managed(using='default')
            pb.update(len(row_keys))
            pb.show()


def _subtract_days(date, days_to_subtract):
    """
    Return a date (string) that is 'days' before 'date'

    @param date: date (string) we are subtracting from.
    @param days_to_subtract: days (int) we are subtracting.
    """
    date_obj = datetime.datetime.strptime(date, '%Y-%m-%d')
    difference = date_obj - datetime.timedelta(days=days_to_subtract)
    return difference.strftime('%Y-%m-%d')


def _delete_all_data_before_date(date):
    """
    Delete all the database data before a given date.

    This function focuses predominately on the data for jobs in tko_jobs.
    However not all jobs in afe_jobs are also in tko_jobs.

    Therefore we delete all the afe_job and foreign key relations prior to two
    days before date. Then we do the queries using tko_jobs and these
    tables to ensure all the related information is gone. Even though we are
    repeating deletes on these tables, the second delete will be quick and
    completely thorough in ensuring we clean up all the foreign key
    dependencies correctly.

    @param date: End date of the information we are trying to delete.
    @param step: Rows to delete per SQL query.
    """
    # First cleanup all afe_job related data (prior to 2 days before date).
    # The reason for this is not all afe_jobs may be in tko_jobs.
    afe_date = _subtract_days(date, 2)
    logging.info('Cleaning up all afe_job data prior to %s.', afe_date)
    _delete_table_data_before_date('afe_aborted_host_queue_entries',
                                   'queue_entry_id',
                                   'afe_jobs', 'id', afe_date,
                                   time_column= 'created_on',
                                   foreign_key='queue_entry_id',
                                   indirection_table='afe_host_queue_entries',
                                   indirection_primary_key='id',
                                   indirection_foreign_key='job_id')
    _delete_table_data_before_date('afe_special_tasks', 'id',
                                   'afe_jobs', 'id',
                                   afe_date, time_column='created_on',
                                   foreign_key='queue_entry_id',
                                   indirection_table='afe_host_queue_entries',
                                   indirection_primary_key='id',
                                   indirection_foreign_key='job_id')
    _delete_table_data_before_date('afe_host_queue_entries', 'id',
                                   'afe_jobs', 'id',
                                   afe_date, time_column='created_on',
                                   foreign_key=JOB_ID)
    _delete_table_data_before_date('afe_job_keyvals', 'id',
                                   'afe_jobs', 'id',
                                   afe_date, time_column='created_on',
                                   foreign_key=JOB_ID)
    _delete_table_data_before_date('afe_jobs_dependency_labels', 'id',
                                   'afe_jobs', 'id',
                                   afe_date, time_column='created_on',
                                   foreign_key=JOB_ID)
    _delete_table_data_before_date('afe_jobs', 'id',
                                   None, None,
                                   afe_date, time_column='created_on')
    # Special tasks that aren't associated with an HQE
    # Since we don't do the queue_entry_id=NULL check, we might wipe out a bit
    # more than we should, but I doubt anyone will notice or care.
    _delete_table_data_before_date('afe_special_tasks', 'id',
                                   None, None,
                                   afe_date, time_column='time_requested')

    # Now go through and clean up all the rows related to tko_jobs prior to
    # date.
    logging.info('Cleaning up all data related to tko_jobs prior to %s.',
                  date)
    _delete_table_data_before_date('tko_test_attributes', 'id',
                                   'tko_tests', TEST_IDX,
                                   date, foreign_key=TEST_IDX)
    _delete_table_data_before_date('tko_test_labels_tests', 'id',
                                   'tko_tests', TEST_IDX,
                                   date, foreign_key= 'test_id')
    _delete_table_data_before_date('tko_iteration_result', TEST_IDX,
                                   'tko_tests', TEST_IDX,
                                   date)
    _delete_table_data_before_date('tko_iteration_perf_value', TEST_IDX,
                                   'tko_tests', TEST_IDX,
                                   date)
    _delete_table_data_before_date('tko_iteration_attributes', TEST_IDX,
                                   'tko_tests', TEST_IDX,
                                   date)
    _delete_table_data_before_date('tko_job_keyvals', 'id',
                                   'tko_jobs', JOB_IDX,
                                   date, foreign_key='job_id')
    _delete_table_data_before_date('afe_aborted_host_queue_entries',
                                   'queue_entry_id',
                                   'tko_jobs', AFE_JOB_ID, date,
                                   foreign_key='queue_entry_id',
                                   indirection_table='afe_host_queue_entries',
                                   indirection_primary_key='id',
                                   indirection_foreign_key='job_id')
    _delete_table_data_before_date('afe_special_tasks', 'id',
                                   'tko_jobs', AFE_JOB_ID,
                                   date, foreign_key='queue_entry_id',
                                   indirection_table='afe_host_queue_entries',
                                   indirection_primary_key='id',
                                   indirection_foreign_key='job_id')
    _delete_table_data_before_date('afe_host_queue_entries', 'id',
                                   'tko_jobs', AFE_JOB_ID,
                                   date, foreign_key='job_id')
    _delete_table_data_before_date('afe_job_keyvals', 'id',
                                   'tko_jobs', AFE_JOB_ID,
                                   date, foreign_key='job_id')
    _delete_table_data_before_date('afe_jobs_dependency_labels', 'id',
                                   'tko_jobs', AFE_JOB_ID,
                                   date, foreign_key='job_id')
    _delete_table_data_before_date('afe_jobs', 'id',
                                   'tko_jobs', AFE_JOB_ID,
                                   date, foreign_key='id')
    _delete_table_data_before_date('tko_tests', TEST_IDX,
                                   'tko_jobs', JOB_IDX,
                                   date, foreign_key=JOB_IDX)
    _delete_table_data_before_date('tko_jobs', JOB_IDX,
                                   None, None, date)


def parse_args():
    """Parse command line arguments"""
    parser = argparse.ArgumentParser()
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='Print SQL commands and results')
    parser.add_argument('--step', type=int, action='store',
                        default=1000,
                        help='Number of rows to delete at once')
    parser.add_argument('-c', '--check_server', action='store_true',
                        help='Check if the server should run db clean up.')
    parser.add_argument('date', help='Keep results newer than')
    return parser.parse_args()


def should_cleanup():
    """Check if the server should run db_cleanup.

    Only shard should clean up db.

    @returns: True if it should run db cleanup otherwise False.
    """
    return utils.is_shard()


def main():
    """main"""
    args = parse_args()

    level = logging.DEBUG if args.verbose else logging.INFO
    logging.basicConfig(level=level, format=LOGGING_FORMAT)
    logging.info('Calling: %s', sys.argv)

    if not re.match(DATE_FORMAT_REGEX, args.date):
        print 'DATE must be in yyyy-mm-dd format!'
        return
    if args.check_server and not should_cleanup():
        print 'Only shard can run db cleanup.'
        return

    global STEP_SIZE
    STEP_SIZE = args.step
    _delete_all_data_before_date(args.date)


if __name__ == '__main__':
    main()