diff --git a/google/cloud/ndb/model.py b/google/cloud/ndb/model.py index dcee3fa6..83cf867a 100644 --- a/google/cloud/ndb/model.py +++ b/google/cloud/ndb/model.py @@ -3022,58 +3022,14 @@ def auth_domain(self): """ return self._auth_domain - def add_to_entity(self, entity, name): - """Add the user value to a datastore entity. - - .. note:: - - This assumes, but does not check, that ``name`` is not already - set on ``entity`` or in the meanings of ``entity``. - - Args: - entity (~google.cloud.datastore.entity.Entity): An entity that - contains a user value as the field ``name``. - name (str): The name of the field containing this user value. - """ - user_entity = ds_entity_module.Entity() - entity[name] = user_entity - entity._meanings[name] = (_MEANING_PREDEFINED_ENTITY_USER, user_entity) - - # Set required fields. - user_entity["email"] = self._email - user_entity.exclude_from_indexes.add("email") - user_entity["auth_domain"] = self._auth_domain - user_entity.exclude_from_indexes.add("auth_domain") - # Set optional field. - if self._user_id: - user_entity["user_id"] = self._user_id - user_entity.exclude_from_indexes.add("user_id") - @classmethod - def read_from_entity(cls, entity, name): + def _from_ds_entity(cls, user_entity): """Convert the user value to a datastore entity. Args: - entity (~google.cloud.datastore.entity.Entity): An entity that - contains a user value as the field ``name``. - name (str): The name of the field containing this user value. - - Raises: - ValueError: If the stored meaning for the ``name`` field is not - equal to ``ENTITY_USER=20``. - ValueError: If the value stored in the meanings for ``entity`` - is not the actual stored value under ``name``. - """ - # NOTE: This may fail in a ``KeyError``. - user_entity = entity[name] - # NOTE: This may result in a ``ValueError`` for failed unpacking. - meaning, value = entity._meanings.get(name, (0, None)) - if meaning != _MEANING_PREDEFINED_ENTITY_USER: - raise ValueError("User values should have meaning=20") - if user_entity is not value: - raise ValueError("Unexpected value stored for meaning") - - # NOTE: We do not check ``exclude_from_indexes``. + user_entity (~google.cloud.datastore.entity.Entity): A user value + datastore entity. + """ kwargs = { "email": user_entity["email"], "_auth_domain": user_entity["auth_domain"], @@ -3235,7 +3191,8 @@ def _validate(self, value): Raises: .BadValueError: If ``value`` is not a :class:`User`. """ - if not isinstance(value, User): + # Might be GAE User or our own version + if type(value).__name__ != "User": raise exceptions.BadValueError( "Expected User, got {!r}".format(value) ) @@ -3251,6 +3208,59 @@ def _prepare_for_put(self, entity): entity (Model): An entity with values. """ + def _to_base_type(self, value): + """Convert the user value to a datastore entity. + + Arguments: + value (User): The user value. + + Returns: + ~google.cloud.datastore.entity.Entity: The datastore entity. + """ + user_entity = ds_entity_module.Entity() + + # Set required fields. + user_entity["email"] = six.ensure_text(value.email()) + user_entity.exclude_from_indexes.add("email") + user_entity["auth_domain"] = six.ensure_text(value.auth_domain()) + user_entity.exclude_from_indexes.add("auth_domain") + # Set optional field. + user_id = value.user_id() + if user_id: + user_entity["user_id"] = six.ensure_text(user_id) + user_entity.exclude_from_indexes.add("user_id") + + return user_entity + + def _from_base_type(self, ds_entity): + """Convert the user value from a datastore entity. + + Arguments: + ds_entity (~google.cloud.datastore.entity.Entity): The datastore + entity. + + Returns: + User: The converted entity. + """ + return User._from_ds_entity(ds_entity) + + def _to_datastore(self, entity, data, prefix="", repeated=False): + """Override of :method:`Property._to_datastore`. + + We just need to set the meaning to indicate value is a User. + """ + keys = super(UserProperty, self)._to_datastore( + entity, data, prefix=prefix, repeated=repeated + ) + + for key in keys: + value = data.get(key) + if value: + data.setdefault("_meanings", {})[key] = ( + _MEANING_PREDEFINED_ENTITY_USER, + value, + ) + class KeyProperty(Property): """A property that contains :class:`.Key` values. diff --git a/tests/system/test_crud.py b/tests/system/test_crud.py index f4c61a7e..c6d1f56d 100644 --- a/tests/system/test_crud.py +++ b/tests/system/test_crud.py @@ -1056,3 +1056,42 @@ class CModel(ndb.Model): entity = key_c.get() assert entity.key_a == key_a assert entity.key_b == key_b + + +@pytest.mark.usefixtures("client_context") +def test_user_property(dispose_of): + class SomeKind(ndb.Model): + user = ndb.UserProperty() + + user = ndb.User("somebody@example.com", "gmail.com") + entity = SomeKind(user=user) + key = entity.put() + dispose_of(key._key) + + retreived = key.get() + assert retreived.user.email() == "somebody@example.com" + assert retreived.user.auth_domain() == "gmail.com" + + +@pytest.mark.usefixtures("client_context") +def test_user_property_different_user_class(dispose_of): + class SomeKind(ndb.Model): + user = ndb.UserProperty() + + class User(object): + def email(self): + return "somebody@example.com" + + def auth_domain(self): + return "gmail.com" + + def user_id(self): + return None + + entity = SomeKind(user=User()) + key = entity.put() + dispose_of(key._key) + + retreived = key.get() + assert retreived.user.email() == "somebody@example.com" + assert retreived.user.auth_domain() == "gmail.com" diff --git a/tests/unit/test_model.py b/tests/unit/test_model.py index 943c4ff5..c29d3733 100644 --- a/tests/unit/test_model.py +++ b/tests/unit/test_model.py @@ -2227,108 +2227,6 @@ def test_auth_domain(self): user_value = self._make_default() assert user_value.auth_domain() == "example.com" - @staticmethod - def _add_to_entity_helper(user_value): - entity = entity_module.Entity() - name = "u" - - user_value.add_to_entity(entity, name) - assert list(entity.keys()) == [name] - user_entity = entity[name] - assert entity._meanings == { - name: (model._MEANING_PREDEFINED_ENTITY_USER, user_entity) - } - assert user_entity["email"] == user_value._email - assert user_entity["auth_domain"] == user_value._auth_domain - return user_entity - - def test_add_to_entity(self): - user_value = self._make_default() - user_entity = self._add_to_entity_helper(user_value) - assert sorted(user_entity.keys()) == ["auth_domain", "email"] - assert user_entity.exclude_from_indexes == set( - ["auth_domain", "email"] - ) - - def test_add_to_entity_with_user_id(self): - user_value = model.User( - email="foo@example.com", - _auth_domain="example.com", - _user_id="197382", - ) - user_entity = self._add_to_entity_helper(user_value) - assert sorted(user_entity.keys()) == [ - "auth_domain", - "email", - "user_id", - ] - assert user_entity["user_id"] == user_value._user_id - assert user_entity.exclude_from_indexes == set( - ["auth_domain", "email", "user_id"] - ) - - @staticmethod - def _prepare_entity(name, email, auth_domain): - entity = entity_module.Entity() - user_entity = entity_module.Entity() - - entity[name] = user_entity - entity._meanings[name] = ( - model._MEANING_PREDEFINED_ENTITY_USER, - user_entity, - ) - user_entity.exclude_from_indexes.update(["auth_domain", "email"]) - user_entity["auth_domain"] = auth_domain - user_entity["email"] = email - - return entity - - def test_read_from_entity(self): - name = "you_sir" - email = "foo@example.com" - auth_domain = "example.com" - entity = self._prepare_entity(name, email, auth_domain) - - user_value = model.User.read_from_entity(entity, name) - assert user_value._auth_domain == auth_domain - assert user_value._email == email - assert user_value._user_id is None - - def test_read_from_entity_bad_meaning(self): - name = "you_sir" - email = "foo@example.com" - auth_domain = "example.com" - entity = self._prepare_entity(name, email, auth_domain) - - # Wrong meaning. - entity._meanings[name] = ("not-20", entity[name]) - with pytest.raises(ValueError): - model.User.read_from_entity(entity, name) - - # Wrong associated value. - entity._meanings[name] = (model._MEANING_PREDEFINED_ENTITY_USER, None) - with pytest.raises(ValueError): - model.User.read_from_entity(entity, name) - - # No meaning. - entity._meanings.clear() - with pytest.raises(ValueError): - model.User.read_from_entity(entity, name) - - def test_read_from_entity_with_user_id(self): - name = "you_sir" - email = "foo@example.com" - auth_domain = "example.com" - entity = self._prepare_entity(name, email, auth_domain) - entity[name].exclude_from_indexes.add("user_id") - user_id = "80131394" - entity[name]["user_id"] = user_id - - user_value = model.User.read_from_entity(entity, name) - assert user_value._auth_domain == auth_domain - assert user_value._email == email - assert user_value._user_id == user_id - def test___str__(self): user_value = self._make_default() assert str(user_value) == "foo" @@ -2380,6 +2278,22 @@ def test___lt__(self): with pytest.raises(TypeError): user_value1 < user_value4 + @staticmethod + def test__from_ds_entity(): + assert model.User._from_ds_entity( + {"email": "foo@example.com", "auth_domain": "gmail.com"} + ) == model.User("foo@example.com", "gmail.com") + + @staticmethod + def test__from_ds_entity_with_user_id(): + assert model.User._from_ds_entity( + { + "email": "foo@example.com", + "auth_domain": "gmail.com", + "user_id": "12345", + } + ) == model.User("foo@example.com", "gmail.com", "12345") + class TestUserProperty: @staticmethod @@ -2429,6 +2343,58 @@ def test__db_get_value(): with pytest.raises(NotImplementedError): prop._db_get_value(None, None) + @staticmethod + def test__to_base_type(): + prop = model.UserProperty(name="u") + entity = prop._to_base_type(model.User("email", "auth_domain",)) + assert entity["email"] == "email" + assert "email" in entity.exclude_from_indexes + assert entity["auth_domain"] == "auth_domain" + assert "auth_domain" in entity.exclude_from_indexes + assert "user_id" not in entity + + @staticmethod + def test__to_base_type_w_user_id(): + prop = model.UserProperty(name="u") + entity = prop._to_base_type( + model.User("email", "auth_domain", "user_id") + ) + assert entity["email"] == "email" + assert "email" in entity.exclude_from_indexes + assert entity["auth_domain"] == "auth_domain" + assert "auth_domain" in entity.exclude_from_indexes + assert entity["user_id"] == "user_id" + assert "user_id" in entity.exclude_from_indexes + + @staticmethod + def test__from_base_type(): + prop = model.UserProperty(name="u") + assert prop._from_base_type( + {"email": "email", "auth_domain": "auth_domain"} + ) == model.User("email", "auth_domain") + + @staticmethod + def test__to_datastore(): + class SomeKind(model.Model): + u = model.UserProperty() + + entity = SomeKind(u=model.User("email", "auth_domain")) + data = {} + SomeKind.u._to_datastore(entity, data) + meaning, ds_entity = data["_meanings"]["u"] + assert meaning == model._MEANING_PREDEFINED_ENTITY_USER + assert data["u"] == ds_entity + + @staticmethod + def test__to_datastore_no_value(): + class SomeKind(model.Model): + u = model.UserProperty() + + entity = SomeKind() + data = {} + SomeKind.u._to_datastore(entity, data) + assert data == {"u": None} + class TestKeyProperty: @staticmethod