[Findit] Extending versioned_model.py to support versioning multiple entities of the same class.

appengine/findit/model/versioned_model.py

 # Copyright 2015 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """Provides a model to support versioned entities in datastore.
 
 Idea: use a root model entity to keep track of the most recent version of a
 versioned entity, and make the versioned entities and the root model entity in
 the same entity group so that they could be read and written in a transaction.
 """
 
 import logging
 
 from google.appengine.api import datastore_errors
 from google.appengine.ext import ndb
 from google.appengine.runtime import apiproxy_errors
 
 
 class _GroupRoot(ndb.Model):
   """Root entity of a group to support versioned children."""
   # Key id of the most recent child entity in the datastore. It is monotonically
   # increasing and is 0 if no child is present.
   current = ndb.IntegerProperty(indexed=False, default=0)
 
 
 class VersionedModel(ndb.Model):
   """A model that supports versioning.
 
   Subclass will automatically be versioned, if use GetVersion() to

[stgao, 2016/09/24 01:43:12]

Update this to reflect the addition of Create below.
With this change, all subclass should use the Create function to create all
instances instead of using ModelName().

Please also update existing testcases too.

[lijeffrey, 2016/09/24 06:13:14]

Done.

   read and use Save() to write.
   """
 
   @property
+  def _root_id(self):
+    return self.key.pairs()[0][1] if self.key else None
+
+  @property
   def version(self):
-    return self.key.integer_id() if self.key else 0
+    # Ndb treats key.integer_id() of 0 as None, so default to 0.
+    return self.key.integer_id() or 0 if self.key else 0
 
   @classmethod
-  def GetVersion(cls, version=None):
+  def Create(cls, root_id=None):
+    """Creates an instance of cls that is to become the first version.
+
+    The calling function of Create() should be responsible first for checking
+    no previous version of the proposed entity already exists.
+
+    Args:
+      root_id: Any user-specified value that will serve as the id for the root
+        entity's key.
+
+    Returns:
+      An instance of cls meant to be the first version. Note for this instance
+        to be committed to the datastore Save() would need to be called on the
+        instance returned by this method.
+    """
+    return cls(key=ndb.Key(cls, 0, parent=cls._GetRootKey(root_id)))
+
+  @classmethod
+  def GetVersion(cls, root_id=None, version=None):
     """Returns a version of the entity, the latest if version=None."""
     assert not ndb.in_transaction()
 
-    root_key = cls._GetRootKey()
+    root_key = cls._GetRootKey(root_id)
     root = root_key.get()
+
     if not root or not root.current:
       return None
 
     if version is None:
       version = root.current
     elif version < 1:
       #  Return None for versions < 1, which causes exceptions in ndb.Key()
       return None
 
     return ndb.Key(cls, version, parent=root_key).get()
 
   @classmethod
-  def GetLatestVersionNumber(cls):
-    root_entity = cls._GetRootKey().get()
+  def GetLatestVersionNumber(cls, root_id=None):
+    root_entity = cls._GetRootKey(root_id).get()
     if not root_entity:
       return -1
     return root_entity.current
 
-  def Save(self):
-    """Saves the current entity, but as a new version."""
-    root_key = self._GetRootKey()
+  def Save(self, retry_on_conflict=True):
+    """Saves the current entity, but as a new version.
+
+    Args:
+      retry_on_conflict (bool): Whether or not the next version number should
+        automatically be tried in case another transaction writes the entity
+        first with the same proposed new version number.
+
+    Returns:
+      The key of the newly written version, and a boolean whether or not this
+      call to Save() was responsible for creating it.
+    """
+    root_key = self._GetRootKey(self._root_id)
     root = root_key.get() or self._GetRootModel()(key=root_key)
 
     def SaveData():
       if self.key.get():
         return False  # The entity exists, should retry.
+
       ndb.put_multi([self, root])
       return True
 
     def SetNewKey():
       root.current += 1
       self.key = ndb.Key(self.__class__, root.current, parent=root_key)
 
     SetNewKey()
     while True:
       while self.key.get():
-        SetNewKey()
+        if retry_on_conflict:
+          SetNewKey()
+        else:
+          # Another transaction had already written the proposed new version, so
+          # return that version's key and False indicating this call to Save()
+          # was not responsible for creating it.
+          return self.key, False
 
       try:
         if ndb.transaction(SaveData, retries=0):
-          return self.key
+          return self.key, True
       except (
           datastore_errors.InternalError,
           datastore_errors.Timeout,
           datastore_errors.TransactionFailedError) as e:
         # https://cloud.google.com/appengine/docs/python/datastore/transactions
         # states the result is ambiguous, it could have succeeded.
         logging.info('Transaction likely failed: %s', e)
       except (
           apiproxy_errors.CancelledError,
           datastore_errors.BadRequestError,
           RuntimeError) as e:
         logging.info('Transaction failure: %s', e)
       else:
-        SetNewKey()
+        if retry_on_conflict:
+          SetNewKey()
+        else:
+          # Another transaction had already written the proposed new version, so
+          # return that version's key and False indicating this call to Save()
+          # was not responsible for creating it.
+          return self.key, False
 
   @classmethod
   def _GetRootModel(cls):
     """Returns a root model that can be used for versioned entities."""
     root_model_name = '%sRoot' % cls.__name__
 
     class _RootModel(_GroupRoot):
 
       @classmethod
       def _get_kind(cls):
         return root_model_name
 
     return _RootModel
 
   @classmethod
-  def _GetRootKey(cls):
-    return ndb.Key(cls._GetRootModel(), 1)
+  def _GetRootKey(cls, root_id=None):
+    return ndb.Key(
+        cls._GetRootModel(), root_id if root_id is not None else 1)