thirdparty/google_appengine/google/appengine/api/datastore.py
changeset 2864 2e0b0af889be
parent 2413 d0b7dac5325c
child 3031 7678f72140e6
--- a/thirdparty/google_appengine/google/appengine/api/datastore.py	Sat Sep 05 14:04:24 2009 +0200
+++ b/thirdparty/google_appengine/google/appengine/api/datastore.py	Sun Sep 06 23:31:53 2009 +0200
@@ -49,12 +49,19 @@
 from google.appengine.runtime import apiproxy_errors
 from google.appengine.datastore import entity_pb
 
+try:
+  from google.appengine.api.labs.taskqueue import taskqueue_service_pb
+except ImportError:
+  from google.appengine.api.taskqueue import taskqueue_service_pb
+
 MAX_ALLOWABLE_QUERIES = 30
 
 DEFAULT_TRANSACTION_RETRIES = 3
 
 _MAX_INDEXED_PROPERTIES = 5000
 
+_MAX_ID_BATCH_SIZE = 1000 * 1000 * 1000
+
 Key = datastore_types.Key
 typename = datastore_types.typename
 
@@ -147,7 +154,7 @@
     return []
 
   for entity in entities:
-    if not entity.kind() or not entity.app():
+    if not entity.kind() or not entity.app_id_namespace():
       raise datastore_errors.BadRequestError(
           'App and kind must not be empty, in entity: %s' % entity)
 
@@ -156,8 +163,6 @@
 
   keys = [e.key() for e in entities]
   tx = _MaybeSetupTransaction(req, keys)
-  if tx:
-    tx.RecordModifiedKeys([k for k in keys if k.has_id_or_name()])
 
   resp = datastore_pb.PutResponse()
   try:
@@ -177,7 +182,6 @@
     entity._Entity__key._Key__reference.CopyFrom(key)
 
   if tx:
-    tx.RecordModifiedKeys([e.key() for e in entities], error_on_repeat=False)
     tx.entity_group = entities[0].entity_group()
 
   if multiple:
@@ -259,8 +263,6 @@
   req.key_list().extend([key._Key__reference for key in keys])
 
   tx = _MaybeSetupTransaction(req, keys)
-  if tx:
-    tx.RecordModifiedKeys(keys)
 
   resp = datastore_pb.DeleteResponse()
   try:
@@ -275,8 +277,8 @@
   Includes read-only accessors for app id, kind, and primary key. Also
   provides dictionary-style access to properties.
   """
-  def __init__(self, kind, parent=None, _app=None, name=None,
-               unindexed_properties=[]):
+  def __init__(self, kind, parent=None, _app=None, name=None, id=None,
+               unindexed_properties=[], _namespace=None):
     """Constructor. Takes the kind and transaction root, which cannot be
     changed after the entity is constructed, and an optional parent. Raises
     BadArgumentError or BadKeyError if kind is invalid or parent is not an
@@ -289,33 +291,41 @@
       parent: Entity or Key
       # if provided, this entity's name.
       name: string
+      # if provided, this entity's id.
+      id: integer
       # if provided, a sequence of property names that should not be indexed
       # by the built-in single property indices.
       unindexed_properties: list or tuple of strings
     """
     ref = entity_pb.Reference()
-    _app = datastore_types.ResolveAppId(_app)
-    ref.set_app(_app)
+    _app_namespace = datastore_types.ResolveAppIdNamespace(_app, _namespace)
+    ref.set_app(_app_namespace.to_encoded())
 
     datastore_types.ValidateString(kind, 'kind',
                                    datastore_errors.BadArgumentError)
-
     if parent is not None:
       parent = _GetCompleteKeyOrError(parent)
-      if _app != parent.app():
+      if _app_namespace != parent.app_id_namespace():
         raise datastore_errors.BadArgumentError(
-            "_app %s doesn't match parent's app %s" % (_app, parent.app()))
+            " %s doesn't match parent's app_namespace %s" %
+            (_app_namespace, parent.app_id_namespace()))
       ref.CopyFrom(parent._Key__reference)
 
     last_path = ref.mutable_path().add_element()
     last_path.set_type(kind.encode('utf-8'))
 
+    if name is not None and id is not None:
+      raise datastore_errors.BadArgumentError(
+          "Cannot set both name and id on an Entity")
+
     if name is not None:
       datastore_types.ValidateString(name, 'name')
-      if name[0] in string.digits:
-        raise datastore_errors.BadValueError('name cannot begin with a digit')
       last_path.set_name(name.encode('utf-8'))
 
+    if id is not None:
+      datastore_types.ValidateInteger(id, 'id')
+      last_path.set_id(id)
+
     unindexed_properties, multiple = NormalizeAndTypeCheck(unindexed_properties, basestring)
     if not multiple:
       raise datastore_errors.BadArgumentError(
@@ -329,15 +339,32 @@
 
   def app(self):
     """Returns the name of the application that created this entity, a
