Move Dataset.get_entity to Key.

Addresses third part of googleapis#477.

gcloud/datastore/__init__.py

@@ -146,18 +146,6 @@ def _require_dataset():
     return _implicit_environ.DATASET
 
 
-def get_entity(key):
-    """Retrieves entity from implicit dataset, along with its attributes.
-
-    :type key: :class:`gcloud.datastore.key.Key`
-    :param key: The name of the item to retrieve.
-
-    :rtype: :class:`gcloud.datastore.entity.Entity` or ``None``
-    :returns: The requested entity, or ``None`` if there was no match found.
-    """
-    return _require_dataset().get_entity(key)
-
-
 def get_entities(keys):
     """Retrieves entities from implied dataset, along with their attributes.
 

gcloud/datastore/dataset.py

@@ -72,21 +72,6 @@ def id(self):
 
         return self._id
 
-    def get_entity(self, key):
-        """Retrieves entity from the dataset, along with its attributes.
-
-        :type key: :class:`gcloud.datastore.key.Key`
-        :param key: The key of the entity to be retrieved.
-
-        :rtype: :class:`gcloud.datastore.entity.Entity` or `NoneType`
-        :returns: The requested entity, or ``None`` if there was no
-                  match found.
-        """
-        entities = self.get_entities([key])
-
-        if entities:
-            return entities[0]
-
     def get_entities(self, keys, missing=None, deferred=None):
         """Retrieves entities from the dataset, along with their attributes.
 

gcloud/datastore/entity.py

@@ -213,8 +213,7 @@ def reload(self):
           exist only locally.
         """
         key = self._must_key
-        dataset = self._must_dataset
-        entity = dataset.get_entity(key.to_protobuf())
+        entity = key.get()
 
         if entity:
             self.update(entity)

gcloud/datastore/key.py

@@ -227,6 +227,33 @@ def to_protobuf(self):
 
         return key
 
+    def get(self, connection=None):
+        """Retrieve entity corresponding to the curretn key.
+
+        :type connection: :class:`gcloud.datastore.connection.Connection`
+        :param connection: Optional connection used to connection to datastore.
+
+        :rtype: :class:`gcloud.datastore.entity.Entity` or `NoneType`
+        :returns: The requested entity, or ``None`` if there was no
+                  match found.
+        :raises: `ValueError` if the current key is partial.
+        """
+        # Temporary import hack until Dataset is removed in #477.
+        from gcloud.datastore.dataset import Dataset
+
+        if self.is_partial:
+            raise ValueError('Can only retrieve complete keys.')
+
+        connection = connection or _implicit_environ.CONNECTION
+        dataset = Dataset(self.dataset_id, connection=connection)
+        entities = dataset.get_entities([self])
+
+        if entities:
+            result = entities[0]
+            # We assume that the backend has not changed the key.
+            result.key(self)
+            return result
+
     @property
     def is_partial(self):
         """Boolean indicating if the key has an ID (or name).

gcloud/datastore/test___init__.py

@@ -166,20 +166,6 @@ def test__require_dataset(self):
         finally:
             _implicit_environ.DATASET = original_dataset
 
-    def test_get_entity(self):
-        import gcloud.datastore
-        from gcloud.datastore import _implicit_environ
-        from gcloud.datastore.test_entity import _Dataset
-        from gcloud._testing import _Monkey
-
-        CUSTOM_DATASET = _Dataset()
-        DUMMY_KEY = object()
-        DUMMY_VAL = object()
-        CUSTOM_DATASET[DUMMY_KEY] = DUMMY_VAL
-        with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET):
-            result = gcloud.datastore.get_entity(DUMMY_KEY)
-        self.assertTrue(result is DUMMY_VAL)
-
     def test_get_entities(self):
         import gcloud.datastore
         from gcloud.datastore import _implicit_environ

gcloud/datastore/test_dataset.py

@@ -40,48 +40,6 @@ def test_ctor_explicit(self):
         self.assertEqual(dataset.id(), DATASET_ID)
         self.assertTrue(dataset.connection() is CONNECTION)
 
