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

"""Run the query using MySQL client and report the runtime.

Report both the runtime as reported by MySQL, which won't include overhead
like time spent waiting for password entry, query parsing, etc but has only
second precision, and the total running time of the subprocess, which has
higher resolution but includes the aforementioned variables.
"""


import subprocess
import sys
import tempfile
import time


FILE_CONTENTS = '''SET @start=UNIX_TIMESTAMP();
select benchmark(%d, (select count(*) from (%s) as r));
SELECT UNIX_TIMESTAMP()-@start AS seconds;
'''


def mysql_bench(repeat, query, *args):
  """Run the query using MySQL client and report the runtime."""
  with tempfile.NamedTemporaryFile() as f:
    f.write(FILE_CONTENTS % (repeat, query))
    f.flush()
    cmd = ['mysql', '-e', 'source %s' % f.name, '-s']
    cmd.extend(*args)
    start = time.time()
    t = int(subprocess.check_output(cmd).splitlines()[-1])
    print 'Elapsed: %d secs (%f sec subprocess runtime) ' % (t, time.time()-start)


def usage():
  """Print the usage string and exit."""
  print >> sys.stderr, 'USAGE: %s <repeat-count> "SELECT * FROM ..." [args to mysql client]'
  sys.exit(1)


if __name__ == '__main__':
  if len(sys.argv) < 3:
    usage()
  try:
    repeats = int(sys.argv[1])
  except:
    usage()
  mysql_bench(repeats, sys.argv[2], sys.argv[3:])