Improve Keys and KeyProperty

.gitignore

@@ -29,3 +29,5 @@ pip-log.txt
 .tox
 nosetests.xml
 
+# PyEnv
+.python-version

ndb_orm/__init__.py

@@ -65,7 +65,7 @@
 real_entity_from_protobuf = None
 real_entity_to_protobuf = None
 
-def enable_use_with_gcd(project=None, namespace=None):
+def enable_use_with_gcd(project=None, namespace=None, client=None):
   from google.cloud import datastore
   from google.cloud.datastore.key import Key as DatastoreKey
   from google.cloud.datastore_v1.proto import entity_pb2
@@ -92,11 +92,18 @@ def model_from_protobuf_datastore(pb):
       return None
     entity = modelclass._from_pb(pb, key=key, set_key=False)
     #entity = modelclass._from_pb(pb, key=key, set_key=True)
-    entity.key = key
+    # NOTE(cmiN): Make sure we use the same augmented custom Key class.
+    entity.key = Key(
+      *key._flat_path,
+      parent=key._parent,
+      namespace=key._namespace,
+      project=key._project
+    )
     return entity
 
-  def model_to_protobuf_datastore(entity_of_ndb_model, project, namespace=None):
-    if namespace and entity_of_ndb_model._key and (entity_of_ndb_model._key.namespace == None):
+  def model_to_protobuf_datastore(entity_of_ndb_model, project, namespace=namespace):
+    if namespace and entity_of_ndb_model._key and (
+            entity_of_ndb_model._key.namespace is None):
         # add namespace
         entity_of_ndb_model._key._namespace = namespace
     entity_of_ndb_model._prepare_for_put()
@@ -142,7 +149,28 @@ def new_entity_to_protobuf(entity):
     datastore.helpers.entity_from_protobuf = real_entity_from_protobuf
     datastore.helpers.entity_to_protobuf = real_entity_to_protobuf
 
-  key_module.KeyBase = DatastoreKey
+  class KeyBase(DatastoreKey):
+
+    """Custom Key class that implements missing legacy methods."""
+
+    @property
+    def client(self):
+      if client:
+        return client
+      raise NotImplementedError(
+        "`KeyBase` class doesn't have a client associated with it; "
+        "please provide a `client` too when enabling with GCD"
+      )
+
+    def get(self):
+      """Get the entity object by using the client with its key."""
+      return self.client.get(self)
+
+    def delete(self):
+      """Remove the entity object by using the client with its key."""
+      self.client.delete(self)
+
+  key_module.KeyBase = KeyBase
   entity_module.Entity = entity_pb2.Entity
   entity_module.Property = entity_module.Property
   entity_module.Reference = entity_module.Reference

ndb_orm/key.py

@@ -45,7 +45,7 @@ def __call__(self, model_cls, *path_args, **kwargs):
       # accept both, plain strings and object instances as path
       if not isinstance(model_cls, six.string_types):
         path_args[i] = path_args[i]._get_kind()
-    
+
     return KeyBase(
       model_cls_str,
       *path_args,

ndb_orm/model.py

@@ -1995,7 +1995,11 @@ def _db_get_value(self, v):
     path_flat = []
     for elm in v.key_value.path:
       path_flat.extend([elm.kind, elm.name if elm.name != '' else elm.id])
-    return key_module.Key(*path_flat, project=v.key_value.partition_id.project_id)
+    return key_module.Key(
+      *path_flat,
+      project=v.key_value.partition_id.project_id,
+      namespace=v.key_value.partition_id.namespace_id or None
+    )
 
 
 class BlobKeyProperty(Property):
@@ -2077,7 +2081,7 @@ def _db_set_value(self, v, value):
 #       raise NotImplementedError('DatetimeProperty %s can only support UTC. '
 #                                 'Please derive a new Property to support '
 #                                 'alternative timezones.' % self._name)
-    
+
     v.timestamp_value.CopyFrom(helpers.datetime_to_pb_timestamp(value))
 
     # TODO old style ndb write - not possible anymore !
@@ -2913,7 +2917,7 @@ def __init__(*args, **kwds):
       self._key = _validate_key(key, entity=self)
 #     elif (id is not None or parent is not None or
 #           project is not None or namespace is not None):
-  
+
     else:
       path = [id] if id else [] # path
       self._key = key_module.Key(
@@ -3191,14 +3195,14 @@ def _from_pb(cls, pb, set_key=True, ent=None, key=None):
     # iterate over properties
     projection = []
     for name, p in six.iteritems(pb.properties):
-      # TODO 
+      # TODO
       if p.meaning == PROPERTY_INDEX_VALUE:
         projection.append(name)
       indexed = not p.exclude_from_indexes
       prop = ent._get_property_for(name, p, indexed)
       prop._deserialize(name, ent, p)
 
-    # TODO 
+    # TODO
 #     ent._set_projection(projection)
     return ent
 

tests/tests.py

@@ -1,11 +1,13 @@
 # -*- coding: utf-8 -*-
 # vim: sts=2:ts=2:sw=2
 
+import unittest
 from unittest import TestCase, main
 
 import os
 import binascii
 import ndb_orm as ndb
+from ndb_orm import key_module
 from protorpc import messages
 from google.cloud.datastore_v1.proto import entity_pb2
 from . import person_pb2
@@ -412,5 +414,54 @@ def test_repeated_structuredproperty(self):
     self.assertEqual(foo_recovered.a[2].b.d, 3)
 
 
+class KeyDepartment(Department):
+
+  dep_key = ndb.KeyProperty(kind=Department, required=True)
+
+
+@unittest.skipIf(not USE_DATASTORE, "not supported without Datastore emulator")
+class TestEntityKey(TestCase):
+
+  NAMESPACE = "test-namespace"
+
+  @classmethod
+  def setUpClass(cls):
+      from google.cloud import datastore
+      from .gcloud_credentials import EmulatorCredentials
+      cls.client = datastore.Client(project=PROJECT, credentials=EmulatorCredentials())
+      ndb.enable_use_with_gcd(project=PROJECT, namespace=cls.NAMESPACE, client=cls.client)
+
+  def setUp(self):
+    self._dep = Department(name="test_dep")
+    self.client.put(self._dep)
+    self._key_dep = KeyDepartment(name="test_key_dep", dep_key=self._dep.key)
+    self.client.put(self._key_dep)
+
+  def tearDown(self):
+    self.client.delete(self._dep.key)
+    self.client.delete(self._key_dep.key)
+
+  def test_retrieval(self):
+    self.assertEqual("test_key_dep", self._key_dep.key.get().name)
+
+  def test_self_retrieval(self):
+    dep = self._dep.key.get()
+    self.assertEqual("test_dep", dep.name)
+    self_key = dep.key
+    self.assertIsInstance(self_key, key_module.KeyBase)
+    self_dep = self_key.get()
+    self.assertEqual("test_dep", self_dep.name)
+
+  def test_removal(self):
+    self.assertTrue(self._key_dep.key.get())
+    self._key_dep.key.delete()
+    self.assertFalse(self._key_dep.key.get())
+
+  def test_keyprop_namespace(self):
+    key = self._key_dep.key.get().dep_key
+    self.assertEqual(self.NAMESPACE, key.namespace)
+    self.assertEqual("test_dep", key.get().name)
+
+
 if __name__ == '__main__':
   main()