-    def test_get_entity_miss(self):
-        from gcloud.datastore.key import Key
-        DATASET_ID = 'DATASET'
-        connection = _Connection()
-        dataset = self._makeOne(DATASET_ID, connection)
-        key = Key('Kind', 1234, dataset_id=DATASET_ID)
-        self.assertEqual(dataset.get_entity(key), None)
-
-    def test_get_entity_hit(self):
-        from gcloud.datastore.connection import datastore_pb
-        from gcloud.datastore.key import Key
-        DATASET_ID = 'DATASET'
-        KIND = 'Kind'
-        ID = 1234
-        PATH = [{'kind': KIND, 'id': ID}]
-        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 = ID
-        prop = entity_pb.property.add()
-        prop.name = 'foo'
-        prop.value.string_value = 'Foo'
-        connection = _Connection(entity_pb)
-        dataset = self._makeOne(DATASET_ID, connection)
-        key = Key(KIND, ID, dataset_id=DATASET_ID)
-        result = dataset.get_entity(key)
-        key = result.key()
-        self.assertEqual(key.dataset_id, DATASET_ID)
-        self.assertEqual(key.path, PATH)
-        self.assertEqual(list(result), ['foo'])
-        self.assertEqual(result['foo'], 'Foo')
-
-    def test_get_entity_bad_type(self):
-        DATASET_ID = 'DATASET'
-        connection = _Connection()
-        dataset = self._makeOne(DATASET_ID, connection)
-        with self.assertRaises(AttributeError):
-            dataset.get_entity(None)
-        with self.assertRaises(AttributeError):
-            dataset.get_entity([])
-
     def test_get_entities_miss(self):
         from gcloud.datastore.key import Key
         DATASET_ID = 'DATASET'

gcloud/datastore/test_entity.py

@@ -118,6 +118,7 @@ def test_reload_no_key(self):
     def test_reload_miss(self):
         dataset = _Dataset()
         key = _Key()
+        key._stored = None  # Explicit miss.
         entity = self._makeOne(dataset)
         entity.key(key)
         entity['foo'] = 'Foo'
@@ -127,13 +128,15 @@ def test_reload_miss(self):
 
     def test_reload_hit(self):
         dataset = _Dataset()
-        dataset['KEY'] = {'foo': 'Bar'}
         key = _Key()
+        NEW_VAL = 'Baz'
+        key._stored = {'foo': NEW_VAL}
         entity = self._makeOne(dataset)
         entity.key(key)
         entity['foo'] = 'Foo'
         self.assertTrue(entity.reload() is entity)
-        self.assertEqual(entity['foo'], 'Bar')
+        self.assertEqual(entity['foo'], NEW_VAL)
+        self.assertEqual(entity.keys(), ['foo'])
 
     def test_save_no_key(self):
         from gcloud.datastore.entity import NoKey
@@ -248,6 +251,7 @@ class _Key(object):
     _partial = False
     _path = None
     _id = None
+    _stored = None
 
     def to_protobuf(self):
         return self._key
@@ -260,6 +264,9 @@ def is_partial(self):
     def path(self):
         return self._path
 
+    def get(self):
+        return self._stored
+
 
 class _Dataset(dict):
 
@@ -280,9 +287,6 @@ def id(self):
     def connection(self):
         return self._connection
 
-    def get_entity(self, key):
-        return self.get(key)
-
     def get_entities(self, keys):
         return [self.get(key) for key in keys]
 

gcloud/datastore/test_key.py

@@ -195,6 +195,63 @@ def test_to_protobuf_w_no_kind(self):
         pb = key.to_protobuf()
         self.assertFalse(pb.path_element[0].HasField('kind'))
 
