feature: add 'command' argument to private upload/download interface (#1082)

* Refactor client.download_blob_to_file

* Chore: clean up code

* refactor blob and client unit tests

* lint reformat

* Rename _prep_and_do_download

* Refactor blob.upload_from_file

* Lint reformat

* feature: add 'command' argument to private upload/download interface

* lint reformat

* reduce duplication and edit docstring

Author: MiaCY

google/cloud/storage/_helpers.py

@@ -599,19 +599,30 @@ def _get_default_headers(
     user_agent,
     content_type="application/json; charset=UTF-8",
     x_upload_content_type=None,
+    command=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.
+    :type user_agent: str
+    :param user_agent: The user-agent for requests.
+
+    :type command: str
+    :param user_agent:
+        (Optional) Information about which interface for upload/download was used, to be included in the X-Goog-API-Client header for traffic analysis purposes. Please leave as None unless otherwise directed.
+
+    :rtype: dict
+    :returns: The headers to be used for the request.
     """
+    x_goog_api_client = f"{user_agent} {_get_invocation_id()}"
+
+    if command:
+        x_goog_api_client += f" gccl-gcs-cmd/{command}"
+
     return {
         "Accept": "application/json",
         "Accept-Encoding": "gzip, deflate",
         "User-Agent": user_agent,
-        "X-Goog-API-Client": f"{user_agent} {_get_invocation_id()}",
+        "X-Goog-API-Client": x_goog_api_client,
         "content-type": content_type,
         "x-upload-content-type": x_upload_content_type or content_type,
     }

google/cloud/storage/blob.py

@@ -1697,7 +1697,7 @@ def _get_writable_metadata(self):
 
         return object_metadata
 
-    def _get_upload_arguments(self, client, content_type):
+    def _get_upload_arguments(self, client, content_type, command=None):
         """Get required arguments for performing an upload.
 
         The content type returned will be determined in order of precedence:
@@ -1709,6 +1709,10 @@ def _get_upload_arguments(self, client, content_type):
         :type content_type: str
         :param content_type: Type of content being uploaded (or :data:`None`).
 
+        :type command: str
+        :param command:
+            (Optional) Information about which interface for upload was used, to be included in the X-Goog-API-Client header for traffic analysis purposes. Please leave as None unless otherwise directed.
+
         :rtype: tuple
         :returns: A triple of
 
@@ -1718,7 +1722,9 @@ def _get_upload_arguments(self, client, content_type):
         """
         content_type = self._get_content_type(content_type)
         headers = {
-            **_get_default_headers(client._connection.user_agent, content_type),
+            **_get_default_headers(
+                client._connection.user_agent, content_type, command=command
+            ),
             **_get_encryption_headers(self._encryption_key),
         }
         object_metadata = self._get_writable_metadata()
@@ -1739,6 +1745,7 @@ def _do_multipart_upload(
         timeout=_DEFAULT_TIMEOUT,
         checksum=None,
         retry=None,
+        command=None,
     ):
         """Perform a multipart upload.
 
@@ -1822,6 +1829,10 @@ def _do_multipart_upload(
             (google.cloud.storage.retry) for information on retry types and how
             to configure them.
 
+        :type command: str
+        :param command:
+            (Optional) Information about which interface for upload was used, to be included in the X-Goog-API-Client header for traffic analysis purposes. Please leave as None unless otherwise directed.
+
         :rtype: :class:`~requests.Response`
         :returns: The "200 OK" response object returned after the multipart
                   upload request.
@@ -1840,7 +1851,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(client, content_type)
+        info = self._get_upload_arguments(client, content_type, command=command)
         headers, object_metadata, content_type = info
 
         hostname = _get_host_name(client._connection)
@@ -1910,6 +1921,7 @@ def _initiate_resumable_upload(
         timeout=_DEFAULT_TIMEOUT,
         checksum=None,
         retry=None,
+        command=None,
     ):
         """Initiate a resumable upload.
 
@@ -2008,6 +2020,10 @@ def _initiate_resumable_upload(
             (google.cloud.storage.retry) for information on retry types and how
             to configure them.
 
+        :type command: str
+        :param command:
+            (Optional) Information about which interface for upload was used, to be included in the X-Goog-API-Client header for traffic analysis purposes. Please leave as None unless otherwise directed.
+
         :rtype: tuple
         :returns:
             Pair of
@@ -2025,7 +2041,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(client, content_type)
+        info = self._get_upload_arguments(client, content_type, command=command)
         headers, object_metadata, content_type = info
         if extra_headers is not None:
             headers.update(extra_headers)
@@ -2103,6 +2119,7 @@ def _do_resumable_upload(
         timeout=_DEFAULT_TIMEOUT,
         checksum=None,
         retry=None,
+        command=None,
     ):
         """Perform a resumable upload.
 
@@ -2191,6 +2208,10 @@ def _do_resumable_upload(
             (google.cloud.storage.retry) for information on retry types and how
             to configure them.
 
+        :type command: str
+        :param command:
+            (Optional) Information about which interface for upload was used, to be included in the X-Goog-API-Client header for traffic analysis purposes. Please leave as None unless otherwise directed.
+
         :rtype: :class:`~requests.Response`
         :returns: The "200 OK" response object returned after the final chunk
                   is uploaded.
@@ -2209,6 +2230,7 @@ def _do_resumable_upload(
             timeout=timeout,
             checksum=checksum,
             retry=retry,
+            command=command,
         )
         while not upload.finished:
             try:
@@ -2234,6 +2256,7 @@ def _do_upload(
         timeout=_DEFAULT_TIMEOUT,
         checksum=None,
         retry=None,
+        command=None,
     ):
         """Determine an upload strategy and then perform the upload.
 
@@ -2333,6 +2356,10 @@ def _do_upload(
             configuration changes for Retry objects such as delays and deadlines
             are respected.
 
+        :type command: str
+        :param command:
+            (Optional) Information about which interface for upload was used, to be included in the X-Goog-API-Client header for traffic analysis purposes. Please leave as None unless otherwise directed.
+
         :rtype: dict
         :returns: The parsed JSON from the "200 OK" response. This will be the
                   **only** response in the multipart case and it will be the
@@ -2366,6 +2393,7 @@ def _do_upload(
                 timeout=timeout,
                 checksum=checksum,
                 retry=retry,
+                command=command,
             )
         else:
             response = self._do_resumable_upload(
@@ -2382,6 +2410,7 @@ def _do_upload(
                 timeout=timeout,
                 checksum=checksum,
                 retry=retry,
+                command=command,
             )
 
         return response.json()
@@ -2402,6 +2431,7 @@ def _prep_and_do_upload(
         timeout=_DEFAULT_TIMEOUT,
         checksum=None,
         retry=DEFAULT_RETRY_IF_GENERATION_SPECIFIED,
+        command=None,
     ):
         """Upload the contents of this blob from a file-like object.
 
@@ -2522,6 +2552,10 @@ def _prep_and_do_upload(
             configuration changes for Retry objects such as delays and deadlines
             are respected.
 
+        :type command: str
+        :param command:
+            (Optional) Information about which interface for upload was used, to be included in the X-Goog-API-Client header for traffic analysis purposes. Please leave as None unless otherwise directed.
+
         :raises: :class:`~google.cloud.exceptions.GoogleCloudError`
                  if the upload response returns an error status.
         """
@@ -2551,6 +2585,7 @@ def _prep_and_do_upload(
                 timeout=timeout,
                 checksum=checksum,
                 retry=retry,
+                command=command,
             )
             self._set_properties(created_json)
         except resumable_media.InvalidResponse as exc:
@@ -4108,6 +4143,7 @@ def _prep_and_do_download(
         timeout=_DEFAULT_TIMEOUT,
         checksum="md5",
         retry=DEFAULT_RETRY,
+        command=None,
     ):
         """Download the contents of a blob object into a file-like object.
 
@@ -4195,6 +4231,10 @@ def _prep_and_do_download(
             predicates in a Retry object. The default will always be used. Other
             configuration changes for Retry objects such as delays and deadlines
             are respected.
+
+        :type command: str
+        :param command:
+            (Optional) Information about which interface for download was used, to be included in the X-Goog-API-Client header for traffic analysis purposes. Please leave as None unless otherwise directed.
         """
         # Handle ConditionalRetryPolicy.
         if isinstance(retry, ConditionalRetryPolicy):
@@ -4224,7 +4264,10 @@ def _prep_and_do_download(
             if_etag_match=if_etag_match,
             if_etag_not_match=if_etag_not_match,
         )
-        headers = {**_get_default_headers(client._connection.user_agent), **headers}
+        headers = {
+            **_get_default_headers(client._connection.user_agent, command=command),
+            **headers,
+        }
 
         transport = client._http
 

tests/unit/test_blob.py

@@ -2251,11 +2251,12 @@ def test__get_upload_arguments(self):
         blob = self._make_one(name, bucket=None, encryption_key=key)
         blob.content_disposition = "inline"
 
+        COMMAND = "tm.upload_many"
         content_type = "image/jpeg"
         with patch.object(
             _helpers, "_get_invocation_id", return_value=GCCL_INVOCATION_TEST_CONST
         ):
-            info = blob._get_upload_arguments(client, content_type)
+            info = blob._get_upload_arguments(client, content_type, command=COMMAND)
 
         headers, object_metadata, new_content_type = info
         header_key_value = "W3BYd0AscEBAQWZCZnJSM3gtMmIyU0NIUiwuP1l3Uk8="
@@ -2264,11 +2265,17 @@ def test__get_upload_arguments(self):
             _helpers, "_get_invocation_id", return_value=GCCL_INVOCATION_TEST_CONST
         ):
             expected_headers = {
-                **_get_default_headers(client._connection.user_agent, content_type),
+                **_get_default_headers(
+                    client._connection.user_agent, content_type, command=COMMAND
+                ),
                 "X-Goog-Encryption-Algorithm": "AES256",
                 "X-Goog-Encryption-Key": header_key_value,
                 "X-Goog-Encryption-Key-Sha256": header_key_hash_value,
             }
+        self.assertEqual(
+            headers["X-Goog-API-Client"],
+            f"{client._connection.user_agent} {GCCL_INVOCATION_TEST_CONST} gccl-gcs-cmd/{COMMAND}",
+        )
         self.assertEqual(headers, expected_headers)
         expected_metadata = {
             "contentDisposition": blob.content_disposition,
@@ -3184,6 +3191,7 @@ def _do_upload_helper(
                 timeout=expected_timeout,
                 checksum=None,
                 retry=retry,
+                command=None,
             )
             blob._do_resumable_upload.assert_not_called()
         else:
@@ -3202,6 +3210,7 @@ def _do_upload_helper(
                 timeout=expected_timeout,
                 checksum=None,
                 retry=retry,
+                command=None,
             )
 
     def test__do_upload_uses_multipart(self):
@@ -3294,6 +3303,7 @@ def _upload_from_file_helper(self, side_effect=None, **kwargs):
             timeout=expected_timeout,
             checksum=None,
             retry=retry,
+            command=None,
         )
         return stream
 
@@ -3385,7 +3395,13 @@ def _do_upload_mock_call_helper(
         if not retry:
             retry = DEFAULT_RETRY_IF_GENERATION_SPECIFIED if not num_retries else None
         self.assertEqual(
-            kwargs, {"timeout": expected_timeout, "checksum": None, "retry": retry}
+            kwargs,
+            {
+                "timeout": expected_timeout,
+                "checksum": None,
+                "retry": retry,
+                "command": None,
+            },
         )
 
         return pos_args[1]