Removing datastore Batch.add_auto_id_entity.

gcloud/datastore/batch.py

@@ -65,7 +65,7 @@ class Batch(object):
     def __init__(self, client):
         self._client = client
         self._mutation = datastore_pb.Mutation()
-        self._auto_id_entities = []
+        self._partial_key_entities = []
 
     def current(self):
         """Return the topmost batch / transaction, or None."""
@@ -114,30 +114,6 @@ def mutation(self):
         """
         return self._mutation
 
-    def add_auto_id_entity(self, entity):
-        """Adds an entity to the list of entities to update with IDs.
-
-        When an entity has a partial key, calling ``save()`` adds an
-        insert_auto_id entry in the mutation.  In order to make sure we
-        update the Entity once the transaction is committed, we need to
-        keep track of which entities to update (and the order is
-        important).
-
-        When you call ``save()`` on an entity inside a transaction, if
-        the entity has a partial key, it adds itself to the list of
-        entities to be updated once the transaction is committed by
-        calling this method.
-
-        :type entity: :class:`gcloud.datastore.entity.Entity`
-        :param entity: The entity to be updated with a completed key.
-
-        :raises: ValueError if the entity's key is alread completed.
-        """
-        if not entity.key.is_partial:
-            raise ValueError("Entity has a completed key")
-
-        self._auto_id_entities.append(entity)
-
     def put(self, entity):
         """Remember an entity's state to be saved during ``commit``.
 
@@ -152,6 +128,11 @@ def put(self, entity):
            Python3) map to 'string_value' in the datastore;  values which are
            "bytes" ('str' in Python2, 'bytes' in Python3) map to 'blob_value'.
 
+        When an entity has a partial key, calling :meth:`commit`` sends it as
+        an ``insert_auto_id`` mutation and the key is completed. On return, the
+        key for the ``entity`` passed in as updated to match the key ID
+        assigned by the server.
+
         :type entity: :class:`gcloud.datastore.entity.Entity`
         :param entity: the entity to be saved.
 
@@ -164,8 +145,13 @@ def put(self, entity):
         if not _dataset_ids_equal(self.dataset_id, entity.key.dataset_id):
             raise ValueError("Key must be from same dataset as batch")
 
-        _assign_entity_to_mutation(
-            self.mutation, entity, self._auto_id_entities)
+        if entity.key.is_partial:
+            entity_pb = self.mutation.insert_auto_id.add()
+            self._partial_key_entities.append(entity)
+        else:
+            entity_pb = self.mutation.upsert.add()
+
+        _assign_entity_to_pb(entity_pb, entity)
 
     def delete(self, key):
         """Remember a key to be deleted durring ``commit``.
@@ -204,9 +190,9 @@ def commit(self):
         # If the back-end returns without error, we are guaranteed that
         # the response's 'insert_auto_id_key' will match (length and order)
         # the request's 'insert_auto_id` entities, which are derived from
-        # our '_auto_id_entities' (no partial success).
+        # our '_partial_key_entities' (no partial success).
         for new_key_pb, entity in zip(response.insert_auto_id_key,
-                                      self._auto_id_entities):
+                                      self._partial_key_entities):
             new_id = new_key_pb.path_element[-1].id
             entity.key = entity.key.completed_key(new_id)
 
@@ -232,47 +218,28 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             self._client._pop_batch()
 
 
-def _assign_entity_to_mutation(mutation_pb, entity, auto_id_entities):
-    """Copy ``entity`` into appropriate slot of ``mutation_pb``.
-
-    If ``entity.key`` is incomplete, append ``entity`` to ``auto_id_entities``
-    for later fixup during ``commit``.
+def _assign_entity_to_pb(entity_pb, entity):
+    """Copy ``entity`` into ``entity_pb``.
 
     Helper method for ``Batch.put``.
 
-    :type mutation_pb: :class:`gcloud.datastore._datastore_v1_pb2.Mutation`
-    :param mutation_pb: The Mutation protobuf for the batch / transaction.
+    :type entity_pb: :class:`gcloud.datastore._datastore_v1_pb2.Entity`
+    :param entity_pb: The entity owned by a mutation.
 
     :type entity: :class:`gcloud.datastore.entity.Entity`
     :param entity: The entity being updated within the batch / transaction.
