client: Support repeated keys in task request.

client/swarming.py

 #!/usr/bin/env python
 # Copyright 2013 The LUCI Authors. All rights reserved.
 # Use of this source code is governed under the Apache License, Version 2.0
 # that can be found in the LICENSE file.
 
 """Client tool to trigger tasks or retrieve results from a Swarming server."""
 
 __version__ = '0.9.1'
 
 import collections
@@ skipping 37 matching lines @@
 class Failure(Exception):
   """Generic failure."""
   pass
 
 
 def default_task_name(options):
   """Returns a default task name if not specified."""
   if not options.task_name:
     task_name = u'%s/%s' % (
         options.user,
-        '_'.join(
-            '%s=%s' % (k, v)
-            for k, v in sorted(options.dimensions.iteritems())))
+        '_'.join('%s=%s' % (k, v) for k, v in options.dimensions))
     if options.isolated:
       task_name += u'/' + options.isolated
     return task_name
   return options.task_name
 
 
 ### Triggering.
 
 
 # See ../appengine/swarming/swarming_rpcs.py.
@@ skipping 84 matching lines @@
   This is for the v1 client Swarming API.
   """
   out = namedtuple_to_dict(task_request)
   if hide_token:
     if out['service_account_token'] not in (None, 'bot', 'none'):
       out['service_account_token'] = '<hidden>'
   # Don't send 'service_account_token' if it is None to avoid confusing older
   # version of the server that doesn't know about 'service_account_token'.
   if out['service_account_token'] in (None, 'none'):
     out.pop('service_account_token')
-  # Maps are not supported until protobuf v3.
   out['properties']['dimensions'] = [
     {'key': k, 'value': v}
-    for k, v in out['properties']['dimensions'].iteritems()
+    for k, v in out['properties']['dimensions']
   ]
-  out['properties']['dimensions'].sort(key=lambda x: x['key'])
   out['properties']['env'] = [
     {'key': k, 'value': v}
     for k, v in out['properties']['env'].iteritems()
   ]
   out['properties']['env'].sort(key=lambda x: x['key'])
   return out
 
 
 def swarming_trigger(swarming, raw_request):
   """Triggers a request on the Swarming server and returns the json data.
@@ skipping 627 matching lines @@
     # URL is of the following form:
     # url = host + (
     #   '/_ah/api/discovery/v1/apis/%s/%s/rest' % (api['id'], api['version'])
     api_data = net.url_read_json(api['discoveryRestUrl'])
     if api_data is None:
       raise APIError('Failed to discover %s on %s' % (api['id'], host))
     out[api['id']] = api_data
   return out
 
 
+def get_yielder(base_url, limit):
+  """Returns the first query and a function that yields following items."""
+  CHUNK_SIZE = 250
+
+  url = base_url
+  if limit:
+    url += '%slimit=%d' % ('&' if '?' in url else '?', min(CHUNK_SIZE, limit))
+  data = net.url_read_json(url)
+  if data is None:
+    # TODO(maruel): Do basic diagnostic.
+    raise Failure('Failed to access %s' % url)
+  org_cursor = data.pop('cursor', None)
+  org_total = len(data.get('items') or [])
+  logging.info('get_yielder(%s) returning %d items', base_url, org_total)
+  if not org_cursor or not org_total:
+    # This is not an iterable resource.
+    return data, lambda: []
+
+  def yielder():
+    cursor = org_cursor
+    total = org_total
+    # Some items support cursors. Try to get automatically if cursors are needed
+    # by looking at the 'cursor' items.
+    while cursor and (not limit or total < limit):
+      merge_char = '&' if '?' in base_url else '?'
+      url = base_url + '%scursor=%s' % (merge_char, urllib.quote(cursor))
+      if limit:
+        url += '&limit=%d' % min(CHUNK_SIZE, limit - total)
+      new = net.url_read_json(url)
+      if new is None:
+        raise Failure('Failed to access %s' % url)
+      cursor = new.get('cursor')
+      new_items = new.get('items')
+      nb_items = len(new_items or [])
+      total += nb_items
+      logging.info('get_yielder(%s) yielding %d items', base_url, nb_items)
+      yield new_items
+
+  return data, yielder
+
+
 ### Commands.
 
 
 def abort_task(_swarming, _manifest):
   """Given a task manifest that was triggered, aborts its execution."""
   # TODO(vadimsh): No supported by the server yet.
 
 
 def add_filter_options(parser):
   parser.filter_group = optparse.OptionGroup(parser, 'Bot selection')
   parser.filter_group.add_option(
       '-d', '--dimension', default=[], action='append', nargs=2,
       dest='dimensions', metavar='FOO bar',
       help='dimension to filter on')
   parser.add_option_group(parser.filter_group)
 
 