-    string.
+    string or None if not set.
     """
     return self.__key.app()
 
+  def namespace(self):
+    """Returns the namespace of this entity, a string or None.
+    """
+    return self.__key.namespace()
+
+  def app_id_namespace(self):
+    """Returns the AppIdNamespace of this entity or None if not set.
+    """
+    return self.__key.app_id_namespace()
+
   def kind(self):
     """Returns this entity's kind, a string.
     """
     return self.__key.kind()
 
+  def is_saved(self):
+    """Returns if this entity has been saved to the datastore
+    """
+    last_path = self.__key._Key__reference.path().element_list()[-1]
+    return ((last_path.has_name() ^ last_path.has_id()) and
+            self.__key.has_id_or_name())
+
   def key(self):
     """Returns this entity's primary key, a Key instance.
     """
@@ -483,7 +510,15 @@
 
     return xml
 
-  def _ToPb(self):
+  def ToPb(self):
+    """Converts this Entity to its protocol buffer representation.
+
+    Returns:
+      entity_pb.Entity
+    """
+    return self._ToPb(False)
+
+  def _ToPb(self, mark_key_as_saved=True):
     """Converts this Entity to its protocol buffer representation. Not
     intended to be used by application developers.
 
@@ -493,6 +528,9 @@
 
     pb = entity_pb.EntityProto()
     pb.mutable_key().CopyFrom(self.key()._ToPb())
+    last_path = pb.key().path().element_list()[-1]
+    if mark_key_as_saved and last_path.has_name() and last_path.has_id():
+      last_path.clear_id()
 
     group = pb.mutable_entity_group()
     if self.__key.has_id_or_name():
@@ -523,7 +561,25 @@
     return pb
 
   @staticmethod
-  def _FromPb(pb):
+  def FromPb(pb):
+    """Static factory method. Returns the Entity representation of the
+    given protocol buffer (datastore_pb.Entity).
+
+    Args:
+      pb: datastore_pb.Entity or str encoding of a datastore_pb.Entity
+
+    Returns:
+      Entity: the Entity representation of pb
+    """
+    if isinstance(pb, str):
+      real_pb = entity_pb.EntityProto()
+      real_pb.ParseFromString(pb)
+      pb = real_pb
+
+    return Entity._FromPb(pb, require_valid_key=False)
+
+  @staticmethod
+  def _FromPb(pb, require_valid_key=True):
     """Static factory method. Returns the Entity representation of the
     given protocol buffer (datastore_pb.Entity). Not intended to be used by
     application developers.
@@ -542,12 +598,13 @@
     assert pb.key().path().element_size() > 0
 
     last_path = pb.key().path().element_list()[-1]
-    assert last_path.has_id() ^ last_path.has_name()
-    if last_path.has_id():
-      assert last_path.id() != 0
-    else:
-      assert last_path.has_name()
-      assert last_path.name()
+    if require_valid_key:
+      assert last_path.has_id() ^ last_path.has_name()
+      if last_path.has_id():
+        assert last_path.id() != 0
+      else:
+        assert last_path.has_name()
+        assert last_path.name()
 
     unindexed_properties = [p.name() for p in pb.raw_property_list()]
 
@@ -701,7 +758,8 @@
   __inequality_prop = None
   __inequality_count = 0
 
-  def __init__(self, kind, filters={}, _app=None, keys_only=False):
+  def __init__(self, kind=None, filters={}, _app=None, keys_only=False,
+               _namespace=None):
     """Constructor.
 
     Raises BadArgumentError if kind is not a string. Raises BadValueError or
