Require that keys for 'put' / 'delete' match the 'dataset_id' of the batch.

Fixes #447.

Author: tseaver

gcloud/datastore/batch.py

@@ -210,11 +210,15 @@ def put(self, entity):
         :type entity: :class:`gcloud.datastore.entity.Entity`
         :param entity: the entity to be saved.
 
-        :raises: ValueError if entity has no key assigned.
+        :raises: ValueError if entity has no key assigned, or if the key's
+                 ``dataset_id`` does not match ours.
         """
         if entity.key is None:
             raise ValueError("Entity must have a key")
 
+        if entity.key.dataset_id != self._dataset_id:
+            raise ValueError("Key must be from same dataset as batch")
+
         _assign_entity_to_mutation(
             self.mutation, entity, self._auto_id_entities)
 
@@ -224,11 +228,15 @@ def delete(self, key):
         :type key: :class:`gcloud.datastore.key.Key`
         :param key: the key to be deleted.
 
-        :raises: ValueError if key is not complete.
+        :raises: ValueError if key is not complete, or if the key's
+                 ``dataset_id`` does not match ours.
         """
         if key.is_partial:
             raise ValueError("Key must be complete")
 
+        if key.dataset_id != self._dataset_id:
+            raise ValueError("Key must be from same dataset as batch")
+
         key_pb = key.to_protobuf()
         helpers._add_keys_to_request(self.mutation.delete, [key_pb])
 

gcloud/datastore/test_batch.py

@@ -116,7 +116,7 @@ def test_add_auto_id_entity_w_partial_key(self):
         connection = _Connection()
         batch = self._makeOne(dataset_id=_DATASET, connection=connection)
         entity = _Entity()
-        key = entity.key = _Key(_Entity)
+        key = entity.key = _Key(_DATASET)
         key._id = None
 
         batch.add_auto_id_entity(entity)
@@ -128,7 +128,7 @@ def test_add_auto_id_entity_w_completed_key(self):
         connection = _Connection()
         batch = self._makeOne(dataset_id=_DATASET, connection=connection)
         entity = _Entity()
-        entity.key = _Key(_Entity)
+        entity.key = _Key(_DATASET)
 
         self.assertRaises(ValueError, batch.add_auto_id_entity, entity)
 
@@ -139,6 +139,15 @@ def test_put_entity_wo_key(self):
 
         self.assertRaises(ValueError, batch.put, _Entity())
 
+    def test_put_entity_w_key_wrong_dataset_id(self):
+        _DATASET = 'DATASET'
+        connection = _Connection()
+        batch = self._makeOne(dataset_id=_DATASET, connection=connection)
+        entity = _Entity()
+        entity.key = _Key('OTHER')
+
+        self.assertRaises(ValueError, batch.put, entity)
+
     def test_put_entity_w_partial_key(self):
         _DATASET = 'DATASET'
         _PROPERTIES = {'foo': 'bar'}
@@ -203,6 +212,14 @@ def test_delete_w_partial_key(self):
 
         self.assertRaises(ValueError, batch.delete, key)
 
+    def test_delete_w_key_wrong_dataset_id(self):
+        _DATASET = 'DATASET'
+        connection = _Connection()
+        batch = self._makeOne(dataset_id=_DATASET, connection=connection)
+        key = _Key('OTHER')
+
+        self.assertRaises(ValueError, batch.delete, key)
+
     def test_delete_w_completed_key(self):
         _DATASET = 'DATASET'
         connection = _Connection()