Fix key property.

src/google/cloud/ndb/model.py

@@ -252,7 +252,8 @@ class Person(Model):
 import pickle
 import zlib
 
-from google.cloud.datastore import entity as entity_module
+from google.cloud.datastore import key as ds_key_module
+from google.cloud.datastore import entity as ds_entity_module
 from google.cloud.datastore import helpers
 from google.cloud.datastore_v1.proto import entity_pb2
 
@@ -543,11 +544,11 @@ def _entity_from_ds_entity(ds_entity, model_class=None):
                     key = key_module.Key(kind, None)
                     if structprop._repeated:
                         value = [
-                            _BaseValue(entity_module.Entity(key._key))
+                            _BaseValue(ds_entity_module.Entity(key._key))
                             for _ in subvalue
                         ]
                     else:
-                        value = entity_module.Entity(key._key)
+                        value = ds_entity_module.Entity(key._key)
                         value = _BaseValue(value)
 
                     structprop._store_value(entity, value)
@@ -644,11 +645,11 @@ def _entity_to_ds_entity(entity, set_key=True):
         if key is None:
             # use _class_name instead of _get_kind, to get PolyModel right
             key = key_module.Key(entity._class_name(), None)
-        ds_entity = entity_module.Entity(
+        ds_entity = ds_entity_module.Entity(
             key._key, exclude_from_indexes=exclude_from_indexes
         )
     else:
-        ds_entity = entity_module.Entity(
+        ds_entity = ds_entity_module.Entity(
             exclude_from_indexes=exclude_from_indexes
         )
     ds_entity.update(data)
@@ -2855,7 +2856,7 @@ def add_to_entity(self, entity, name):
                 contains a user value as the field ``name``.
             name (str): The name of the field containing this user value.
         """
-        user_entity = entity_module.Entity()
+        user_entity = ds_entity_module.Entity()
         entity[name] = user_entity
         entity._meanings[name] = (_MEANING_PREDEFINED_ENTITY_USER, user_entity)
 
@@ -3321,6 +3322,36 @@ def _db_get_value(self, v, unused_p):
         """
         raise exceptions.NoLongerImplementedError()
 
+    def _to_base_type(self, value):
+        """Convert a value to the "base" value type for this property.
+
+        Args:
+            value (~key.Key): The value to be converted.
+
+        Returns:
+            google.cloud.datastore.Key: The converted value.
+
+        Raises:
+            TypeError: If ``value`` is not a :class:`~key.Key`.
+        """
+        if not isinstance(value, key_module.Key):
+            raise TypeError(
+                "Cannot convert to datastore key, expected Key value; "
+                "received {}".format(value)
+            )
+        return value._key
+
+    def _from_base_type(self, value):
+        """Convert a value from the "base" value type for this property.
+
+        Args:
+            value (google.cloud.datastore.Key): The value to be converted.
+
+        Returns:
+            key.Key: The converted value.
+        """
+        return key_module.Key._from_ds_key(value)
+
 
 class BlobKeyProperty(Property):
     """A property containing :class:`~google.cloud.ndb.model.BlobKey` values.
@@ -3873,7 +3904,7 @@ def _from_base_type(self, value):
         Returns:
             The converted value with given class.
         """
-        if isinstance(value, entity_module.Entity):
+        if isinstance(value, ds_entity_module.Entity):
             value = _entity_from_ds_entity(
                 value, model_class=self._model_class
             )

tests/system/test_crud.py

@@ -210,6 +210,21 @@ class SomeKind(ndb.Model):
     dispose_of(key._key)
 
 
+@pytest.mark.usefixtures("client_context")
+def test_key_property(dispose_of, ds_client):
+    class SomeKind(ndb.Model):
+        foo = ndb.KeyProperty()
+
+    key_value = ndb.Key("Whatevs", 123)
+    entity = SomeKind(foo=key_value)
+    key = entity.put()
+
+    retrieved = key.get()
+    assert retrieved.foo == key_value
+
+    dispose_of(key._key)
+
+
 def test_insert_entity_with_caching(dispose_of, client_context):
     class SomeKind(ndb.Model):
         foo = ndb.IntegerProperty()

tests/unit/test_model.py

@@ -20,6 +20,7 @@
 
 from google.cloud import datastore
 from google.cloud.datastore import entity as entity_module
+from google.cloud.datastore import key as ds_key_module
 from google.cloud.datastore import helpers
 from google.cloud.datastore_v1 import types as ds_types
 from google.cloud.datastore_v1.proto import entity_pb2
@@ -2316,6 +2317,30 @@ def test__db_get_value():
         with pytest.raises(NotImplementedError):
             prop._db_get_value(None, None)
 
+    @staticmethod
+    @pytest.mark.usefixtures("in_context")
+    def test__to_base_type():
+        prop = model.KeyProperty("keyp")
+        value = key_module.Key("Kynd", 123)
+        assert prop._to_base_type(value) is value._key
+
+    @staticmethod
+    @pytest.mark.usefixtures("in_context")
+    def test__to_base_type_wrong_type():
+        prop = model.KeyProperty("keyp")
+        value = ("Kynd", 123)
+        with pytest.raises(TypeError):
+            assert prop._to_base_type(value) is value._key
+
+    @staticmethod
+    @pytest.mark.usefixtures("in_context")
+    def test__from_base_type():
+        prop = model.KeyProperty("keyp")
+        ds_value = ds_key_module.Key("Kynd", 123, project="testing")
+        value = prop._from_base_type(ds_value)
+        assert value.kind() == "Kynd"
+        assert value.id() == 123
+
 
 class TestBlobKeyProperty:
     @staticmethod