-
-    :type auto_id_entities: list of :class:`gcloud.datastore.entity.Entity`
-    :param auto_id_entities: Entities with partial keys, to be fixed up
-                             during commit.
     """
-    auto_id = entity.key.is_partial
-
     key_pb = entity.key.to_protobuf()
     key_pb = helpers._prepare_key_for_request(key_pb)
-
-    if auto_id:
-        insert = mutation_pb.insert_auto_id.add()
-        auto_id_entities.append(entity)
-    else:
-        # We use ``upsert`` for entities with completed keys, rather than
-        # ``insert`` or ``update``, in order not to create race conditions
-        # based on prior existence / removal of the entity.
-        insert = mutation_pb.upsert.add()
-
-    insert.key.CopyFrom(key_pb)
+    entity_pb.key.CopyFrom(key_pb)
 
     for name, value in entity.items():
 
         value_is_list = isinstance(value, list)
         if value_is_list and len(value) == 0:
             continue
 
-        prop = insert.property.add()
+        prop = entity_pb.property.add()
         # Set the name of the property.
         prop.name = name
 

gcloud/datastore/test_batch.py

@@ -38,7 +38,7 @@ def test_ctor(self):
         self.assertEqual(batch.namespace, _NAMESPACE)
         self.assertTrue(batch._id is None)
         self.assertTrue(isinstance(batch.mutation, Mutation))
-        self.assertEqual(batch._auto_id_entities, [])
+        self.assertEqual(batch._partial_key_entities, [])
 
     def test_current(self):
         _DATASET = 'DATASET'
@@ -59,29 +59,6 @@ def test_current(self):
         self.assertTrue(batch1.current() is None)
         self.assertTrue(batch2.current() is None)
 
-    def test_add_auto_id_entity_w_partial_key(self):
-        _DATASET = 'DATASET'
-        connection = _Connection()
-        client = _Client(_DATASET, connection)
-        batch = self._makeOne(client)
-        entity = _Entity()
-        key = entity.key = _Key(_DATASET)
-        key._id = None
-
-        batch.add_auto_id_entity(entity)
-
-        self.assertEqual(batch._auto_id_entities, [entity])
-
-    def test_add_auto_id_entity_w_completed_key(self):
-        _DATASET = 'DATASET'
-        connection = _Connection()
-        client = _Client(_DATASET, connection)
-        batch = self._makeOne(client)
-        entity = _Entity()
-        entity.key = _Key(_DATASET)
-
-        self.assertRaises(ValueError, batch.add_auto_id_entity, entity)
-
     def test_put_entity_wo_key(self):
         _DATASET = 'DATASET'
         connection = _Connection()
@@ -119,7 +96,7 @@ def test_put_entity_w_partial_key(self):
         self.assertEqual(len(upserts), 0)
         deletes = list(batch.mutation.delete)
         self.assertEqual(len(deletes), 0)
-        self.assertEqual(batch._auto_id_entities, [entity])
+        self.assertEqual(batch._partial_key_entities, [entity])
 
     def test_put_entity_w_completed_key(self):
         _DATASET = 'DATASET'
@@ -257,7 +234,7 @@ def test_commit(self):
         self.assertEqual(connection._committed,
                          [(_DATASET, batch.mutation, None)])
 
-    def test_commit_w_auto_id_entities(self):
+    def test_commit_w_partial_key_entities(self):
         _DATASET = 'DATASET'
         _NEW_ID = 1234
         connection = _Connection(_NEW_ID)
@@ -266,7 +243,7 @@ def test_commit_w_auto_id_entities(self):
         entity = _Entity({})
         key = entity.key = _Key(_DATASET)
         key._id = None
-        batch._auto_id_entities.append(entity)
+        batch._partial_key_entities.append(entity)
 
         batch.commit()
 

gcloud/datastore/test_transaction.py

@@ -36,7 +36,7 @@ def test_ctor_defaults(self):
         self.assertEqual(xact.id, None)
         self.assertEqual(xact._status, self._getTargetClass()._INITIAL)
         self.assertTrue(isinstance(xact.mutation, Mutation))
-        self.assertEqual(len(xact._auto_id_entities), 0)
+        self.assertEqual(len(xact._partial_key_entities), 0)
 
     def test_current(self):
         from gcloud.datastore.test_client import _NoCommitBatch
@@ -97,7 +97,7 @@ def test_rollback(self):
         self.assertEqual(xact.id, None)
         self.assertEqual(connection._rolled_back, (_DATASET, 234))
 
-    def test_commit_no_auto_ids(self):
+    def test_commit_no_partial_keys(self):
         _DATASET = 'DATASET'
         connection = _Connection(234)
         client = _Client(_DATASET, connection)
@@ -108,7 +108,7 @@ def test_commit_no_auto_ids(self):
         self.assertEqual(connection._committed, (_DATASET, mutation, 234))
         self.assertEqual(xact.id, None)
 
-    def test_commit_w_auto_ids(self):
+    def test_commit_w_partial_keys(self):
         _DATASET = 'DATASET'
         _KIND = 'KIND'
         _ID = 123
@@ -118,7 +118,7 @@ def test_commit_w_auto_ids(self):
         client = _Client(_DATASET, connection)
         xact = self._makeOne(client)
         entity = _Entity()
-        xact.add_auto_id_entity(entity)
+        xact.put(entity)
         xact._mutation = mutation = object()
         xact.begin()
         xact.commit()
@@ -197,9 +197,10 @@ def __init__(self, *new_keys):
         self.insert_auto_id_key = new_keys
 
 
-class _Entity(object):
+class _Entity(dict):
 
     def __init__(self):
+        super(_Entity, self).__init__()
         from gcloud.datastore.key import Key
         self.key = Key('KIND', dataset_id='DATASET')