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)
+
+  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

[Vadim Sh., 2017/06/05 21:33:13]

instead of returning data, you can yield data['items'] from yielder before doing
'while' loop there

In fact, I don't understand why this function returns a generator instead of
being a generator itself

def fetch_with_cursor(base_url, limit):
  ...
  yield data['items']
  ...
  while cursor and (....):
     ...
     yield items


At even closer look, it doesn't have to be generator at all. Just add 'progress'
argument and move the entire fetching loop here, instead of doing it in
individual subcommands.

[M-A Ruel, 2017/06/07 15:36:03]

It's because some queries may not return 'items'. For example doing
swarming.py query task/1234/result

I added a more specific example in CMDquery and skipped the function when no
cursor or no items is provided.

+
+
 ### 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'
+  if options.dead_only:
+    url += '?is_dead=TRUE'
+  elif options.keep_dead:
+    url += '?is_dead=NONE'
+  else:
+    url += '?is_dead=FALSE'
 
+  if options.busy:
+    url += '&is_busy=TRUE'
+  elif options.idle:
+    url += '&is_busy=FALSE'
+  else:
+    url += '&is_busy=NONE'
+
+  if options.mp:
+    url += '&is_mp=TRUE'
+  elif options.non_mp:
+    url += '&is_mp=FALSE'
+  else:
+    url += '&is_mp=NONE'
+
+  for key, value in options.dimensions:
+    url += '&dimensions=%s:%s' % (key, value)

[Vadim Sh., 2017/06/05 21:33:13]

this needs url encoding

[M-A Ruel, 2017/06/07 15:36:03]

Argh, thanks. Changed to encode everything systematically.

+  try:
+    data, yielder = get_yielder(url, 0)
+    bots = data.get('items') or []
+    for items in yielder():
+      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 98 matching lines @@
       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 with tags os:Ubuntu-12.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'
   """
-  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():
+      # Some items support cursors. Try to get automatically if cursors are

[Vadim Sh., 2017/06/05 21:33:13]

this comment doesn't seem relevant any more

[M-A Ruel, 2017/06/07 15:36:03]

Done.

+      # needed by looking at the 'cursor' items.
+      data['items'].extend(items)
+      if options.progress:
+        sys.stdout.write('.')
+        sys.stdout.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.stdout.write('\n')
+    sys.stdout.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)

[Vadim Sh., 2017/06/05 21:33:13]

this progress output will screw up json

consider emitting it to stderr

[M-A Ruel, 2017/06/07 15:36:03]

Done.

       sys.stdout.write('\n')
     except IOError:
       pass
   return 0
 
 
 def CMDquery_list(parser, args):
   """Returns list of all the Swarming APIs that can be used with command
   'query'.
   """
@@ skipping 312 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:]))