@@ -714,15 +772,17 @@
       filters: dict
       keys_only: boolean
     """
-    datastore_types.ValidateString(kind, 'kind',
-                                   datastore_errors.BadArgumentError)
+    if kind is not None:
+      datastore_types.ValidateString(kind, 'kind',
+                                     datastore_errors.BadArgumentError)
 
     self.__kind = kind
     self.__orderings = []
     self.__filter_order = {}
     self.update(filters)
 
-    self.__app = datastore_types.ResolveAppId(_app)
+    self.__app = datastore_types.ResolveAppIdNamespace(_app,
+                                                       _namespace).to_encoded()
     self.__keys_only = keys_only
 
   def Order(self, *orderings):
@@ -794,6 +854,13 @@
             str(direction))
         direction = Query.ASCENDING
 
+      if (self.__kind is None and
+          (property != datastore_types._KEY_SPECIAL_PROPERTY or
+          direction != Query.ASCENDING)):
+        raise datastore_errors.BadArgumentError(
+            'Only %s ascending orders are supported on kindless queries' %
+            datastore_types._KEY_SPECIAL_PROPERTY)
+
       orderings[i] = (property, direction)
 
     if (orderings and self.__inequality_prop and
@@ -884,16 +951,17 @@
     """
     return self._Run()
 
-  def _Run(self, limit=None, offset=None):
+  def _Run(self, limit=None, offset=None,
+           prefetch_count=None, next_count=None):
     """Runs this query, with an optional result limit and an optional offset.
 
-    Identical to Run, with the extra optional limit and offset parameters.
-    limit and offset must both be integers >= 0.
+    Identical to Run, with the extra optional limit, offset, prefetch_count,
+    next_count parameters. These parameters must be integers >= 0.
 
     This is not intended to be used by application developers. Use Get()
     instead!
     """
-    pb = self._ToPb(limit, offset)
+    pb = self._ToPb(limit, offset, prefetch_count)
     result = datastore_pb.QueryResult()
 
     try:
@@ -907,7 +975,7 @@
         raise datastore_errors.NeedIndexError(
           str(exc) + '\nThis query needs this index:\n' + yaml)
 
-    return Iterator(result)
+    return Iterator(result, batch_size=next_count)
 
   def Get(self, limit, offset=0):
     """Fetches and returns a maximum number of results from the query.
@@ -956,7 +1024,8 @@
           'Argument to Get named \'offset\' must be an int greater than or '
           'equal to 0; received %s (a %s)' % (offset, typename(offset)))
 
-    return self._Run(limit, offset)._Get(limit)
+    return self._Run(limit=limit, offset=offset,
+                     prefetch_count=limit)._Get(limit)
 
   def Count(self, limit=None):
     """Returns the number of entities that this query matches. The returned
@@ -1108,6 +1177,12 @@
           'first sort order, if any sort orders are supplied' %
           ', '.join(self.INEQUALITY_OPERATORS))
 
+    if (self.__kind is None and
+        property != datastore_types._KEY_SPECIAL_PROPERTY):
+      raise datastore_errors.BadFilterError(
+          'Only %s filters are allowed on kindless queries.' %
+          datastore_types._KEY_SPECIAL_PROPERTY)
+
     if property in datastore_types._SPECIAL_PROPERTIES:
       if property == datastore_types._KEY_SPECIAL_PROPERTY:
         for value in values:
@@ -1118,7 +1193,7 @@
 
     return match
 
-  def _ToPb(self, limit=None, offset=None):
+  def _ToPb(self, limit=None, offset=None, count=None):
     """Converts this Query to its protocol buffer representation. Not
     intended to be used by application developers. Enforced by hiding the
     datastore_pb classes.
@@ -1129,6 +1204,8 @@
       # number of results that match the query to skip.  limit is applied
       # after the offset is fulfilled
       offset: int
+      # the requested initial batch size
+      count: int
 
     Returns:
       # the PB representation of this Query
@@ -1138,6 +1215,7 @@
       BadRequestError if called inside a transaction and the query does not
       include an ancestor.
     """
+
     if not self.__ancestor and _CurrentTransactionKey():
       raise datastore_errors.BadRequestError(
         'Only ancestor queries are allowed inside transactions.')
@@ -1145,7 +1223,8 @@
     pb = datastore_pb.Query()
     _MaybeSetupTransaction(pb, [self.__ancestor])
 
