From e1290f523808c7ef5be7dd335a5c94cd1739e6e3 Mon Sep 17 00:00:00 2001 From: William Brode Date: Thu, 23 Dec 2021 11:30:46 -0800 Subject: [PATCH] feat: add support for signed resumable upload URLs (#290) Co-authored-by: Aaron Gabriel Neyer Co-authored-by: Cathy Ouyang --- google/resumable_media/_upload.py | 16 ++++++++++++---- tests/unit/test__upload.py | 21 +++++++++++++++++++-- 2 files changed, 31 insertions(+), 6 deletions(-) diff --git a/google/resumable_media/_upload.py b/google/resumable_media/_upload.py index 9cf40202..dc22cd1c 100644 --- a/google/resumable_media/_upload.py +++ b/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 diff --git a/tests/unit/test__upload.py b/tests/unit/test__upload.py index 2c2ca924..4e948f0d 100644 --- a/tests/unit/test__upload.py +++ b/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(