fix: location prefix support

google/cloud/ndb/key.py

@@ -454,6 +454,10 @@ def __setstate__(self, state):
 
         flat = _get_path(None, kwargs["pairs"])
         project = _project_from_app(kwargs["app"])
+        location_prefix = _location_prefix_from_app(kwargs["app"])
+        if location_prefix:
+            project = "{}~{}".format(location_prefix, project)
+
         self._key = _key_module.Key(
             *flat, project=project, namespace=kwargs["namespace"]
         )
@@ -729,6 +733,22 @@ def urlsafe(self):
         raw_bytes = self.serialized()
         return base64.urlsafe_b64encode(raw_bytes).strip(b"=")
 
+    def legacy_urlsafe(self, location_prefix=None):
+        """A ``Reference`` protobuf encoded as urlsafe base 64.
+
+        .. doctest:: key-urlsafe
+
+            >>> key = ndb.Key("Kind", 1337, project="example")
+            >>> key.legacy_urlsafe()
+            b'agdleGFtcGxlcgsLEgRLaW5kGLkKDA'
+        """
+        return google.cloud.datastore.Key(
+            self._key.kind,
+            self._key.id,
+            namespace=self._key.namespace,
+            project=self._key.project,
+        ).to_legacy_urlsafe(location_prefix=location_prefix)
+
     @_options.ReadOptions.options
     @utils.positional(1)
     def get(
@@ -1071,6 +1091,33 @@ def _project_from_app(app, allow_empty=False):
     return parts[-1]
 
 
+def _location_prefix_from_app(app):
+    """Get the location prefix from a legacy Google App Engine
+    app string.
+
+    Args:
+        app (str): The application value to be used. If the caller passes
+            :data:`None` and ``allow_empty`` is :data:`False`, then this will
+            use the project set by the current client context. (See
+            :meth:`~client.Client.context`.)
+
+    Returns:
+        str: The location prefix.
+    """
+    # Avoid circular import in Python 2.7
+    from google.cloud.ndb import context as context_module
+
+    if app is None:
+        client = context_module.get_context().client
+        app = client.project
+
+    parts = app.split("~", 1)
+    if len(parts) == 2:
+        return parts[0]
+
+    return ""
+
+
 def _from_reference(reference, app, namespace):
     """Convert Reference protobuf to :class:`~google.cloud.datastore.key.Key`.
 
@@ -1100,6 +1147,7 @@ def _from_reference(reference, app, namespace):
             ``reference.name_space``.
     """
     project = _project_from_app(reference.app)
+    location_prefix = _location_prefix_from_app(reference.app)
     if app is not None:
         if _project_from_app(app) != project:
             raise RuntimeError(
@@ -1117,6 +1165,10 @@ def _from_reference(reference, app, namespace):
 
     _key_module._check_database_id(reference.database_id)
     flat_path = _key_module._get_flat_path(reference.path)
+
+    if location_prefix:
+        project = "{}~{}".format(location_prefix, project)
+
     return google.cloud.datastore.Key(
         *flat_path, project=project, namespace=parsed_namespace
     )
@@ -1340,15 +1392,20 @@ def _parse_from_args(
     parent_ds_key = None
     if parent is None:
         project = _project_from_app(app)
+        location_prefix = _location_prefix_from_app(app)
     else:
         project = _project_from_app(app, allow_empty=True)
+        location_prefix = _location_prefix_from_app(app)
         if not isinstance(parent, Key):
             raise exceptions.BadValueError(
                 "Expected Key instance, got {!r}".format(parent)
             )
         # Offload verification of parent to ``google.cloud.datastore.Key()``.
         parent_ds_key = parent._key
 
+    if location_prefix:
+        project = "{}~{}".format(location_prefix, project)
+
     return google.cloud.datastore.Key(
         *flat, parent=parent_ds_key, project=project, namespace=namespace
     )

noxfile.py

@@ -76,7 +76,7 @@ def cover(session):
     session.run("coverage", "erase")
 
 
-def run_black(session, use_check=False):
+def run_black(use_check=False):
     args = ["black"]
     if use_check:
         args.append("--check")

tests/system/test_crud.py

@@ -1034,7 +1034,7 @@ class CModel(ndb.Model):
         s_foobar = ndb.StringProperty()
         key_b = ndb.KeyProperty(kind="BModel", indexed=True)
         key_a = ndb.ComputedProperty(  # Issue here
-            lambda self: self.key_b.get().key_a if self.key_b else None,
+            lambda self: self.key_b.get().key_a if self.key_b else None
         )
 
     key_a = AModel(s_foo="test").put()

tests/unit/test_key.py

@@ -567,6 +567,14 @@ def test_urlsafe():
         key = key_module.Key("d", None, app="f")
         assert key.urlsafe() == b"agFmcgULEgFkDA"
 
+    @staticmethod
+    @pytest.mark.usefixtures("in_context")
+    def test_legacy_urlsafe():
+        key = key_module.Key("d", 123, app="f")
+        assert (
+            key.legacy_urlsafe(location_prefix="s~") == b"agNzfmZyBwsSAWQYeww"
+        )
+
     @staticmethod
     @pytest.mark.usefixtures("in_context")
     @mock.patch("google.cloud.ndb._datastore_api")
@@ -847,6 +855,29 @@ def test_app_fallback(context):
             assert key_module._project_from_app(None) == "jectpro"
 
 
+class Test__location_prefix_from_app:
+    @staticmethod
+    def test_has_prefix():
+        app = "my-prahjekt"
+        prefix = ""
+        assert key_module._location_prefix_from_app(app) == prefix
+
+    @staticmethod
+    def test_no_prefix():
+        project = "my-prahjekt"
+        for prefix in ("s", "e", "dev"):
+            app = "{}~{}".format(prefix, project)
+            assert key_module._location_prefix_from_app(app) == prefix
+
+    @staticmethod
+    def test_app_fallback(context):
+        prefix = "s"
+        context.client.project = "{}~jectpro".format(prefix)
+
+        with context.use():
+            assert key_module._location_prefix_from_app(None) == "s"
+
+
 class Test__from_reference:
     def test_basic(self):
         reference = make_reference()

tests/unit/test_model.py

@@ -2346,7 +2346,7 @@ def test__db_get_value():
     @staticmethod
     def test__to_base_type():
         prop = model.UserProperty(name="u")
-        entity = prop._to_base_type(model.User("email", "auth_domain",))
+        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"

tests/unit/test_query.py

@@ -2213,7 +2213,7 @@ def next(self):
         _datastore_query.iterate.return_value = DummyQueryIterator()
         query = query_module.Query()
         query.filters = mock.Mock(
-            _multiquery=False, _post_filters=mock.Mock(return_value=False),
+            _multiquery=False, _post_filters=mock.Mock(return_value=False)
         )
         results, cursor, more = query.fetch_page(5)
         assert results == [0, 1, 2, 3, 4]