diff --git a/datastore/google/cloud/datastore/key.py b/datastore/google/cloud/datastore/key.py index 74d23e49265c..2824e57d3ebf 100644 --- 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 index e95d756013cc..ddd04d161967 100644 --- 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)