fix: make sure key.Key uses namespace from client when not specified (#339)

Fixes #337

Author: Chris Rossi

packages/google-cloud-ndb/google/cloud/ndb/key.py

@@ -100,7 +100,7 @@
 from google.cloud.ndb import utils
 
 
-__all__ = ["Key"]
+__all__ = ["Key", "UNDEFINED"]
 _APP_ID_ENVIRONMENT = "APPLICATION_ID"
 _APP_ID_DEFAULT = "_"
 _WRONG_TYPE = "Cannot construct Key reference on non-Key class; received {!r}"
@@ -126,6 +126,14 @@
     "Key name strings must be non-empty strings up to {:d} bytes; received {}"
 )
 
+UNDEFINED = object()
+"""Sentinel value.
+
+Used to indicate a namespace hasn't been explicitly set in key construction.
+Used to distinguish between not passing a value and passing `None`, which
+indicates the default namespace.
+"""
+
 
 class Key(object):
     """An immutable datastore key.
@@ -278,11 +286,15 @@ def __new__(cls, *path_args, **kwargs):
 
         _constructor_handle_positional(path_args, kwargs)
         instance = super(Key, cls).__new__(cls)
+
         # Make sure to pass in the namespace if it's not explicitly set.
-        if "namespace" not in kwargs:
+        if kwargs.get("namespace", UNDEFINED) is UNDEFINED:
             client = context_module.get_context().client
             if client.namespace:
                 kwargs["namespace"] = client.namespace
+            else:
+                kwargs["namespace"] = None  # default namespace
+
         if (
             "reference" in kwargs
             or "serialized" in kwargs

packages/google-cloud-ndb/google/cloud/ndb/model.py

@@ -4526,7 +4526,7 @@ def __init__(_self, **kwargs):
         id_ = self._get_arg(kwargs, "id")
         project = self._get_arg(kwargs, "project")
         app = self._get_arg(kwargs, "app")
-        namespace = self._get_arg(kwargs, "namespace")
+        namespace = self._get_arg(kwargs, "namespace", key_module.UNDEFINED)
         parent = self._get_arg(kwargs, "parent")
         projection = self._get_arg(kwargs, "projection")
 
@@ -4542,7 +4542,7 @@ def __init__(_self, **kwargs):
             id_ is None
             and parent is None
             and project is None
-            and namespace is None
+            and namespace is key_module.UNDEFINED
         )
         if key is not None:
             if not key_parts_unspecified:
@@ -4567,7 +4567,7 @@ def __init__(_self, **kwargs):
             self._set_projection(projection)
 
     @classmethod
-    def _get_arg(cls, kwargs, keyword):
+    def _get_arg(cls, kwargs, keyword, default=None):
         """Parse keywords for fields that aren't user-defined properties.
 
         This is used to re-map special keyword arguments in the presence
@@ -4581,9 +4581,11 @@ def _get_arg(cls, kwargs, keyword):
         Args:
             kwargs (Dict[str, Any]): A keyword arguments dictionary.
             keyword (str): A keyword to be converted.
+            default (Any): Returned if argument isn't found.
 
         Returns:
-            Optional[Any]: The ``keyword`` argument, if found.
+            Optional[Any]: The ``keyword`` argument, if found, otherwise
+                ``default``.
         """
         alt_keyword = "_" + keyword
         if alt_keyword in kwargs:
@@ -4594,7 +4596,7 @@ def _get_arg(cls, kwargs, keyword):
             if not isinstance(obj, Property) or isinstance(obj, ModelKey):
                 return kwargs.pop(keyword)
 
-        return None
+        return default
 
     def _set_attributes(self, kwargs):
         """Set attributes from keyword arguments.

packages/google-cloud-ndb/tests/system/test_query.py

@@ -301,6 +301,31 @@ class SomeKind(ndb.Model):
     assert results[0].key.namespace() == other_namespace
 
 
+def test_namespace_set_on_client_with_id(dispose_of, other_namespace):
+    """Regression test for #337
+
+    https://github.com/googleapis/python-ndb/issues/337
+    """
+
+    class SomeKind(ndb.Model):
+        foo = ndb.IntegerProperty()
+        bar = ndb.StringProperty()
+
+    client = ndb.Client(namespace=other_namespace)
+    with client.context(cache_policy=False):
+        id = test_utils.system.unique_resource_id()
+        entity1 = SomeKind(id=id, foo=1, bar="a")
+        key = entity1.put()
+        dispose_of(key._key)
+        assert key.namespace() == other_namespace
+
+        results = eventually(SomeKind.query().fetch, _length_equals(1))
+
+        assert results[0].foo == 1
+        assert results[0].bar == "a"
+        assert results[0].key.namespace() == other_namespace
+
+
 @pytest.mark.usefixtures("client_context")
 def test_filter_equal(ds_entity):
     for i in range(5):