+def process_filter_options(parser, options):
+  for key, value in options.dimensions:
+    if ':' in key:
+      parser.error('--dimension key cannot contain ":"')
+    if key.strip() != key:
+      parser.error('--dimension key has whitespace')
+    if not key:
+      parser.error('--dimension key is empty')
+
+    if value.strip() != value:
+      parser.error('--dimension value has whitespace')
+    if not value:
+      parser.error('--dimension value is empty')
+  options.dimensions.sort()
+
+
 def add_sharding_options(parser):
   parser.sharding_group = optparse.OptionGroup(parser, 'Sharding options')
   parser.sharding_group.add_option(
       '--shards', type='int', default=1, metavar='NUMBER',
       help='Number of shards to trigger and collect.')
   parser.add_option_group(parser.sharding_group)
 
 
 def add_trigger_options(parser):
   """Adds all options to trigger a task on Swarming."""
@@ skipping 71 matching lines @@
       '--deadline', type='int', dest='expiration',
       help=optparse.SUPPRESS_HELP)
   parser.add_option_group(group)
 
 
 def process_trigger_options(parser, options, args):
   """Processes trigger options and does preparatory steps.
 
   Generates service account tokens if necessary.
   """
-  options.dimensions = dict(options.dimensions)
+  process_filter_options(parser, options)
   options.env = dict(options.env)
   if args and args[0] == '--':
     args = args[1:]
 
   if not options.dimensions:
     parser.error('Please at least specify one --dimension')
   if not all(len(t.split(':', 1)) == 2 for t in options.tags):
     parser.error('--tags must be in the format key:value')
   if options.raw_cmd and not args:
     parser.error(
@@ skipping 145 matching lines @@
       print('Deleting %s failed. Probably already gone' % bot)
       result = 1
   return result
 
 
 def CMDbots(parser, args):
   """Returns information about the bots connected to the Swarming server."""
   add_filter_options(parser)
   parser.filter_group.add_option(
       '--dead-only', action='store_true',
-      help='Only print dead bots, useful to reap them and reimage broken bots')
+      help='Filter out bots alive, useful to reap them and reimage broken bots')
   parser.filter_group.add_option(
       '-k', '--keep-dead', action='store_true',
-      help='Do not filter out dead bots')
+      help='Keep both dead and alive bots')
+  parser.filter_group.add_option(
+      '--busy', action='store_true', help='Keep only busy bots')
+  parser.filter_group.add_option(
+      '--idle', action='store_true', help='Keep only idle bots')
+  parser.filter_group.add_option(
+      '--mp', action='store_true',
+      help='Keep only Machine Provider managed bots')
+  parser.filter_group.add_option(
+      '--non-mp', action='store_true',
+      help='Keep only non Machine Provider managed bots')
   parser.filter_group.add_option(
       '-b', '--bare', action='store_true',
       help='Do not print out dimensions')
   options, args = parser.parse_args(args)
+  process_filter_options(parser, options)
 
   if options.keep_dead and options.dead_only:
-    parser.error('Use only one of --keep-dead and --dead-only')
+    parser.error('Use only one of --keep-dead or --dead-only')
+  if options.busy and options.idle:
+    parser.error('Use only one of --busy or --idle')
+  if options.mp and options.non_mp:
+    parser.error('Use only one of --mp or --non-mp')
 