-    pb.set_kind(self.__kind.encode('utf-8'))
+    if self.__kind is not None:
+      pb.set_kind(self.__kind.encode('utf-8'))
     pb.set_keys_only(bool(self.__keys_only))
     if self.__app:
       pb.set_app(self.__app.encode('utf-8'))
@@ -1153,6 +1232,8 @@
       pb.set_limit(limit)
     if offset is not None:
       pb.set_offset(offset)
+    if count is not None:
+      pb.set_count(count)
     if self.__ancestor:
       pb.mutable_ancestor().CopyFrom(self.__ancestor._Key__reference)
 
@@ -1193,6 +1274,44 @@
     return pb
 
 
+def AllocateIds(model_key, size):
+  """Allocates a range of IDs of size for the key defined by model_key
+
+  Allocates a range of IDs in the datastore such that those IDs will not
+  be automatically assigned to new entities. You can only allocate IDs
+  for model keys from your app. If there is an error, raises a subclass of
+  datastore_errors.Error.
+
+  Args:
+    model_key: Key or string to serve as a model specifying the ID sequence
+               in which to allocate IDs
+
+  Returns:
+    (start, end) of the allocated range, inclusive.
+  """
+  keys, multiple = NormalizeAndTypeCheckKeys(model_key)
+
+  if len(keys) > 1:
+    raise datastore_errors.BadArgumentError(
+        'Cannot allocate IDs for more than one model key at a time')
+
+  if size > _MAX_ID_BATCH_SIZE:
+    raise datastore_errors.BadArgumentError(
+        'Cannot allocate more than %s ids at a time' % _MAX_ID_BATCH_SIZE)
+
+  req = datastore_pb.AllocateIdsRequest()
+  req.mutable_model_key().CopyFrom(keys[0]._Key__reference)
+  req.set_size(size)
+
+  resp = datastore_pb.AllocateIdsResponse()
+  try:
+    apiproxy_stub_map.MakeSyncCall('datastore_v3', 'AllocateIds', req, resp)
+  except apiproxy_errors.ApplicationError, err:
+    raise _ToDatastoreError(err)
+
+  return resp.start(), resp.end()
+
+
 class MultiQuery(Query):
   """Class representing a query which requires multiple datastore queries.
 
@@ -1517,9 +1636,10 @@
   > for person in it:
   >   print 'Hi, %s!' % person['name']
   """
-  def __init__(self, query_result_pb):
+  def __init__(self, query_result_pb, batch_size=None):
     self.__cursor = query_result_pb.cursor()
     self.__keys_only = query_result_pb.keys_only()
+    self.__batch_size = batch_size
     self.__buffer = self._ProcessQueryResult(query_result_pb)
 
   def _Get(self, count):
@@ -1547,16 +1667,16 @@
       # a list of entities or keys
       [Entity or Key, ...]
     """
-    entityList = self._Next(count)
-    while len(entityList) < count and self.__more_results:
-      next_results = self._Next(count - len(entityList))
+    entity_list = self._Next(count)
+    while len(entity_list) < count and self.__more_results:
+      next_results = self._Next(count - len(entity_list), self.__batch_size)
       if not next_results:
         break
-      entityList += next_results
-    return entityList;
-
-  def _Next(self, count):
-    """Returns the next result(s) of the query.
+      entity_list += next_results
+    return entity_list;
+
+  def _Next(self, count=None):
+    """Returns the next batch of results.
 
     Not intended to be used by application developers. Use the python
     iterator protocol instead.
@@ -1565,11 +1685,14 @@
     results. If the query specified a sort order, results are returned in that
     order. Otherwise, the order is undefined.
 
-    The argument, count, specifies the number of results to return. However, the
-    length of the returned list may be smaller than count. This is the case if
-    count is greater than the number of remaining results or the size of the
-    remaining results exciteds the RPC buffer limit. Use _Get to insure all
-    possible entities are retrieved.
+    The optional argument, count, specifies the number of results to return.
+    However, the length of the returned list may be smaller than count. This is
+    the case if count is greater than the number of remaining results or the
+    size of the remaining results exceeds the RPC buffer limit. Use _Get to
+    insure all possible entities are retrieved.
+
+    If the count is omitted, the datastore backend decides how many entities to
+    send.
 
     There is an internal buffer for use with the next() method. If this buffer
     is not empty, up to 'count' values are removed from this buffer and
@@ -1580,19 +1703,23 @@
 
     Args:
       # the number of results to return; must be >= 1
-      count: int or long
+      count: int or long or None
 
     Returns:
       # a list of entities or keys
       [Entity or Key, ...]
     """
