fix: add user agent in python-storage when calling resumable media (WIP) (#715)

* fix: add user agent in python-storage when calling resumable media (WIP) * put the things in the right places * starting to get some tests working * a bit closer, still some things not quite working * almost there * first cleanup * ensure up to date resumable media * next round of cleanup * lint * update resumable media * get tests passing * lint

Author: Aaron Gabriel Neyer

google/cloud/storage/_helpers.py

@@ -581,3 +581,25 @@ def _api_core_retry_to_resumable_media_retry(retry, num_retries=None):
  return resumable_media.RetryStrategy(max_retries=num_retries)
  else:
  return resumable_media.RetryStrategy(max_retries=0)
+
+
+def _get_default_headers(
+ user_agent,
+ content_type="application/json; charset=UTF-8",
+ x_upload_content_type=None,
+):
+ """Get the headers for a request.
+
+ Args:
+ user_agent (str): The user-agent for requests.
+ Returns:
+ Dict: The headers to be used for the request.
+ """
+ return {
+ "Accept": "application/json",
+ "Accept-Encoding": "gzip, deflate",
+ "User-Agent": user_agent,
+ "x-goog-api-client": user_agent,
+ "content-type": content_type,
+ "x-upload-content-type": x_upload_content_type or content_type,
+ }

google/cloud/storage/blob.py

@@ -64,6 +64,7 @@
 from google.cloud.storage._helpers import _bucket_bound_hostname_url
 from google.cloud.storage._helpers import _raise_if_more_than_one_set
 from google.cloud.storage._helpers import _api_core_retry_to_resumable_media_retry
+from google.cloud.storage._helpers import _get_default_headers
 from google.cloud.storage._signing import generate_signed_url_v2
 from google.cloud.storage._signing import generate_signed_url_v4
 from google.cloud.storage._helpers import _NUM_RETRIES_MESSAGE
@@ -1720,7 +1721,7 @@ def _get_writable_metadata(self):
 
  return object_metadata
 
- def _get_upload_arguments(self, content_type):
+ def _get_upload_arguments(self, client, content_type):
  """Get required arguments for performing an upload.
 
  The content type returned will be determined in order of precedence:
@@ -1739,9 +1740,12 @@ def _get_upload_arguments(self, content_type):
  * An object metadata dictionary
  * The ``content_type`` as a string (according to precedence)
  """
- headers = _get_encryption_headers(self._encryption_key)
- object_metadata = self._get_writable_metadata()
  content_type = self._get_content_type(content_type)
+ headers = {
+ **_get_default_headers(client._connection.user_agent, content_type),
+ **_get_encryption_headers(self._encryption_key),
+ }
+ object_metadata = self._get_writable_metadata()
  return headers, object_metadata, content_type
 
  def _do_multipart_upload(
@@ -1860,7 +1864,7 @@ def _do_multipart_upload(
  transport = self._get_transport(client)
  if "metadata" in self._properties and "metadata" not in self._changes:
  self._changes.add("metadata")
- info = self._get_upload_arguments(content_type)
+ info = self._get_upload_arguments(client, content_type)
  headers, object_metadata, content_type = info
 
  hostname = _get_host_name(client._connection)
@@ -2045,7 +2049,7 @@ def _initiate_resumable_upload(
  transport = self._get_transport(client)
  if "metadata" in self._properties and "metadata" not in self._changes:
  self._changes.add("metadata")
- info = self._get_upload_arguments(content_type)
+ info = self._get_upload_arguments(client, content_type)
  headers, object_metadata, content_type = info
  if extra_headers is not None:
  headers.update(extra_headers)
@@ -2230,15 +2234,13 @@ def _do_resumable_upload(
  checksum=checksum,
  retry=retry,
  )
-
  while not upload.finished:
  try:
  response = upload.transmit_next_chunk(transport, timeout=timeout)
  except resumable_media.DataCorruption:
  # Attempt to delete the corrupted object.
  self.delete()
  raise
-
  return response
 
  def _do_upload(

google/cloud/storage/client.py

@@ -31,6 +31,7 @@
 from google.cloud._helpers import _LocalStack, _NOW
 from google.cloud.client import ClientWithProject
 from google.cloud.exceptions import NotFound
+from google.cloud.storage._helpers import _get_default_headers
 from google.cloud.storage._helpers import _get_environ_project
 from google.cloud.storage._helpers import _get_storage_host
 from google.cloud.storage._helpers import _BASE_STORAGE_URI
@@ -1131,6 +1132,7 @@ def download_blob_to_file(
  _add_etag_match_headers(
  headers, if_etag_match=if_etag_match, if_etag_not_match=if_etag_not_match,
  )
+ headers = {**_get_default_headers(self._connection.user_agent), **headers}
 
  transport = self._http
  try:

setup.py

@@ -31,7 +31,7 @@
  "google-auth >= 1.25.0, < 3.0dev",
  "google-api-core >= 1.31.5, <3.0.0dev,!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.0",
  "google-cloud-core >= 1.6.0, < 3.0dev",
- "google-resumable-media >= 1.3.0",
+ "google-resumable-media >= 2.3.2",
  "requests >= 2.18.0, < 3.0.0dev",
  "protobuf",
 ]

tests/unit/test_blob.py

@@ -26,6 +26,7 @@
 import mock
 import pytest
 
+from google.cloud.storage._helpers import _get_default_headers
 from google.cloud.storage.retry import (
  DEFAULT_RETRY,
  DEFAULT_RETRY_IF_METAGENERATION_SPECIFIED,
@@ -2212,16 +2213,19 @@ def test__set_metadata_to_none(self):
  def test__get_upload_arguments(self):
  name = u"blob-name"
  key = b"[pXw@,p@@AfBfrR3x-2b2SCHR,.?YwRO"
+ client = mock.Mock(_connection=_Connection)
+ client._connection.user_agent = "testing 1.2.3"
  blob = self._make_one(name, bucket=None, encryption_key=key)
  blob.content_disposition = "inline"
 
  content_type = u"image/jpeg"
- info = blob._get_upload_arguments(content_type)
+ info = blob._get_upload_arguments(client, content_type)
 
  headers, object_metadata, new_content_type = info
  header_key_value = "W3BYd0AscEBAQWZCZnJSM3gtMmIyU0NIUiwuP1l3Uk8="
  header_key_hash_value = "G0++dxF4q5rG4o9kE8gvEKn15RH6wLm0wXV1MgAlXOg="
  expected_headers = {
+ **_get_default_headers(client._connection.user_agent, content_type),
  "X-Goog-Encryption-Algorithm": "AES256",
  "X-Goog-Encryption-Key": header_key_value,
  "X-Goog-Encryption-Key-Sha256": header_key_hash_value,
@@ -2368,7 +2372,11 @@ def _do_multipart_success(
  + data_read
  + b"\r\n--==0==--"
  )
- headers = {"content-type": b'multipart/related; boundary="==0=="'}
+ headers = _get_default_headers(
+ client._connection.user_agent,
+ b'multipart/related; boundary="==0=="',
+ "application/xml",
+ )
  client._http.request.assert_called_once_with(
  "POST", upload_url, data=payload, headers=headers, timeout=expected_timeout
  )
@@ -2614,10 +2622,17 @@ def _initiate_resumable_helper(
 
  self.assertEqual(upload.upload_url, upload_url)
  if extra_headers is None:
- self.assertEqual(upload._headers, {})
+ self.assertEqual(
+ upload._headers,
+ _get_default_headers(client._connection.user_agent, content_type),
+ )
  else:
- self.assertEqual(upload._headers, extra_headers)
- self.assertIsNot(upload._headers, extra_headers)
+ expected_headers = {
+ **_get_default_headers(client._connection.user_agent, content_type),
+ **extra_headers,
+ }
+ self.assertEqual(upload._headers, expected_headers)
+ self.assertIsNot(upload._headers, expected_headers)
  self.assertFalse(upload.finished)
  if chunk_size is None:
  if blob_chunk_size is None:
@@ -2656,10 +2671,9 @@ def _initiate_resumable_helper(
  # Check the mocks.
  blob._get_writable_metadata.assert_called_once_with()
  payload = json.dumps(object_metadata).encode("utf-8")
- expected_headers = {
- "content-type": "application/json; charset=UTF-8",
- "x-upload-content-type": content_type,
- }
+ expected_headers = _get_default_headers(
+ client._connection.user_agent, x_upload_content_type=content_type
+ )
  if size is not None:
  expected_headers["x-upload-content-length"] = str(size)
  if extra_headers is not None:
@@ -2778,6 +2792,7 @@ def _make_resumable_transport(
 
  @staticmethod
  def _do_resumable_upload_call0(
+ client,
  blob,
  content_type,
  size=None,
@@ -2796,10 +2811,9 @@ def _do_resumable_upload_call0(
  )
  if predefined_acl is not None:
  upload_url += "&predefinedAcl={}".format(predefined_acl)
- expected_headers = {
- "content-type": "application/json; charset=UTF-8",
- "x-upload-content-type": content_type,
- }
+ expected_headers = _get_default_headers(
+ client._connection.user_agent, x_upload_content_type=content_type
+ )
  if size is not None:
  expected_headers["x-upload-content-length"] = str(size)
  payload = json.dumps({"name": blob.name}).encode("utf-8")
@@ -2809,6 +2823,7 @@ def _do_resumable_upload_call0(
 
  @staticmethod
  def _do_resumable_upload_call1(
+ client,
  blob,
  content_type,
  data,
@@ -2828,6 +2843,9 @@ def _do_resumable_upload_call1(
  content_range = "bytes 0-{:d}/{:d}".format(blob.chunk_size - 1, size)
 
  expected_headers = {
+ **_get_default_headers(
+ client._connection.user_agent, x_upload_content_type=content_type
+ ),
  "content-type": content_type,
  "content-range": content_range,
  }
@@ -2842,6 +2860,7 @@ def _do_resumable_upload_call1(
 
  @staticmethod
  def _do_resumable_upload_call2(
+ client,
  blob,
  content_type,
  data,
@@ -2859,6 +2878,9 @@ def _do_resumable_upload_call2(
  blob.chunk_size, total_bytes - 1, total_bytes
  )
  expected_headers = {
+ **_get_default_headers(
+ client._connection.user_agent, x_upload_content_type=content_type
+ ),
  "content-type": content_type,
  "content-range": content_range,
  }
@@ -2884,13 +2906,11 @@ def _do_resumable_helper(
  data_corruption=False,
  retry=None,
  ):
- bucket = _Bucket(name="yesterday")
- blob = self._make_one(u"blob-name", bucket=bucket)
- blob.chunk_size = blob._CHUNK_SIZE_MULTIPLE
- self.assertIsNotNone(blob.chunk_size)
-
+ CHUNK_SIZE = 256 * 1024
+ USER_AGENT = "testing 1.2.3"
+ content_type = u"text/html"
  # Data to be uploaded.
- data = b"<html>" + (b"A" * blob.chunk_size) + b"</html>"
+ data = b"<html>" + (b"A" * CHUNK_SIZE) + b"</html>"
  total_bytes = len(data)
  if use_size:
  size = total_bytes
@@ -2899,17 +2919,29 @@ def _do_resumable_helper(
 
  # Create mocks to be checked for doing transport.
  resumable_url = "http://test.invalid?upload_id=and-then-there-was-1"
- headers1 = {"location": resumable_url}
- headers2 = {"range": "bytes=0-{:d}".format(blob.chunk_size - 1)}
+ headers1 = {
+ **_get_default_headers(USER_AGENT, content_type),
+ "location": resumable_url,
+ }
+ headers2 = {
+ **_get_default_headers(USER_AGENT, content_type),
+ "range": "bytes=0-{:d}".format(CHUNK_SIZE - 1),
+ }
+ headers3 = _get_default_headers(USER_AGENT, content_type)
  transport, responses = self._make_resumable_transport(
- headers1, headers2, {}, total_bytes, data_corruption=data_corruption
+ headers1, headers2, headers3, total_bytes, data_corruption=data_corruption
  )
 
  # Create some mock arguments and call the method under test.
  client = mock.Mock(_http=transport, _connection=_Connection, spec=["_http"])
  client._connection.API_BASE_URL = "https://storage.googleapis.com"
+ client._connection.user_agent = USER_AGENT
  stream = io.BytesIO(data)
- content_type = u"text/html"
+
+ bucket = _Bucket(name="yesterday")
+ blob = self._make_one(u"blob-name", bucket=bucket)
+ blob.chunk_size = blob._CHUNK_SIZE_MULTIPLE
+ self.assertIsNotNone(blob.chunk_size)
 
  if timeout is None:
  expected_timeout = self._get_default_timeout()
@@ -2939,6 +2971,7 @@ def _do_resumable_helper(
 
  # Check the mocks.
  call0 = self._do_resumable_upload_call0(
+ client,
  blob,
  content_type,
  size=size,
@@ -2950,6 +2983,7 @@ def _do_resumable_helper(
  timeout=expected_timeout,
  )
  call1 = self._do_resumable_upload_call1(
+ client,
  blob,
  content_type,
  data,
@@ -2963,6 +2997,7 @@ def _do_resumable_helper(
  timeout=expected_timeout,
  )
  call2 = self._do_resumable_upload_call2(
+ client,
  blob,
  content_type,
  data,
@@ -3510,6 +3545,7 @@ def _create_resumable_upload_session_helper(
  size = 10000
  client = mock.Mock(_http=transport, _connection=_Connection, spec=[u"_http"])
  client._connection.API_BASE_URL = "https://storage.googleapis.com"
+ client._connection.user_agent = "testing 1.2.3"
 
  if timeout is None:
  expected_timeout = self._get_default_timeout()
@@ -3556,7 +3592,9 @@ def _create_resumable_upload_session_helper(
  upload_url += "?" + urlencode(qs_params)
  payload = b'{"name": "blob-name"}'
  expected_headers = {
- "content-type": "application/json; charset=UTF-8",
+ **_get_default_headers(
+ client._connection.user_agent, x_upload_content_type=content_type
+ ),
  "x-upload-content-length": str(size),
  "x-upload-content-type": content_type,
  }
@@ -5739,6 +5777,7 @@ class _Connection(object):
 
  API_BASE_URL = "http://example.com"
  USER_AGENT = "testing 1.2.3"
+ user_agent = "testing 1.2.3"
  credentials = object()
 
 

tests/unit/test_client.py

@@ -28,6 +28,7 @@
 from google.oauth2.service_account import Credentials
 
 from google.cloud.storage._helpers import STORAGE_EMULATOR_ENV_VAR
+from google.cloud.storage._helpers import _get_default_headers
 from google.cloud.storage.retry import DEFAULT_RETRY
 from google.cloud.storage.retry import DEFAULT_RETRY_IF_GENERATION_SPECIFIED
 
@@ -1567,7 +1568,10 @@ def test_download_blob_to_file_with_failure(self):
 
  self.assertEqual(file_obj.tell(), 0)
 
- headers = {"accept-encoding": "gzip"}
+ headers = {
+ **_get_default_headers(client._connection.user_agent),
+ "accept-encoding": "gzip",
+ }
  blob._do_download.assert_called_once_with(
  client._http,
  file_obj,
@@ -1598,7 +1602,10 @@ def test_download_blob_to_file_with_uri(self):
  ):
  client.download_blob_to_file("gs://bucket_name/path/to/object", file_obj)
 
- headers = {"accept-encoding": "gzip"}
+ headers = {
+ **_get_default_headers(client._connection.user_agent),
+ "accept-encoding": "gzip",
+ }
  blob._do_download.assert_called_once_with(
  client._http,
  file_obj,
@@ -1714,6 +1721,7 @@ def _download_blob_to_file_helper(
  if_etag_not_match = [if_etag_not_match]
  headers["If-None-Match"] = ", ".join(if_etag_not_match)
 
+ headers = {**_get_default_headers(client._connection.user_agent), **headers}
  blob._do_download.assert_called_once_with(
  client._http,
  file_obj,