Removing dataset_id from get_entities.

dhermes · dhermes · commit daaa26e76c99 · 2015-01-08T12:05:45.000-08:00
It can be implied from the keys passed in.
diff --git a/gcloud/datastore/__init__.py b/gcloud/datastore/__init__.py
@@ -153,8 +153,7 @@ def _require_connection(connection=None):
     return connection
 
 
-def get_entities(keys, missing=None, deferred=None,
-                 connection=None, dataset_id=None):
+def get_entities(keys, missing=None, deferred=None, connection=None):
     """Retrieves entities, along with their attributes.
 
     :type keys: list of :class:`gcloud.datastore.key.Key`
@@ -173,14 +172,23 @@ def get_entities(keys, missing=None, deferred=None,
     :type connection: :class:`gcloud.datastore.connection.Connection`
     :param connection: Optional. The connection used to connect to datastore.
 
-    :type dataset_id: string
-    :param dataset_id: Optional. The ID of the dataset.
-
     :rtype: list of :class:`gcloud.datastore.entity.Entity`
     :returns: The requested entities.
+    :raises: :class:`ValueError` if the key dataset IDs don't agree.
     """
+    if not keys:
+        return []
+
     connection = _require_connection(connection)
-    dataset_id = _require_dataset_id(dataset_id)
+    dataset_id = keys[0].dataset_id
+    # Rather than creating a list or set of all dataset IDs, we iterate
+    # and check. We could allow the backend to check this for us if IDs
+    # with no prefix worked (GoogleCloudPlatform/google-cloud-datastore#59)
+    # or if we made sure that a prefix s~ or e~ was on each key.
+    for key in keys[1:]:
+        if key.dataset_id != dataset_id:
+            raise ValueError('All keys in get_entities must be from the '
+                             'same dataset.')
 
     entity_pbs = connection.lookup(
         dataset_id=dataset_id,
diff --git a/gcloud/datastore/key.py b/gcloud/datastore/key.py
@@ -232,8 +232,7 @@ def get(self, connection=None):
 
         # We allow partial keys to attempt a get, the backend will fail.
         connection = connection or _implicit_environ.CONNECTION
-        entities = datastore.get_entities(
-            [self], connection=connection, dataset_id=self.dataset_id)
+        entities = datastore.get_entities([self], connection=connection)
 
         if entities:
             result = entities[0]
diff --git a/gcloud/datastore/test___init__.py b/gcloud/datastore/test___init__.py
@@ -213,11 +213,14 @@ def test__require_connection_implicit_set_passed_explicitly(self):
 
 class Test_get_entities_function(unittest2.TestCase):
 
-    def _callFUT(self, keys, missing=None, deferred=None,
-                 connection=None, dataset_id=None):
+    def _callFUT(self, keys, missing=None, deferred=None, connection=None):
         from gcloud.datastore import get_entities
         return get_entities(keys, missing=missing, deferred=deferred,
-                            connection=connection, dataset_id=dataset_id)
+                            connection=connection)
+
+    def test_get_entities_no_keys(self):
+        results = self._callFUT([])
+        self.assertEqual(results, [])
 
     def test_get_entities_miss(self):
         from gcloud.datastore.key import Key
@@ -226,8 +229,7 @@ def test_get_entities_miss(self):
         DATASET_ID = 'DATASET'
         connection = _Connection()
         key = Key('Kind', 1234, dataset_id=DATASET_ID)
-        results = self._callFUT([key], connection=connection,
-                                dataset_id=DATASET_ID)
+        results = self._callFUT([key], connection=connection)
         self.assertEqual(results, [])
 
     def test_get_entities_miss_w_missing(self):
@@ -252,8 +254,7 @@ def test_get_entities_miss_w_missing(self):
 
         key = Key(KIND, ID, dataset_id=DATASET_ID)
         missing = []
