[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.
 """
@@ skipping 13 matching lines @@
 
 
 class VersionedModel(ndb.Model):
   """A model that supports versioning.
 
   Subclass will automatically be versioned, if use GetVersion() to
   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
+    return self.key.integer_id() or 0 if self.key else 0
 
   @classmethod
-  def GetVersion(cls, version=None):
+  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)