diff --git a/.gitignore b/.gitignore index 2790314..8e74eb8 100644 --- a/.gitignore +++ b/.gitignore @@ -29,3 +29,5 @@ pip-log.txt .tox nosetests.xml +# PyEnv +.python-version diff --git a/ndb_orm/__init__.py b/ndb_orm/__init__.py index 7c0937d..3143a65 100644 --- a/ndb_orm/__init__.py +++ b/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 diff --git a/ndb_orm/key.py b/ndb_orm/key.py index 6f8d847..b9ae71a 100644 --- a/ndb_orm/key.py +++ b/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, diff --git a/ndb_orm/model.py b/ndb_orm/model.py index 6f6cba0..a5d180d 100644 --- a/ndb_orm/model.py +++ b/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 diff --git a/tests/tests.py b/tests/tests.py index d676e62..140cc42 100644 --- a/tests/tests.py +++ b/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()