-        entities = self._callFUT([key], connection=connection,
-                                 dataset_id=DATASET_ID, missing=missing)
+        entities = self._callFUT([key], connection=connection, missing=missing)
         self.assertEqual(entities, [])
         self.assertEqual([missed.key.to_protobuf() for missed in missing],
                          [key.to_protobuf()])
@@ -271,22 +272,24 @@ def test_get_entities_miss_w_deferred(self):
 
         deferred = []
         entities = self._callFUT([key], connection=connection,
-                                 dataset_id=DATASET_ID, deferred=deferred)
+                                 deferred=deferred)
         self.assertEqual(entities, [])
         self.assertEqual([def_key.to_protobuf() for def_key in deferred],
                          [key.to_protobuf()])
 
-    def _make_entity_pb(self, dataset_id, kind, integer_id, name, str_val):
+    def _make_entity_pb(self, dataset_id, kind, integer_id,
+                        name=None, str_val=None):
         from gcloud.datastore import datastore_v1_pb2 as datastore_pb
 
         entity_pb = datastore_pb.Entity()
         entity_pb.key.partition_id.dataset_id = dataset_id
         path_element = entity_pb.key.path_element.add()
         path_element.kind = kind
         path_element.id = integer_id
-        prop = entity_pb.property.add()
-        prop.name = name
-        prop.value.string_value = str_val
+        if name is not None and str_val is not None:
+            prop = entity_pb.property.add()
+            prop.name = name
+            prop.value.string_value = str_val
 
         return entity_pb
 
@@ -307,8 +310,7 @@ def test_get_entities_hit(self):
         connection = _Connection(entity_pb)
 
         key = Key(KIND, ID, dataset_id=DATASET_ID)
-        result, = self._callFUT([key], connection=connection,
-                                dataset_id=DATASET_ID)
+        result, = self._callFUT([key], connection=connection)
         new_key = result.key
 
         # Check the returned value is as expected.
@@ -318,6 +320,47 @@ def test_get_entities_hit(self):
         self.assertEqual(list(result), ['foo'])
         self.assertEqual(result['foo'], 'Foo')
 
+    def test_get_entities_hit_multiple_keys_same_dataset(self):
+        from gcloud.datastore.key import Key
+        from gcloud.datastore.test_connection import _Connection
+
+        DATASET_ID = 'DATASET'
+        KIND = 'Kind'
+        ID1 = 1234
+        ID2 = 2345
+
+        # Make a found entity pb to be returned from mock backend.
+        entity_pb1 = self._make_entity_pb(DATASET_ID, KIND, ID1)
+        entity_pb2 = self._make_entity_pb(DATASET_ID, KIND, ID2)
+
+        # Make a connection to return the entity pbs.
+        connection = _Connection(entity_pb1, entity_pb2)
+
+        key1 = Key(KIND, ID1, dataset_id=DATASET_ID)
+        key2 = Key(KIND, ID2, dataset_id=DATASET_ID)
+        retrieved1, retrieved2 = self._callFUT(
+            [key1, key2], connection=connection)
+
+        # Check values match.
+        self.assertEqual(retrieved1.key.path, key1.path)
+        self.assertEqual(dict(retrieved1), {})
+        self.assertEqual(retrieved2.key.path, key2.path)
+        self.assertEqual(dict(retrieved2), {})
+
+    def test_get_entities_hit_multiple_keys_different_dataset(self):
+        from gcloud.datastore.key import Key
+
+        DATASET_ID1 = 'DATASET'
+        DATASET_ID2 = 'DATASET-ALT'
+
+        # Make sure our IDs are actually different.
+        self.assertNotEqual(DATASET_ID1, DATASET_ID2)
+
+        key1 = Key('KIND', 1234, dataset_id=DATASET_ID1)
+        key2 = Key('KIND', 1234, dataset_id=DATASET_ID2)
+        with self.assertRaises(ValueError):
+            self._callFUT([key1, key2], connection=object())
+
     def test_get_entities_implicit(self):
         from gcloud.datastore import _implicit_environ
         from gcloud.datastore.key import Key
