Make backward-forward compatible version.

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 33 matching lines @@
                    |
                    v
     <See task_to_run.py and task_result.py>
 
 TaskProperties is embedded in TaskRequest. TaskProperties is still declared as a
 separate entity to clearly declare the boundary for task request deduplication.
 """
 
 import datetime
 import hashlib
+import itertools
 import posixpath
 import random
 import re
 import urlparse
 
 from google.appengine.api import datastore_errors
 from google.appengine.ext import ndb
 
 from components import auth
 from components import datastore_utils
@@ skipping 79 matching lines @@
   """Validates TaskProperties.env."""
   # pylint: disable=W0212
   if not all(
          isinstance(k, unicode) and isinstance(v, unicode)
          for k, v in value.iteritems()):
     # pylint: disable=W0212
     raise TypeError(
         '%s must be a dict of strings, not %r' % (prop._name, value))
 
 
-def _validate_dimensions(prop, value):
-  """Validates TaskProperties.dimensions."""
-  # pylint: disable=W0212
+def _validate_dimension_flat(_prop, value):
+  """Validates TaskProperties.dimensions_flat."""
   if not value:
-    raise datastore_errors.BadValueError(u'%s must be specified' % prop._name)
+    raise datastore_errors.BadValueError(u'dimensions must be specified')
+  key, val = value.split(u':', 1)
+  if not config.validate_dimension_key(key):
+    raise datastore_errors.BadValueError(u'dimension %r isn\'t valid' % key)
+  if not config.validate_dimension_value(val):
+    raise datastore_errors.BadValueError(
+        u'dimension %r:%r isn\'t valid' % (key, val))
+
+
+def _validate_dimensions_dict(prop, value):
+  """Validates old TaskProperties.dimensions_dict."""
   _validate_dict_of_strings(prop, value)
   for key, val in value.iteritems():
     if not config.validate_dimension_key(key):
       raise datastore_errors.BadValueError(u'dimension %r isn\'t valid' % key)
     if not config.validate_dimension_value(val):
       raise datastore_errors.BadValueError(
           u'dimension %r:%r isn\'t valid' % (key, val))
-  if u'pool' not in value and u'id' not in value:
-    raise datastore_errors.BadValueError(
-        u'At least one of \'id\' or \'pool\' must be used as %s' % prop._name)
-  if len(value) > 64:
-    raise datastore_errors.BadValueError(
-        '%s can have up to 64 keys' % prop._name)
 
 
 def _validate_env(prop, value):
   _validate_dict_of_strings(prop, value)
   maxlen = 1024
   for k, v in value.iteritems():
     if not k:
       raise datastore_errors.BadValueError(
           'valid key are required in %s' % prop._name)
     if len(k) > maxlen:
@@ skipping 300 matching lines @@
   # output too. See TODO at the top of this class.
   # May be non-None even if task input is not isolated.
   #
   # Only inputs_ref.isolated or command can be specified.
   inputs_ref = ndb.LocalStructuredProperty(FilesRef)
 
   # CIPD packages to install.
   cipd_input = ndb.LocalStructuredProperty(CipdInput)
 
   # Filter to use to determine the required properties on the bot to run on. For
-  # example, Windows or hostname. Encoded as json. Either 'pool' or 'id'
-  # dimension are required (see _validate_dimensions).
-  dimensions = datastore_utils.DeterministicJsonProperty(
-      validator=_validate_dimensions, json_type=dict, indexed=False)
+  # example, Windows or hostname. Either 'pool' or 'id' dimension is required.
+  dimensions_flat = ndb.StringProperty(
+      validator=_validate_dimension_flat, indexed=False, repeated=True)
+  # Old way to express dimensions. Not used anymore as this doesn't support
+  # repeated keys.
+  dimensions_dict = datastore_utils.DeterministicJsonProperty(
+      validator=_validate_dimensions_dict, json_type=dict, indexed=False,
+      name='dimensions')
 
   # Environment variables. Encoded as json. Optional.
   env = datastore_utils.DeterministicJsonProperty(
       validator=_validate_env, json_type=dict, indexed=False)
 
   # Maximum duration the bot can take to run this task. It's named hard_timeout
   # in the bot.
   execution_timeout_secs = ndb.IntegerProperty(
       validator=_validate_timeout, required=True, indexed=False)
 
@@ skipping 21 matching lines @@
   # $(ISOLATED_OUTDIR) will be returned; otherwise, the files in this list
   # 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)."""
+    if self.dimensions_dict:
+      return self.dimensions_dict
+    return dict(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
         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:
+      raise datastore_errors.BadValueError(u'dimensions are required')
+    keys = self.dimensions.keys()
+    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:
+      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.
     isolated_input = self.inputs_ref and self.inputs_ref.isolated
     if not self.command and not isolated_input:
       raise datastore_errors.BadValueError(
           'use at least one of command or inputs_ref.isolated')
     if self.command and self.extra_args:
       raise datastore_errors.BadValueError(
           'can\'t use both command and extra_args')
     if self.extra_args and not isolated_input:
       raise datastore_errors.BadValueError(
           'extra_args require inputs_ref.isolated')
@@ skipping 189 matching lines @@
       raise datastore_errors.BadValueError(
           'bad pubsub topic name - %s' % self.pubsub_topic)
     if self.pubsub_auth_token and not self.pubsub_topic:
       raise datastore_errors.BadValueError(
           'pubsub_auth_token requires pubsub_topic')
     if self.pubsub_userdata and not self.pubsub_topic:
       raise datastore_errors.BadValueError(
           'pubsub_userdata requires pubsub_topic')
 
 
+### Private stuff.
+
+
+def _is_sorted(l):
+  return all(a <= b for a, b in itertools.izip(l[:-1], l[1:]))
+
+
 ### 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={u'id': unicode(bot_id)},
+      dimensions_dict={u'id': 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 187 matching lines @@
   New request uses same 'service_account_token', but 'init_new_request' will
   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')
-  properties = TaskProperties(**original_request.properties.to_dict())
+  orig_props = original_request.properties.to_dict()
+  dimensions_dict = orig_props.pop(u'dimensions')
+  properties = TaskProperties(dimensions_dict=dimensions_dict, **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))