fix: do not allow empty key parts for key constructor in namespaced model (#401)

cguardia · web-flow · commit 4fbcf7c598c8 · 2020-04-28T20:25:35.000-05:00
* fix: do not allow empty key parts for key constructor in namespaced model refs #384
diff --git a/packages/google-cloud-ndb/google/cloud/ndb/model.py b/packages/google-cloud-ndb/google/cloud/ndb/model.py
@@ -4751,13 +4751,13 @@ def __init__(_self, **kwargs):
             project = app
 
         key_parts_unspecified = (
-            id_ is None
-            and parent is None
-            and project is None
-            and namespace is key_module.UNDEFINED
+            id_ is None and parent is None and project is None
         )
         if key is not None:
-            if not key_parts_unspecified:
+            if (
+                not key_parts_unspecified
+                or namespace is not key_module.UNDEFINED
+            ):
                 raise exceptions.BadArgumentError(
                     "Model constructor given 'key' does not accept "
                     "'id', 'project', 'app', 'namespace', or 'parent'."
diff --git a/packages/google-cloud-ndb/tests/system/test_crud.py b/packages/google-cloud-ndb/tests/system/test_crud.py
@@ -19,6 +19,7 @@
 import functools
 import operator
 import os
+import pickle
 import random
 import threading
 import zlib
@@ -1343,3 +1344,39 @@ class SomeKind(ndb.Model):
 
         retrieved = key.get()
         assert isinstance(retrieved.entry, OtherKind)
+
+
+@pytest.mark.usefixtures("client_context")
+def test_serialization(dispose_of):
+    """Regression test for #384
+
+    https://github.com/googleapis/python-ndb/issues/384
+    """
+
+    # THis is needed because pickle can't serialize local objects
+    global SomeKind, OtherKind
+
+    class OtherKind(ndb.Model):
+        foo = ndb.IntegerProperty()
+
+        @classmethod
+        def _get_kind(cls):
+            return "OtherKind"
+
+    class SomeKind(ndb.Model):
+        other = ndb.StructuredProperty(OtherKind)
+
+        @classmethod
+        def _get_kind(cls):
+            return "SomeKind"
+
+    entity = SomeKind(
+        other=OtherKind(foo=1, namespace="Test"), namespace="Test"
+    )
+    key = entity.put()
+    dispose_of(key._key)
+
+    retrieved = key.get()
+    assert retrieved.other.key is None or retrieved.other.key.id() is None
+    entity = pickle.loads(pickle.dumps(retrieved))
+    assert entity.other.foo == 1
diff --git a/packages/google-cloud-ndb/tests/system/test_metadata.py b/packages/google-cloud-ndb/tests/system/test_metadata.py
@@ -39,11 +39,11 @@ class AnyKind(ndb.Model):
     class MyKind(ndb.Model):
         bar = ndb.StringProperty()
 
-    entity1 = AnyKind(foo=1, namespace="_test_namespace_")
+    entity1 = AnyKind(foo=1, id="x", namespace="_test_namespace_")
     entity1.put()
     dispose_of(entity1.key._key)
 
-    entity2 = MyKind(bar="x", namespace="_test_namespace_")
+    entity2 = MyKind(bar="x", id="x", namespace="_test_namespace_")
     entity2.put()
     dispose_of(entity2.key._key)
 
diff --git a/packages/google-cloud-ndb/tests/system/test_query.py b/packages/google-cloud-ndb/tests/system/test_query.py
@@ -284,11 +284,11 @@ class SomeKind(ndb.Model):
         foo = ndb.IntegerProperty()
         bar = ndb.StringProperty()
 
-    entity1 = SomeKind(foo=1, bar="a", namespace=other_namespace)
+    entity1 = SomeKind(foo=1, bar="a", id="x", namespace=other_namespace)
     entity1.put()
     dispose_of(entity1.key._key)
 
-    entity2 = SomeKind(foo=2, bar="b")
+    entity2 = SomeKind(foo=2, bar="b", id="x")
     entity2.put()
     dispose_of(entity2.key._key)
 
diff --git a/packages/google-cloud-ndb/tests/unit/test_model.py b/packages/google-cloud-ndb/tests/unit/test_model.py
@@ -4043,6 +4043,12 @@ def test_constructor_key_parts():
         key = key_module.Key("Model", 124)
         assert entity.__dict__ == {"_values": {}, "_entity_key": key}
 
+    @staticmethod
+    @pytest.mark.usefixtures("in_context")
+    def test_constructor_namespace_no_key_parts():
+        entity = model.Model(namespace="myspace")
+        assert entity.__dict__ == {"_values": {}}
+
     @staticmethod
     @pytest.mark.usefixtures("in_context")
     def test_constructor_app():
@@ -4070,6 +4076,13 @@ def test_constructor_key_and_key_parts():
         with pytest.raises(exceptions.BadArgumentError):
             model.Model(key=key, id=124)
 
+    @staticmethod
+    @pytest.mark.usefixtures("in_context")
+    def test_constructor_key_and_key_parts_with_namespace():
+        key = key_module.Key("Foo", "bar")
+        with pytest.raises(exceptions.BadArgumentError):
+            model.Model(key=key, namespace="myspace")
+
     @staticmethod
     def test_constructor_user_property_collision():
         class SecretMap(model.Model):
@@ -5754,6 +5767,34 @@ def test_get_indexes():
         model.get_indexes()
 
 
+@pytest.mark.usefixtures("in_context")
+def test_serialization():
+
+    # THis is needed because pickle can't serialize local objects
+    global SomeKind, OtherKind
+
+    class OtherKind(model.Model):
+        foo = model.IntegerProperty()
+
+        @classmethod
+        def _get_kind(cls):
+            return "OtherKind"
+
+    class SomeKind(model.Model):
+        other = model.StructuredProperty(OtherKind)
+
+        @classmethod
+        def _get_kind(cls):
+            return "SomeKind"
+
+    entity = SomeKind(
+        other=OtherKind(foo=1, namespace="Test"), namespace="Test"
+    )
+    assert entity.other.key is None or entity.other.key.id() is None
+    entity = pickle.loads(pickle.dumps(entity))
+    assert entity.other.foo == 1
+
+
 def ManyFieldsFactory():
     """Model type class factory.
 
