contrib/backfill_tko_task_references.py

"""A simple script to backfill tko_task_references table with throttling."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import collections
import contextlib
import logging
import time

import MySQLdb


class BackfillException(Exception):
  pass


def _parse_args():
  parser = argparse.ArgumentParser(
      description=__doc__)
  parser.add_argument('--host', required=True, help='mysql server host')
  parser.add_argument('--user', required=True, help='mysql server user')
  parser.add_argument('--password', required=True, help='mysql server password')
  parser.add_argument('--dryrun', action='store_true', default=False)
  parser.add_argument(
      '--num-iterations',
      default=None,
      type=int,
      help='If set, total number of iterations. Default is no limit.',
  )
  parser.add_argument(
      '--batch-size',
      default=1000,
      help='Number of tko_jobs rows to read in one iteration',
  )
  parser.add_argument(
      '--sleep-seconds',
      type=int,
      default=1,
      help='Time to sleep between iterations',
  )

  args = parser.parse_args()
  if args.dryrun:
    if not args.num_iterations:
      logging.info('DRYRUN: Limiting to 5 iterations in dryrun mode.')
      args.num_iterations = 5
  return args



@contextlib.contextmanager
def _mysql_connection(args):
  conn = MySQLdb.connect(user=args.user, host=args.host, passwd=args.password)
  with _mysql_cursor(conn) as c:
    c.execute('USE chromeos_autotest_db;')
  try:
    yield conn
  finally:
    conn.close()


@contextlib.contextmanager
def _autocommit(conn):
  try:
    yield conn
  except:
    conn.rollback()
  else:
    conn.commit()


@contextlib.contextmanager
def _mysql_cursor(conn):
  c = conn.cursor()
  try:
    yield c
  finally:
    c.close()


def _latest_unfilled_job_idx(conn):
  with _mysql_cursor(conn) as c:
    c.execute("""
SELECT tko_job_idx
FROM tko_task_references
ORDER BY tko_job_idx
LIMIT 1
;""")
    r = c.fetchall()
    if r:
      return str(long(r[0][0]) - 1)
  logging.debug('tko_task_references is empty.'
               ' Grabbing the latest tko_job_idx to fill.')
  with _mysql_cursor(conn) as c:
    c.execute("""
SELECT job_idx
FROM tko_jobs
ORDER BY job_idx DESC
LIMIT 1
;""")
    r = c.fetchall()
    if r:
      return r[0][0]
  return None


_TKOTaskReference = collections.namedtuple(
    '_TKOTaskReference',
    ['tko_job_idx', 'task_reference', 'parent_task_reference'],
)

_SQL_SELECT_TASK_REFERENCES = """
SELECT job_idx, afe_job_id, afe_parent_job_id
FROM tko_jobs
WHERE job_idx <= %(latest_job_idx)s
ORDER BY job_idx DESC
LIMIT %(batch_size)s
;"""
_SQL_INSERT_TASK_REFERENCES = """
INSERT INTO tko_task_references(reference_type, tko_job_idx, task_id, parent_task_id)
VALUES %(values)s
;"""
_SQL_SELECT_TASK_REFERENCE = """
SELECT tko_job_idx FROM tko_task_references WHERE tko_job_idx = %(tko_job_idx)s
;"""


def _compute_task_references(conn, latest_job_idx, batch_size):
  with _mysql_cursor(conn) as c:
    sql = _SQL_SELECT_TASK_REFERENCES % {
        'latest_job_idx': latest_job_idx,
        'batch_size': batch_size,
    }
    c.execute(sql)
    rs = c.fetchall()
    if rs is None:
      return []

    return [_TKOTaskReference(r[0], r[1], r[2]) for r in rs]


def _insert_task_references(conn, task_references, dryrun):
  values = ', '.join([
      '("afe", %s, "%s", "%s")' %
      (tr.tko_job_idx, tr.task_reference, tr.parent_task_reference)
      for tr in task_references
  ])
  sql = _SQL_INSERT_TASK_REFERENCES % {'values': values}
  if dryrun:
    if len(sql) < 200:
      sql_log = sql
    else:
      sql_log = '%s... [SNIP] ...%s' % (sql[:150], sql[-49:])
    logging.debug('Would have run: %s', sql_log)
  with _autocommit(conn) as conn:
    with _mysql_cursor(conn) as c:
      c.execute(sql)


def _verify_task_references(conn, task_references):
  # Just verify that the last one was inserted.
  if not task_references:
    return
  tko_job_idx = task_references[-1].tko_job_idx
  sql = _SQL_SELECT_TASK_REFERENCE % {'tko_job_idx': tko_job_idx}
  with _mysql_cursor(conn) as c:
    c.execute(sql)
    r = c.fetchall()
    if not r or r[0][0] != tko_job_idx:
      raise BackfillException(
          'Failed to insert task reference for tko_job_id %s' % tko_job_idx)


def _next_job_idx(task_references):
  return str(long(task_references[-1].tko_job_idx) - 1)

def main():
  logging.basicConfig(level=logging.DEBUG)
  args = _parse_args()
  with _mysql_connection(args) as conn:
    tko_job_idx = _latest_unfilled_job_idx(conn)
    if tko_job_idx is None:
      raise BackfillException('Failed to get last unfilled tko_job_idx')
    logging.info('First tko_job_idx to fill: %s', tko_job_idx)

  while True:
    logging.info('####################################')
    logging.info('Start backfilling from tko_job_idx: %s', tko_job_idx)

    task_references = ()
    with _mysql_connection(args) as conn:
      task_references = _compute_task_references(
          conn, tko_job_idx, args.batch_size)
    if not task_references:
      logging.info('No more unfilled task references. All done!')
      break

    logging.info(
        'Inserting %d task references. tko_job_ids: %d...%d',
        len(task_references),
        task_references[0].tko_job_idx,
        task_references[-1].tko_job_idx,
    )
    with _mysql_connection(args) as conn:
      _insert_task_references(conn, task_references, args.dryrun)
    if not args.dryrun:
      with _mysql_connection(args) as conn:
        _verify_task_references(conn, task_references)

    tko_job_idx = _next_job_idx(task_references)

    if args.num_iterations is not None:
      args.num_iterations -= 1
      if args.num_iterations <= 0:
        break
      logging.info('%d more iterations left', args.num_iterations)
    logging.info('Iteration done. Sleeping for %d seconds', args.sleep_seconds)
    time.sleep(args.sleep_seconds)


if __name__ == '__main__':
  main()