Adds optional location_prefix kwarg in to_legacy_urlsafe (googleapis#4635)

egalpin · dhermes · commit 4e50bba46ec5 · 2017-12-20T11:50:27.000-08:00
diff --git a/datastore/google/cloud/datastore/key.py b/datastore/google/cloud/datastore/key.py
@@ -298,7 +298,7 @@ def to_protobuf(self):
 
         return key
 
-    def to_legacy_urlsafe(self):
+    def to_legacy_urlsafe(self, location_prefix=None):
         """Convert to a base64 encode urlsafe string for App Engine.
 
         This is intended to work with the "legacy" representation of a
@@ -310,13 +310,26 @@ def to_legacy_urlsafe(self):
         .. note::
 
             The string returned by ``to_legacy_urlsafe`` is equivalent, but
-            not identical, to the string returned by ``ndb``.
+            not identical, to the string returned by ``ndb``. The location
+            prefix may need to be specified to obtain identical urlsafe
+            keys.
+
+        :type location_prefix: str
+        :param location_prefix: The location prefix of an App Engine project
+                                ID. Often this value is 's~', but may also be
+                                'e~', or other location prefixes currently
+                                unknown.
 
         :rtype: bytes
         :returns: A bytestring containing the key encoded as URL-safe base64.
         """
+        if location_prefix is None:
+            project_id = self.project
+        else:
+            project_id = location_prefix + self.project
+
         reference = _app_engine_key_pb2.Reference(
-            app=self.project,
+            app=project_id,
             path=_to_legacy_path(self._path),  # Avoid the copy.
             name_space=self.namespace,
         )
diff --git a/datastore/tests/unit/test_key.py b/datastore/tests/unit/test_key.py
@@ -35,6 +35,9 @@ class TestKey(unittest.TestCase):
     _URLSAFE_EXAMPLE2 = b'agZzfmZpcmVyDwsSBEtpbmQiBVRoaW5nDA'
     _URLSAFE_APP2 = 's~fire'
     _URLSAFE_FLAT_PATH2 = ('Kind', 'Thing')
+    _URLSAFE_EXAMPLE3 = b'ahhzfnNhbXBsZS1hcHAtbm8tbG9jYXRpb25yCgsSBFpvcnAYWAw'
+    _URLSAFE_APP3 = 'sample-app-no-location'
+    _URLSAFE_FLAT_PATH3 = ('Zorp', 88)
 
     @staticmethod
     def _get_target_class():
@@ -408,6 +411,13 @@ def test_to_legacy_urlsafe_strip_padding(self):
         # Make sure it started with base64 padding.
         self.assertNotEqual(len(self._URLSAFE_EXAMPLE2) % 4, 0)
 
+    def test_to_legacy_urlsafe_with_location_prefix(self):
+        key = self._make_one(
+            *self._URLSAFE_FLAT_PATH3,
+            project=self._URLSAFE_APP3)
+        urlsafe = key.to_legacy_urlsafe(location_prefix='s~')
+        self.assertEqual(urlsafe, self._URLSAFE_EXAMPLE3)
+
     def test_from_legacy_urlsafe(self):
         klass = self._get_target_class()
         key = klass.from_legacy_urlsafe(self._URLSAFE_EXAMPLE1)
@@ -430,6 +440,15 @@ def test_from_legacy_urlsafe_needs_padding(self):
         self.assertIsNone(key.namespace)
         self.assertEqual(key.flat_path, self._URLSAFE_FLAT_PATH2)
 
+    def test_from_legacy_urlsafe_with_location_prefix(self):
+        klass = self._get_target_class()
+        # Make sure it will have base64 padding added.
+        key = klass.from_legacy_urlsafe(self._URLSAFE_EXAMPLE3)
+
+        self.assertEqual(key.project, self._URLSAFE_APP3)
+        self.assertIsNone(key.namespace)
+        self.assertEqual(key.flat_path, self._URLSAFE_FLAT_PATH3)
+
     def test_is_partial_no_name_or_id(self):
         key = self._make_one('KIND', project=self._DEFAULT_PROJECT)
         self.assertTrue(key.is_partial)