+    def test_get_explicit_connection_miss(self):
+        from gcloud.datastore.test_dataset import _Connection
+
+        cnxn_lookup_result = []
+        cnxn = _Connection(*cnxn_lookup_result)
+        key = self._makeOne('KIND', 1234)
+        entity = key.get(connection=cnxn)
+        self.assertEqual(entity, None)
+
+    def test_get_implicit_connection_miss(self):
+        from gcloud._testing import _Monkey
+        from gcloud.datastore import _implicit_environ
+        from gcloud.datastore.test_dataset import _Connection
+
+        cnxn_lookup_result = []
+        cnxn = _Connection(*cnxn_lookup_result)
+        key = self._makeOne('KIND', 1234)
+        with _Monkey(_implicit_environ, CONNECTION=cnxn):
+            entity = key.get()
+        self.assertEqual(entity, None)
+
+    def test_get_explicit_connection_hit(self):
+        from gcloud.datastore import datastore_v1_pb2
+        from gcloud.datastore.test_dataset import _Connection
+
+        KIND = 'KIND'
+        ID = 1234
+
+        # Make a bogus entity PB to be returned from fake Connection.
+        entity_pb = datastore_v1_pb2.Entity()
+        entity_pb.key.partition_id.dataset_id = self._DEFAULT_DATASET
+        path_element = entity_pb.key.path_element.add()
+        path_element.kind = KIND
+        path_element.id = ID
+        prop = entity_pb.property.add()
+        prop.name = 'foo'
+        prop.value.string_value = 'Foo'
+
+        # Make fake connection.
+        cnxn_lookup_result = [entity_pb]
+        cnxn = _Connection(*cnxn_lookup_result)
+
+        # Create key and look-up.
+        key = self._makeOne(KIND, ID)
+        entity = key.get(connection=cnxn)
+        self.assertEqual(entity.items(), [('foo', 'Foo')])
+        self.assertTrue(entity.key() is key)
+
+    def test_get_explicit_connection_partial_key(self):
+        from gcloud.datastore.test_dataset import _Connection
+
+        cnxn_lookup_result = []
+        cnxn = _Connection(*cnxn_lookup_result)
+        key = self._makeOne('KIND')
+        with self.assertRaises(ValueError):
+            key.get(connection=cnxn)
+
     def test_is_partial_no_name_or_id(self):
         key = self._makeOne('KIND')
         self.assertTrue(key.is_partial)

pylintrc_default

@@ -73,13 +73,16 @@ ignore =
 #                  identical implementation but different docstrings.
 # - star-args:  standard Python idioms for varargs:
 #                  ancestor = Query().filter(*order_props)
+# - cyclic-import:  temporary workaround to support Key.get until Dataset
+#                   is removed in #477.
 disable =
     maybe-no-member,
     no-member,
     protected-access,
     redefined-builtin,
     similarities,
     star-args,
+    cyclic-import,
 
 
 [REPORTS]

regression/datastore.py

@@ -29,6 +29,7 @@
 
 DATASET_ID = os.getenv('GCLOUD_TESTS_DATASET_ID')
 datastore.set_default_dataset(dataset_id=DATASET_ID)
+datastore.set_default_connection()
 
 
 class TestDatastore(unittest2.TestCase):
@@ -97,11 +98,9 @@ def _generic_test_post(self, name=None, key_id=None):
             self.assertEqual(entity.key().name, name)
         if key_id is not None:
             self.assertEqual(entity.key().id, key_id)
-        retrieved_entity = datastore.get_entity(entity.key())
+        retrieved_entity = entity.key().get()
         # Check the keys are the same.
-        self.assertEqual(retrieved_entity.key().path, entity.key().path)
-        self.assertEqual(retrieved_entity.key().namespace,
-                         entity.key().namespace)
+        self.assertTrue(retrieved_entity.key() is entity.key())
 
         # Check the data is the same.
         retrieved_dict = dict(retrieved_entity.items())
@@ -352,13 +351,13 @@ def test_transaction(self):
         entity['url'] = u'www.google.com'
 
         with Transaction():
-            retrieved_entity = datastore.get_entity(key)
+            retrieved_entity = key.get()
             if retrieved_entity is None:
                 entity.save()
                 self.case_entities_to_delete.append(entity)
 
         # This will always return after the transaction.
-        retrieved_entity = datastore.get_entity(key)
+        retrieved_entity = key.get()
         retrieved_dict = dict(retrieved_entity.items())
         entity_dict = dict(entity.items())
         self.assertEqual(retrieved_dict, entity_dict)