feat: add support for signed resumable upload URLs (#290)

Co-authored-by: Aaron Gabriel Neyer <[email protected]>
Co-authored-by: Cathy Ouyang <[email protected]>

Author: 3 people

google/resumable_media/_upload.py

@@ -27,6 +27,7 @@
 import random
 import re
 import sys
+import urllib.parse
 
 from google import resumable_media
 from google.resumable_media import _helpers
@@ -462,10 +463,17 @@ def _prepare_initiate_request(
 
         self._stream = stream
         self._content_type = content_type
-        headers = {
-            _CONTENT_TYPE_HEADER: "application/json; charset=UTF-8",
-            "x-upload-content-type": content_type,
-        }
+
+        # Signed URL requires content type set directly - not through x-upload-content-type
+        parse_result = urllib.parse.urlparse(self.upload_url)
+        parsed_query = urllib.parse.parse_qs(parse_result.query)
+        if "x-goog-signature" in parsed_query or "X-Goog-Signature" in parsed_query:
+            headers = {_CONTENT_TYPE_HEADER: content_type}
+        else:
+            headers = {
+                _CONTENT_TYPE_HEADER: "application/json; charset=UTF-8",
+                "x-upload-content-type": content_type,
+            }
         # Set the total bytes if possible.
         if total_bytes is not None:
             self._total_bytes = total_bytes

tests/unit/test__upload.py

@@ -393,12 +393,14 @@ def test_total_bytes_property(self):
         upload._total_bytes = 8192
         assert upload.total_bytes == 8192
 
-    def _prepare_initiate_request_helper(self, upload_headers=None, **method_kwargs):
+    def _prepare_initiate_request_helper(
+        self, upload_url=RESUMABLE_URL, upload_headers=None, **method_kwargs
+    ):
         data = b"some really big big data."
         stream = io.BytesIO(data)
         metadata = {"name": "big-data-file.txt"}
 
-        upload = _upload.ResumableUpload(RESUMABLE_URL, ONE_MB, headers=upload_headers)
+        upload = _upload.ResumableUpload(upload_url, ONE_MB, headers=upload_headers)
         orig_headers = upload._headers.copy()
         # Check ``upload``-s state before.
         assert upload._stream is None
@@ -435,6 +437,21 @@ def test__prepare_initiate_request(self):
         }
         assert headers == expected_headers
 
+    def test_prepare_initiate_request_with_signed_url(self):
+        signed_urls = [
+            "https://storage.googleapis.com/b/o?x-goog-signature=123abc",
+            "https://storage.googleapis.com/b/o?X-Goog-Signature=123abc",
+        ]
+        for signed_url in signed_urls:
+            data, headers = self._prepare_initiate_request_helper(
+                upload_url=signed_url,
+            )
+            expected_headers = {
+                "content-type": BASIC_CONTENT,
+                "x-upload-content-length": "{:d}".format(len(data)),
+            }
+            assert headers == expected_headers
+
     def test__prepare_initiate_request_with_headers(self):
         headers = {"caviar": "beluga", "top": "quark"}
         data, new_headers = self._prepare_initiate_request_helper(