1#!/usr/bin/env python
2
3
4""" Trigger a Cluster Telemetry job with the given Lua script. """
5
6
7import argparse
8import base64
9import getpass
10import httplib2
11import json
12import subprocess
13import urllib
14
15
16CT_URL = 'https://skia-tree-status.appspot.com/skia-telemetry/'
17CT_ADD_LUA_TASK_URL = CT_URL + 'add_lua_task'
18CT_GET_SKP_REPOS_URL = CT_URL + 'get_skp_repos'
19CT_PENDING_TASKS_URL = CT_URL + 'pending_tasks'
20POST_DATA = ('username=%s'
21             '&password=%s'
22             '&description=%s'
23             '&lua_script=%s'
24             '&pagesets_type_and_chromium_build=%s')
25
26
27def trigger_ct_run(user, password, description, script, skp_repo,
28                   aggregator=None):
29  """Trigger a Cluster Telemetry run of the given script."""
30  with open(script) as f:
31    script_contents = urllib.quote(base64.b64encode(f.read()))
32
33  body = POST_DATA % (user, password, description, script_contents, skp_repo)
34
35  if aggregator:
36    with open(aggregator) as f:
37      body += '&lua_aggregator=%s' % urllib.quote(base64.b64encode(f.read()))
38
39  resp, content = httplib2.Http().request(
40      CT_ADD_LUA_TASK_URL, 'POST', body=body)
41  if resp['status'] != '200':
42    raise Exception(
43        'Failed to trigger Cluster Telemetry job: (%s): %s' % (
44            resp['status'], content))
45
46
47def parse_args():
48  """Parse command-line flags and obtain any additional information."""
49  parser = argparse.ArgumentParser(
50      description='Trigger a Cluster Telemetry job with the given Lua script.')
51  parser.add_argument('--script', help='Lua script to run', required=True)
52  parser.add_argument('--aggregator', help='Aggregator script')
53  parser.add_argument('--description', help='Description of the job.')
54  parser.add_argument('--email',
55                      help=('Email address to send results. If not specified, '
56                            'the value of `git config user.email` is used.'))
57  parser.add_argument('--password_file',
58                      help=('File in which the CT password is stored. Will '
59                            'prompt for password if not specified.'))
60  parser.add_argument('--skp_repo', default='10k',
61                      help='Which set of SKPs to use, eg. "10k", "All"')
62  args = parser.parse_args()
63
64  # If the user provided their email address, use that. Otherwise obtain it
65  # from the Git config.
66  user = args.email
67  if not user:
68    user = subprocess.check_output(['git', 'config', 'user.email']).rstrip()
69
70  # Read the password from the password file, if provided, otherwise prompt.
71  if args.password_file:
72    with open(args.password_file) as f:
73      password = f.read().rstrip()
74  else:
75    password = getpass.getpass(
76        'Enter the skia_status_password '
77        '(on https://valentine.corp.google.com/): ')
78
79  # Find an SKP repo to use.
80  resp, content = httplib2.Http().request(CT_GET_SKP_REPOS_URL, "GET")
81  if resp['status'] != '200':
82    raise Exception('Failed to obtain SKP repos from %s' % CT_GET_SKP_REPOS_URL)
83  skp_repos = json.loads(content)
84  chosen_skp_repo = skp_repos.get(args.skp_repo)[0]
85  if not chosen_skp_repo:
86    raise Exception('No generated SKPs exist for "%s"' % args.skp_repo)
87  skp_repo = '-'.join((args.skp_repo,
88                       chosen_skp_repo[0],
89                       chosen_skp_repo[1]))
90
91  return (user, password, args.description, args.script, skp_repo,
92          args.aggregator)
93
94
95def main():
96  user, password, description, script, skp_repo, aggregator = parse_args()
97  trigger_ct_run(user, password, description, script, skp_repo, aggregator)
98  print ('Successfully triggered Cluster Telemetry job. View the queue at %s' %
99         CT_PENDING_TASKS_URL)
100
101
102if __name__ == '__main__':
103  main()
104