Add 'Batch.add_auto_id_entity' method.

FBO the proposed global 'put'.
tseaver · Jan 8, 2015 · Jan 8, 2015 · Jan 8, 2015 · Jan 8, 2015 · Jan 8, 2015
commit d5d28c94277ed97cf6c55fcd21dd9e3886d62759
diff --git a/gcloud/datastore/batch.py b/gcloud/datastore/batch.py
@@ -108,6 +108,25 @@ 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.
+        """
+        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``.
 

diff --git a/gcloud/datastore/test_batch.py b/gcloud/datastore/test_batch.py
@@ -65,6 +65,27 @@ def test_ctor_implicit(self):
         self.assertTrue(isinstance(batch.mutation, Mutation))
         self.assertEqual(batch._auto_id_entities, [])
 
+    def test_add_auto_id_entity_w_partial_key(self):
+        _DATASET = 'DATASET'
+        connection = _Connection()
+        batch = self._makeOne(dataset_id=_DATASET, connection=connection)
+        entity = _Entity()
+        key = entity.key = _Key(_Entity)
+        key._partial = True
+
+        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()
+        batch = self._makeOne(dataset_id=_DATASET, connection=connection)
+        entity = _Entity()
+        key = entity.key = _Key(_Entity)
+
+        self.assertRaises(ValueError, batch.add_auto_id_entity, entity)
+
     def test_put_entity_wo_key(self):
         _DATASET = 'DATASET'
         connection = _Connection()