-    if not isinstance(count, (int, long)) or count <= 0:
+    if count is not None and (not isinstance(count, (int, long)) or count <= 0):
       raise datastore_errors.BadArgumentError(
         'Argument to _Next must be an int greater than 0; received %s (a %s)' %
         (count, typename(count)))
 
     if self.__buffer:
-      if count <= len(self.__buffer):
+      if count is None:
+        entity_list = self.__buffer
+        self.__buffer = []
+        return entity_list
+      elif count <= len(self.__buffer):
         entity_list = self.__buffer[:count]
         del self.__buffer[:count]
         return entity_list
@@ -1601,13 +1728,15 @@
         self.__buffer = []
         count -= len(entity_list)
     else:
-        entity_list=[]
+        entity_list = []
+
 
     if not self.__more_results:
       return entity_list
 
     req = datastore_pb.NextRequest()
-    req.set_count(count)
+    if count is not None:
+      req.set_count(count)
     req.mutable_cursor().CopyFrom(self.__cursor)
     result = datastore_pb.QueryResult()
     try:
@@ -1642,11 +1771,9 @@
     else:
       return [Entity._FromPb(e) for e in result.result_list()]
 
-  _BUFFER_SIZE = 20
-
   def next(self):
     if not self.__buffer:
-      self.__buffer = self._Next(self._BUFFER_SIZE)
+      self.__buffer = self._Next(self.__batch_size)
     try:
       return self.__buffer.pop(0)
     except IndexError:
@@ -1657,44 +1784,28 @@
 class _Transaction(object):
   """Encapsulates a transaction currently in progress.
 
-  If we've sent a BeginTransaction call, then handle will be a
-  datastore_pb.Transaction that holds the transaction handle.
-
   If we know the entity group for this transaction, it's stored in the
-  entity_group attribute, which is set by RecordModifiedKeys().
+  entity_group attribute, which is set by RunInTransaction().
 
   modified_keys is a set containing the Keys of all entities modified (ie put
   or deleted) in this transaction. If an entity is modified more than once, a
   BadRequestError is raised.
   """
-  def __init__(self):
-    """Initializes modified_keys to the empty set."""
-    self.handle = None
+  def __init__(self, handle):
+    """Initializes the transaction.
+
+    Args:
+      handle: a datastore_pb.Transaction returned by a BeginTransaction call
+    """
+    assert isinstance(handle, datastore_pb.Transaction)
+    explanation = []
+    assert handle.IsInitialized(explanation), explanation
+
+    self.handle = handle
     self.entity_group = None
     self.modified_keys = None
     self.modified_keys = set()
 
