diff --git a/google/cloud/storage/_helpers.py b/google/cloud/storage/_helpers.py index 29968a9aa..3ff52dd75 100644 --- a/google/cloud/storage/_helpers.py +++ b/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, } diff --git a/google/cloud/storage/blob.py b/google/cloud/storage/blob.py index 6d6bd7986..32b8a62db 100644 --- a/google/cloud/storage/blob.py +++ b/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 diff --git a/tests/unit/test_blob.py b/tests/unit/test_blob.py index 658dba81e..a8d024176 100644 --- a/tests/unit/test_blob.py +++ b/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]