-  bots = []
-  cursor = None
-  limit = 250
-  # Iterate via cursors.
-  base_url = (
-      options.swarming + '/api/swarming/v1/bots/list?limit=%d' % limit)
-  while True:
-    url = base_url
-    if cursor:
-      url += '&cursor=%s' % urllib.quote(cursor)
-    data = net.url_read_json(url)
-    if data is None:
-      print >> sys.stderr, 'Failed to access %s' % options.swarming
-      return 1
-    bots.extend(data['items'])
-    cursor = data.get('cursor')
-    if not cursor:
-      break
+  url = options.swarming + '/api/swarming/v1/bots/list?'
+  values = []
+  if options.dead_only:
+    values.append(('is_dead', 'TRUE'))
+  elif options.keep_dead:
+    values.append(('is_dead', 'NONE'))
+  else:
+    values.append(('is_dead', 'FALSE'))
 
+  if options.busy:
+    values.append(('is_busy', 'TRUE'))
+  elif options.idle:
+    values.append(('is_busy', 'FALSE'))
+  else:
+    values.append(('is_busy', 'NONE'))
+
+  if options.mp:
+    values.append(('is_mp', 'TRUE'))
+  elif options.non_mp:
+    values.append(('is_mp', 'FALSE'))
+  else:
+    values.append(('is_mp', 'NONE'))
+
+  for key, value in options.dimensions:
+    values.append(('dimensions', '%s:%s' % (key, value)))
+  url += urllib.urlencode(values)
+  try:
+    data, yielder = get_yielder(url, 0)
+    bots = data.get('items') or []
+    for items in yielder():
+      if items:
+        bots.extend(items)
+  except Failure as e:
+    sys.stderr.write('\n%s\n' % e)
+    return 1
   for bot in natsort.natsorted(bots, key=lambda x: x['bot_id']):
-    if options.dead_only:
-      if not bot.get('is_dead'):
-        continue
-    elif not options.keep_dead and bot.get('is_dead'):
-      continue
-
-    # If the user requested to filter on dimensions, ensure the bot has all the
-    # dimensions requested.
-    dimensions = {i['key']: i.get('value') for i in bot.get('dimensions', {})}
-    for key, value in options.dimensions:
-      if key not in dimensions:
-        break
-      # A bot can have multiple value for a key, for example,
-      # {'os': ['Windows', 'Windows-6.1']}, so that --dimension os=Windows will
-      # be accepted.
-      if isinstance(dimensions[key], list):
-        if value not in dimensions[key]:
-          break
-      else:
-        if value != dimensions[key]:
-          break
-    else:
-      print bot['bot_id']
-      if not options.bare:
-        print '  %s' % json.dumps(dimensions, sort_keys=True)
-        if bot.get('task_id'):
-          print '  task: %s' % bot['task_id']
+    print bot['bot_id']
+    if not options.bare:
+      dimensions = {i['key']: i.get('value') for i in bot.get('dimensions', {})}
+      print '  %s' % json.dumps(dimensions, sort_keys=True)
+      if bot.get('task_id'):
+        print '  task: %s' % bot['task_id']
   return 0
 
 
 @subcommand.usage('task_id')
 def CMDcancel(parser, args):
   """Cancels a task."""
   options, args = parser.parse_args(args)
   if not args:
     parser.error('Please specify the task to cancel')
   for task_id in args:
@@ skipping 87 matching lines @@
   print data
   return 0
 
 
 @subcommand.usage('[method name]')
 def CMDquery(parser, args):
   """Returns raw JSON information via an URL endpoint. Use 'query-list' to
   gather the list of API methods from the server.
 
   Examples:
+    Raw task request and results:
+      swarming.py query -S server-url.com task/123456/request
+      swarming.py query -S server-url.com task/123456/result
+
     Listing all bots:
       swarming.py query -S server-url.com bots/list
 
-    Listing last 10 tasks on a specific bot named 'swarm1':
-      swarming.py query -S server-url.com --limit 10 bot/swarm1/tasks
+    Listing last 10 tasks on a specific bot named 'bot1':
+      swarming.py query -S server-url.com --limit 10 bot/bot1/tasks
 
-    Listing last 10 tasks with tags os:Ubuntu-12.04 and pool:Chrome. Note that
+    Listing last 10 tasks with tags os:Ubuntu-14.04 and pool:Chrome. Note that
     quoting is important!:
       swarming.py query -S server-url.com --limit 10 \\
-          'tasks/list?tags=os:Ubuntu-12.04&tags=pool:Chrome'
+          'tasks/list?tags=os:Ubuntu-14.04&tags=pool:Chrome'
   """
-  CHUNK_SIZE = 250
-
   parser.add_option(
       '-L', '--limit', type='int', default=200,
       help='Limit to enforce on limitless items (like number of tasks); '
            'default=%default')
   parser.add_option(
       '--json', help='Path to JSON output file (otherwise prints to stdout)')
   parser.add_option(
       '--progress', action='store_true',
       help='Prints a dot at each request to show progress')
   options, args = parser.parse_args(args)
   if len(args) != 1:
     parser.error(
         'Must specify only method name and optionally query args properly '
         'escaped.')
   base_url = options.swarming + '/api/swarming/v1/' + args[0]
-  url = base_url
-  if options.limit:
-    # Check check, change if not working out.
-    merge_char = '&' if '?' in url else '?'
-    url += '%slimit=%d' % (merge_char, min(CHUNK_SIZE, options.limit))
-  data = net.url_read_json(url)
-  if data is None:
-    # TODO(maruel): Do basic diagnostic.
-    print >> sys.stderr, 'Failed to access %s' % url
+  try:
+    data, yielder = get_yielder(base_url, options.limit)
+    for items in yielder():
+      if items:
+        data['items'].extend(items)
+      if options.progress:
+        sys.stderr.write('.')
+        sys.stderr.flush()
+  except Failure as e:
+    sys.stderr.write('\n%s\n' % e)
     return 1
-
-  # Some items support cursors. Try to get automatically if cursors are needed
-  # by looking at the 'cursor' items.
-  while (
-      data.get('cursor') and
-      (not options.limit or len(data['items']) < options.limit)):
-    merge_char = '&' if '?' in base_url else '?'
-    url = base_url + '%scursor=%s' % (merge_char, urllib.quote(data['cursor']))
-    if options.limit:
-      url += '&limit=%d' % min(CHUNK_SIZE, options.limit - len(data['items']))
-    if options.progress:
-      sys.stdout.write('.')
-      sys.stdout.flush()
-    new = net.url_read_json(url)
-    if new is None:
-      if options.progress:
-        print('')
-      print >> sys.stderr, 'Failed to access %s' % options.swarming
-      return 1
-    data['items'].extend(new.get('items', []))
-    data['cursor'] = new.get('cursor')
-
   if options.progress:
-    print('')
-  if options.limit and len(data.get('items', [])) > options.limit:
-    data['items'] = data['items'][:options.limit]
-  data.pop('cursor', None)
-
+    sys.stderr.write('\n')
+    sys.stderr.flush()
   if options.json:
     options.json = unicode(os.path.abspath(options.json))
     tools.write_json(options.json, data, True)
   else:
     try:
       tools.write_json(sys.stdout, data, False)
       sys.stdout.write('\n')
     except IOError:
       pass
   return 0
@@ skipping 318 matching lines @@
   dispatcher = subcommand.CommandDispatcher(__name__)
   return dispatcher.execute(OptionParserSwarming(version=__version__), args)
 
 
 if __name__ == '__main__':
   subprocess42.inhibit_os_error_reporting()
   fix_encoding.fix_encoding()
   tools.disable_buffering()
   colorama.init()
   sys.exit(main(sys.argv[1:]))