Work around the fact that keys updated from the server have prefixes.

tseaver · tseaver · commit 1b87a7e139cf · 2015-01-15T14:51:49.000-05:00
See #528.
diff --git a/gcloud/datastore/batch.py b/gcloud/datastore/batch.py
@@ -193,6 +193,22 @@ def add_auto_id_entity(self, entity):
 
         self._auto_id_entities.append(entity)
 
+    def _match_dataset_id(self, other_id):
+        """Ensure that `other_id` matches our `dataset_id`.
+
+        Helper for :meth:`put` and :meth:`delete`.
+
+        :type other_id: string
+        :param other_id: the dataset ID to compare
+
+        :raises: ValueError if `other_id` does not match (even after stripping
+                 any prefix).
+        """
+        other_id = other_id.rsplit('~', 1)[-1]
+        our_id = self._dataset_id.rsplit('~', 1)[-1]
+        if other_id != our_id:
+            raise ValueError("Key must be from same dataset as batch")
+
     def put(self, entity):
         """Remember an entity's state to be saved during ``commit``.
 
@@ -216,8 +232,7 @@ def put(self, entity):
         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")
+        self._match_dataset_id(entity.key.dataset_id)
 
         _assign_entity_to_mutation(
             self.mutation, entity, self._auto_id_entities)
@@ -234,8 +249,7 @@ def delete(self, key):
         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")
+        self._match_dataset_id(key.dataset_id)
 
         key_pb = key.to_protobuf()
         helpers._add_keys_to_request(self.mutation.delete, [key_pb])
diff --git a/gcloud/datastore/test_batch.py b/gcloud/datastore/test_batch.py
@@ -203,6 +203,41 @@ def test_put_entity_w_completed_key(self):
         deletes = list(batch.mutation.delete)
         self.assertEqual(len(deletes), 0)
 
+    def test_put_entity_w_completed_key_prefixed_dataset_id(self):
+        _DATASET = 'DATASET'
+        _PROPERTIES = {
+            'foo': 'bar',
+            'baz': 'qux',
+            'spam': [1, 2, 3],
+            'frotz': [],  # will be ignored
+            }
+        connection = _Connection()
+        batch = self._makeOne(dataset_id=_DATASET, connection=connection)
+        entity = _Entity(_PROPERTIES)
+        entity.exclude_from_indexes = ('baz', 'spam')
+        key = entity.key = _Key('s~' + _DATASET)
+
+        batch.put(entity)
+
+        insert_auto_ids = list(batch.mutation.insert_auto_id)
+        self.assertEqual(len(insert_auto_ids), 0)
+        upserts = list(batch.mutation.upsert)
+        self.assertEqual(len(upserts), 1)
+
+        upsert = upserts[0]
+        self.assertEqual(upsert.key, key._key)
+        props = dict([(prop.name, prop.value) for prop in upsert.property])
+        self.assertTrue(props['foo'].indexed)
+        self.assertFalse(props['baz'].indexed)
+        self.assertTrue(props['spam'].indexed)
+        self.assertFalse(props['spam'].list_value[0].indexed)
+        self.assertFalse(props['spam'].list_value[1].indexed)
+        self.assertFalse(props['spam'].list_value[2].indexed)
+        self.assertFalse('frotz' in props)
+
+        deletes = list(batch.mutation.delete)
+        self.assertEqual(len(deletes), 0)
+
     def test_delete_w_partial_key(self):
         _DATASET = 'DATASET'
         connection = _Connection()
@@ -236,6 +271,22 @@ def test_delete_w_completed_key(self):
         self.assertEqual(len(deletes), 1)
         self.assertEqual(deletes[0], key._key)
 
+    def test_delete_w_completed_key_w_prefixed_dataset_id(self):
+        _DATASET = 'DATASET'
+        connection = _Connection()
+        batch = self._makeOne(dataset_id=_DATASET, connection=connection)
+        key = _Key('s~' + _DATASET)
+
+        batch.delete(key)
+
+        insert_auto_ids = list(batch.mutation.insert_auto_id)
+        self.assertEqual(len(insert_auto_ids), 0)
+        upserts = list(batch.mutation.upsert)
+        self.assertEqual(len(upserts), 0)
+        deletes = list(batch.mutation.delete)
+        self.assertEqual(len(deletes), 1)
+        self.assertEqual(deletes[0], key._key)
+
     def test_commit(self):
         _DATASET = 'DATASET'
         connection = _Connection()
