diff --git a/datastore/google/cloud/datastore/key.py b/datastore/google/cloud/datastore/key.py index 166a5afde46b..f1733f8f5d8e 100644 --- a/datastore/google/cloud/datastore/key.py +++ b/datastore/google/cloud/datastore/key.py @@ -304,7 +304,8 @@ def to_legacy_urlsafe(self): This is intended to work with the "legacy" representation of a datastore "Key" used within Google App Engine (a so-called "Reference"). The returned string can be used as the ``urlsafe`` - argument to ``ndb.Key(urlsafe=...)``. + argument to ``ndb.Key(urlsafe=...)``. The base64 encoded values + will have padding removed. :rtype: bytes :returns: A bytestring containing the key encoded as URL-safe base64. @@ -315,7 +316,7 @@ def to_legacy_urlsafe(self): name_space=self.namespace, ) raw_bytes = reference.SerializeToString() - return base64.urlsafe_b64encode(raw_bytes) + return base64.urlsafe_b64encode(raw_bytes).strip(b'=') @classmethod def from_legacy_urlsafe(cls, urlsafe): @@ -334,6 +335,8 @@ def from_legacy_urlsafe(cls, urlsafe): :returns: The key corresponding to ``urlsafe``. """ urlsafe = _to_bytes(urlsafe, encoding='ascii') + padding = b'=' * (-len(urlsafe) % 4) + urlsafe += padding raw_bytes = base64.urlsafe_b64decode(urlsafe) reference = _app_engine_key_pb2.Reference() diff --git a/datastore/tests/unit/test_key.py b/datastore/tests/unit/test_key.py index 5b89e146254d..4fb7b89911b2 100644 --- a/datastore/tests/unit/test_key.py +++ b/datastore/tests/unit/test_key.py @@ -26,12 +26,15 @@ class TestKey(unittest.TestCase): # 'Parent', 59, 'Child', 'Feather', # namespace='space', app='s~sample-app') # urlsafe = key.urlsafe() - _URLSAFE_EXAMPLE = ( + _URLSAFE_EXAMPLE1 = ( b'agxzfnNhbXBsZS1hcHByHgsSBlBhcmVudBg7DAsSBUNoaWxkIgdGZ' b'WF0aGVyDKIBBXNwYWNl') - _URLSAFE_APP = 's~sample-app' - _URLSAFE_NAMESPACE = 'space' - _URLSAFE_FLAT_PATH = ('Parent', 59, 'Child', 'Feather') + _URLSAFE_APP1 = 's~sample-app' + _URLSAFE_NAMESPACE1 = 'space' + _URLSAFE_FLAT_PATH1 = ('Parent', 59, 'Child', 'Feather') + _URLSAFE_EXAMPLE2 = b'agZzfmZpcmVyDwsSBEtpbmQiBVRoaW5nDA' + _URLSAFE_APP2 = 's~fire' + _URLSAFE_FLAT_PATH2 = ('Kind', 'Thing') @staticmethod def _get_target_class(): @@ -388,25 +391,45 @@ def test_to_protobuf_w_no_kind(self): def test_to_legacy_urlsafe(self): key = self._make_one( - *self._URLSAFE_FLAT_PATH, - project=self._URLSAFE_APP, - namespace=self._URLSAFE_NAMESPACE) + *self._URLSAFE_FLAT_PATH1, + project=self._URLSAFE_APP1, + namespace=self._URLSAFE_NAMESPACE1) # NOTE: ``key.project`` is somewhat "invalid" but that is OK. urlsafe = key.to_legacy_urlsafe() - self.assertEqual(urlsafe, self._URLSAFE_EXAMPLE) + self.assertEqual(urlsafe, self._URLSAFE_EXAMPLE1) + + def test_to_legacy_urlsafe_strip_padding(self): + key = self._make_one( + *self._URLSAFE_FLAT_PATH2, + project=self._URLSAFE_APP2) + # NOTE: ``key.project`` is somewhat "invalid" but that is OK. + urlsafe = key.to_legacy_urlsafe() + self.assertEqual(urlsafe, self._URLSAFE_EXAMPLE2) + # Make sure it started with base64 padding. + self.assertNotEqual(len(self._URLSAFE_EXAMPLE2) % 4, 0) def test_from_legacy_urlsafe(self): klass = self._get_target_class() - key = klass.from_legacy_urlsafe(self._URLSAFE_EXAMPLE) + key = klass.from_legacy_urlsafe(self._URLSAFE_EXAMPLE1) - self.assertEqual('s~' + key.project, self._URLSAFE_APP) - self.assertEqual(key.namespace, self._URLSAFE_NAMESPACE) - self.assertEqual(key.flat_path, self._URLSAFE_FLAT_PATH) + self.assertEqual('s~' + key.project, self._URLSAFE_APP1) + self.assertEqual(key.namespace, self._URLSAFE_NAMESPACE1) + self.assertEqual(key.flat_path, self._URLSAFE_FLAT_PATH1) # Also make sure we didn't accidentally set the parent. self.assertIsNone(key._parent) self.assertIsNotNone(key.parent) self.assertIs(key._parent, key.parent) + def test_from_legacy_urlsafe_needs_padding(self): + klass = self._get_target_class() + # Make sure it will have base64 padding added. + self.assertNotEqual(len(self._URLSAFE_EXAMPLE2) % 4, 0) + key = klass.from_legacy_urlsafe(self._URLSAFE_EXAMPLE2) + + self.assertEqual('s~' + key.project, self._URLSAFE_APP2) + self.assertIsNone(key.namespace) + self.assertEqual(key.flat_path, self._URLSAFE_FLAT_PATH2) + def test_is_partial_no_name_or_id(self): key = self._make_one('KIND', project=self._DEFAULT_PROJECT) self.assertTrue(key.is_partial)