-  def RecordModifiedKeys(self, keys, error_on_repeat=True):
-    """Updates the modified keys seen so far.
-
-    If error_on_repeat is True and any of the given keys have already been
-    modified, raises BadRequestError.
-
-    Args:
-      keys: sequence of Keys
-    """
-    keys, _ = NormalizeAndTypeCheckKeys(keys)
-    keys = set(keys)
-
-    if error_on_repeat:
-      already_modified = self.modified_keys.intersection(keys)
-      if already_modified:
-        raise datastore_errors.BadRequestError(
-          "Can't update entity more than once in a transaction: %r" %
-          already_modified.pop())
-
-    self.modified_keys.update(keys)
-
 
 def RunInTransaction(function, *args, **kwargs):
   """Runs a function inside a datastore transaction.
@@ -1799,26 +1910,31 @@
 
   try:
     tx_key = _NewTransactionKey()
-    tx = _Transaction()
-    _txes[tx_key] = tx
 
     for i in range(0, retries + 1):
-      tx.modified_keys.clear()
+      handle = datastore_pb.Transaction()
+      try:
+        apiproxy_stub_map.MakeSyncCall('datastore_v3', 'BeginTransaction',
+                                       api_base_pb.VoidProto(), handle)
+      except apiproxy_errors.ApplicationError, err:
+        raise _ToDatastoreError(err)
+
+      tx = _Transaction(handle)
+      _txes[tx_key] = tx
 
       try:
         result = function(*args, **kwargs)
       except:
         original_exception = sys.exc_info()
 
-        if tx.handle:
-          try:
-            resp = api_base_pb.VoidProto()
-            apiproxy_stub_map.MakeSyncCall('datastore_v3', 'Rollback',
-                                           tx.handle, resp)
-          except:
-            exc_info = sys.exc_info()
-            logging.info('Exception sending Rollback:\n' +
-                         ''.join(traceback.format_exception(*exc_info)))
+        try:
+          resp = api_base_pb.VoidProto()
+          apiproxy_stub_map.MakeSyncCall('datastore_v3', 'Rollback',
+                                         tx.handle, resp)
+        except:
+          exc_info = sys.exc_info()
+          logging.info('Exception sending Rollback:\n' +
+                       ''.join(traceback.format_exception(*exc_info)))
 
         type, value, trace = original_exception
         if type is datastore_errors.Rollback:
@@ -1826,21 +1942,20 @@
         else:
           raise type, value, trace
 
-      if tx.handle:
-        try:
-          resp = datastore_pb.CommitResponse()
-          apiproxy_stub_map.MakeSyncCall('datastore_v3', 'Commit',
-                                         tx.handle, resp)
-        except apiproxy_errors.ApplicationError, err:
-          if (err.application_error ==
-              datastore_pb.Error.CONCURRENT_TRANSACTION):
-            logging.warning('Transaction collision for entity group with '
-                            'key %r. Retrying...', tx.entity_group)
-            tx.handle = None
-            tx.entity_group = None
-            continue
-          else:
-            raise _ToDatastoreError(err)
+      try:
+        resp = datastore_pb.CommitResponse()
+        apiproxy_stub_map.MakeSyncCall('datastore_v3', 'Commit',
+                                       tx.handle, resp)
+      except apiproxy_errors.ApplicationError, err:
+        if (err.application_error ==
+            datastore_pb.Error.CONCURRENT_TRANSACTION):
+          logging.warning('Transaction collision for entity group with '
+                          'key %r. Retrying...', tx.entity_group)
+          tx.handle = None
+          tx.entity_group = None
+          continue
+        else:
+          raise _ToDatastoreError(err)
 
       return result
 
@@ -1854,12 +1969,11 @@
 
 
 def _MaybeSetupTransaction(request, keys):
-  """Begins a transaction, if necessary, and populates it in the request.
+  """If we're in a transaction, validates and populates it in the request.
 
   If we're currently inside a transaction, this records the entity group,
-  checks that the keys are all in that entity group, creates the transaction
-  PB, and sends the BeginTransaction. It then populates the transaction handle
-  in the request.
+  checks that the keys are all in that entity group, and populates the
+  transaction handle in the request.
 
   Raises BadRequestError if the entity has a different entity group than the
   current transaction.
@@ -1872,7 +1986,9 @@
     _Transaction if we're inside a transaction, otherwise None
   """
   assert isinstance(request, (datastore_pb.GetRequest, datastore_pb.PutRequest,
-                              datastore_pb.DeleteRequest, datastore_pb.Query))
+                              datastore_pb.DeleteRequest, datastore_pb.Query,
+                              taskqueue_service_pb.TaskQueueAddRequest,
+                              )), request.__class__
   tx_key = None
 
   try:
@@ -1883,8 +1999,10 @@
       groups = [k.entity_group() for k in keys]
       if tx.entity_group:
         expected_group = tx.entity_group
+      elif groups:
+        expected_group = groups[0]
       else:
-        expected_group = groups[0]
+        expected_group = None
 
       for group in groups:
         if (group != expected_group or
@@ -1901,12 +2019,7 @@
         if not tx.entity_group and group.has_id_or_name():
           tx.entity_group = group
 
-      if not tx.handle:
-        tx.handle = datastore_pb.Transaction()
-        req = api_base_pb.VoidProto()
-        apiproxy_stub_map.MakeSyncCall('datastore_v3', 'BeginTransaction', req,
-                                       tx.handle)
-
+      assert tx.handle.IsInitialized()
       request.mutable_transaction().CopyFrom(tx.handle)
 
       return tx