Strip base64 padding characters from urlsafe in Datastore's (to|from)_legacy_urlsafe (googleapis#3560)

Michael Englo · dhermes · commit eb033665252c · 2017-06-28T10:49:14.000-07:00
Also

* add padding characters in `from_legacy_urlsafe` if needed
* add an extra example in the unit tests that actually requires base64 padding
diff --git 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
@@ -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)
