Merge pull request #1226 from mpeg/datastore-fix-unprefixed-id

Fix datastore: unprefixed ids in client

Author: dhermes

gcloud/datastore/client.py

@@ -24,6 +24,7 @@
 from gcloud.datastore.batch import Batch
 from gcloud.datastore.entity import Entity
 from gcloud.datastore.key import Key
+from gcloud.datastore.key import _dataset_ids_equal
 from gcloud.datastore.query import Query
 from gcloud.datastore.transaction import Transaction
 from gcloud.environment_vars import DATASET
@@ -288,9 +289,10 @@ def get_multi(self, keys, missing=None, deferred=None):
         if not keys:
             return []
 
-        ids = list(set([key.dataset_id for key in keys]))
-        if ids != [self.dataset_id]:
-            raise ValueError('Keys do not match dataset ID')
+        ids = set(key.dataset_id for key in keys)
+        for current_id in ids:
+            if not _dataset_ids_equal(current_id, self.dataset_id):
+                raise ValueError('Keys do not match dataset ID')
 
         transaction = self.current_transaction
 

gcloud/datastore/test_client.py

@@ -500,6 +500,55 @@ def test_get_multi_hit_multiple_keys_different_dataset(self):
         with self.assertRaises(ValueError):
             client.get_multi([key1, key2])
 
+    def test_get_multi_diff_prefixes(self):
+        from gcloud.datastore.key import Key
+
+        DATASET_ID1 = 'DATASET'
+        DATASET_ID2 = 'e~DATASET'
+        DATASET_ID3 = 's~DATASET'
+        KIND = 'Kind'
+        ID1 = 1234
+        ID2 = 2345
+        ID3 = 3456
+
+        # Make found entity pbs to be returned from mock backend.
+        entity_pb1 = _make_entity_pb(DATASET_ID1, KIND, ID1)
+        entity_pb2 = _make_entity_pb(DATASET_ID2, KIND, ID2)
+        entity_pb3 = _make_entity_pb(DATASET_ID3, KIND, ID3)
+
+        creds = object()
+        client = self._makeOne(credentials=creds)
+        client.connection._add_lookup_result([entity_pb1,
+                                              entity_pb2,
+                                              entity_pb3])
+
+        key1 = Key(KIND, ID1, dataset_id=DATASET_ID1)
+        key2 = Key(KIND, ID2, dataset_id=DATASET_ID2)
+        key3 = Key(KIND, ID3, dataset_id=DATASET_ID3)
+
+        retrieved_all = client.get_multi([key1, key2, key3])
+        retrieved1, retrieved2, retrieved3 = retrieved_all
+
+        # Check values & keys match.
+        self.assertEqual(retrieved1.key.path, key1.path)
+        self.assertEqual(retrieved2.key.path, key2.path)
+        self.assertEqual(retrieved3.key.path, key3.path)
+
+    def test_get_multi_diff_datasets_w_prefix(self):
+        from gcloud.datastore.key import Key
+
+        DATASET_ID1 = 'e~DATASET'
+        DATASET_ID2 = 's~DATASET-ALT'
+
+        key1 = Key('KIND', 1234, dataset_id=DATASET_ID1)
+        key2 = Key('KIND', 1234, dataset_id=DATASET_ID2)
+
+        creds = object()
+        client = self._makeOne(credentials=creds)
+
+        with self.assertRaises(ValueError):
+            client.get_multi([key1, key2])
+
     def test_get_multi_max_loops(self):
         from gcloud._testing import _Monkey
         from gcloud.datastore import client as _MUT