Add support for repeated keys in TaskRequest.

appengine/swarming/server/task_request.py

 # coding: utf-8
 # Copyright 2014 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.
 
 """Tasks definition.
 
 Each user request creates a new TaskRequest. The TaskRequest instance saves the
 metadata of the request, e.g. who requested it, when why, etc. It links to the
 actual data of the request in a TaskProperties. The TaskProperties represents
@@ skipping 524 matching lines @@
   # will be added to those in that directory.
   outputs = ndb.StringProperty(repeated=True, indexed=False,
       validator=_validate_output_path)
 
   # If True, the TaskRequest embedding these TaskProperties has an associated
   # SecretBytes entity.
   has_secret_bytes = ndb.BooleanProperty(default=False, indexed=False)
 
   @property
   def dimensions(self):
-    """Returns a dict(key, value)."""
+    """Returns a list of tuple(key, value)."""
     if self.dimensions_dict:
-      return self.dimensions_dict
-    return dict(i.split(u':', 1) for i in self.dimensions_flat)
+      return sorted(self.dimensions_dict.items())
+    # Data is already sorted.
+    return [tuple(i.split(u':', 1)) for i in self.dimensions_flat]
 
   @property
   def is_terminate(self):
     """If True, it is a terminate request."""
     return (
         not self.caches and
         not self.command and
-        self.dimensions.keys() == [u'id'] and
+        [i[0] for i in self.dimensions] == [u'id'] and
         not (self.inputs_ref and self.inputs_ref.isolated) and
         not self.cipd_input and
         not self.env and
         not self.execution_timeout_secs and
         not self.extra_args and
         not self.grace_period_secs and
         not self.io_timeout_secs and
         not self.idempotent and
         not self.outputs and
         not self.has_secret_bytes)
 
   def to_dict(self):
     out = super(TaskProperties, self).to_dict(
         exclude=['dimensions_dict', 'dimensions_flat'])
     out['dimensions'] = self.dimensions
     return out
 
   def _pre_put_hook(self):
     super(TaskProperties, self)._pre_put_hook()
     if self.is_terminate:
       # Most values are not valid with a terminate task. self.is_terminate
       # already check those.
       return
 
     # Dimensions.
-    if not self.dimensions:
+    if self.dimensions_dict:
+      raise datastore_errors.BadValueError(
+          u'internal error: dimensions_dict cannot be stored anymore')
+    if not self.dimensions_flat:
       raise datastore_errors.BadValueError(u'dimensions are required')
-    keys = self.dimensions.keys()
+    keys = [i.split(u':', 1)[0] for i in self.dimensions_flat]
     count_pool = sum(1 for k in keys if k == u'pool')
     count_id = sum(1 for k in keys if k == u'id')
     if not count_pool and not count_id:
       raise datastore_errors.BadValueError(
           u'At least one of \'id\' or \'pool\' must be used as dimensions')
     if count_id > 1:
       raise datastore_errors.BadValueError(
           u'\'id\' cannot be specified more than once in dimensions')
-    if len(keys) > 64:
+    if len(self.dimensions_flat) > 64:
       raise datastore_errors.BadValueError(
           'dimensions can have up to 64 entries')
     if not _is_sorted(self.dimensions_flat):
       # Better check here than silently sort since operations (like calculating
       # dimensions_hash) is done before storage, which would cause silent
       # corruption.
       raise datastore_errors.BadValueError(
           u'internal error: dimensions_flat must be sorted')
 
     # Isolated input and commands.
@@ skipping 216 matching lines @@
 ### Public API.
 
 
 def create_termination_task(bot_id, allow_high_priority):
   """Returns a task to terminate the given bot.
 
   Returns:
     TaskRequest for priority 0 (highest) termination task.
   """
   properties = TaskProperties(
-      dimensions_dict={u'id': unicode(bot_id)},
+      dimensions_flat=[u'id:%s' % unicode(bot_id)],
       execution_timeout_secs=0,
       grace_period_secs=0,
       io_timeout_secs=0)
   now = utils.utcnow()
   request = TaskRequest(
       created_ts=now,
       expiration_ts=now + datetime.timedelta(days=1),
       name=u'Terminate %s' % bot_id,
       priority=0,
       properties=properties,
@@ skipping 146 matching lines @@
     if request.service_account_token == 'bot':
       request.service_account = 'bot'
     else:
       # TODO(vadimsh): Check the token signature, verify it can be used by the
       # current user, extract service account email.
       raise auth.AuthorizationError('service_account_token is not implemented')
 
   request.tags.append('priority:%s' % request.priority)
   request.tags.append('user:%s' % request.user)
   request.tags.append('service_account:%s' % request.service_account)
-  for key, value in request.properties.dimensions.iteritems():
+  for key, value in request.properties.dimensions:
     request.tags.append('%s:%s' % (key, value))
   request.tags = sorted(set(request.tags))
 
   if request.properties.idempotent:
     props = request.properties.to_dict()
     if secret_bytes_ent is not None:
       props['secret_bytes'] = secret_bytes_ent.secret_bytes.encode('hex')
     request.properties_hash = request.HASHING_ALGO(
       utils.encode_to_json(props)).digest()
   else:
@@ skipping 21 matching lines @@
   recheck it again, ensuring the current user (the one that triggers the retry)
   is able to use the token.
 
   Returns:
     The newly created TaskRequest.
   """
   now = utils.utcnow()
   if original_request.properties.is_terminate:
     raise ValueError('cannot clone a terminate request')
   orig_props = original_request.properties.to_dict()
-  dimensions_dict = orig_props.pop(u'dimensions')
-  properties = TaskProperties(dimensions_dict=dimensions_dict, **orig_props)
+  dimensions_flat = [
+    u'%s:%s' % (k, v) for k, v in orig_props.pop(u'dimensions')
+  ]
+  properties = TaskProperties(dimensions_flat=dimensions_flat, **orig_props)
   properties.idempotent = False
   expiration_ts = (
       now + (original_request.expiration_ts - original_request.created_ts))
   name = original_request.name
   match = re.match(r'^(.*) \(Retry #(\d+)\)$', name)
   if match:
     name = '%s (Retry #%d)' % (match.group(1), int(match.group(2)) + 1)
   else:
     name += ' (Retry #1)'
   user = auth.get_current_identity()
@@ skipping 16 matching lines @@
   init_new_request(request, allow_high_priority, original_secret_bytes)
   return request
 
 
 def validate_priority(priority):
   """Throws ValueError if priority is not a valid value."""
   if 0 > priority or MAXIMUM_PRIORITY < priority:
     raise datastore_errors.BadValueError(
         'priority (%d) must be between 0 and %d (inclusive)' %
         (priority, MAXIMUM_PRIORITY))