feat: add support for custom headers (#1121)

cojenco · MiaCY · web-flow · commit 2f92c3a2a3a1 · 2023-10-19T11:52:11.000-07:00
* Chore: refactor client.download_blob_to_file (#1052) * Refactor client.download_blob_to_file * Chore: clean up code * refactor blob and client unit tests * lint reformat * Rename _prep_and_do_download * Chore: refactor blob.upload_from_file (#1063) * 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 (#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 * feat: add support for custom headers starting with metadata op * add custom headers to downloads in client blob modules * add custom headers to uploads with tests * update mocks and tests * test custom headers support tm mpu uploads * update tm test * update test --------- Co-authored-by: MiaCY <97990237+MiaCY@users.noreply.github.com>
diff --git a/google/cloud/storage/blob.py b/google/cloud/storage/blob.py
@@ -1738,11 +1738,13 @@ def _get_upload_arguments(self, client, content_type, filename=None, command=Non
                   * The ``content_type`` as a string (according to precedence)
         """
         content_type = self._get_content_type(content_type, filename=filename)
+        # Add any client attached custom headers to the upload headers.
         headers = {
             **_get_default_headers(
                 client._connection.user_agent, content_type, command=command
             ),
             **_get_encryption_headers(self._encryption_key),
+            **client._extra_headers,
         }
         object_metadata = self._get_writable_metadata()
         return headers, object_metadata, content_type
@@ -4313,9 +4315,11 @@ def _prep_and_do_download(
             if_etag_match=if_etag_match,
             if_etag_not_match=if_etag_not_match,
         )
+        # Add any client attached custom headers to be sent with the request.
         headers = {
             **_get_default_headers(client._connection.user_agent, command=command),
             **headers,
+            **client._extra_headers,
         }
 
         transport = client._http
diff --git a/google/cloud/storage/client.py b/google/cloud/storage/client.py
@@ -94,6 +94,11 @@ class Client(ClientWithProject):
         (Optional) Whether authentication is required under custom endpoints.
         If false, uses AnonymousCredentials and bypasses authentication.
         Defaults to True. Note this is only used when a custom endpoint is set in conjunction.
+
+    :type extra_headers: dict
+    :param extra_headers:
+        (Optional) Custom headers to be sent with the requests attached to the client.
+        For example, you can add custom audit logging headers.
     """
 
     SCOPE = (
@@ -111,6 +116,7 @@ def __init__(
         client_info=None,
         client_options=None,
         use_auth_w_custom_endpoint=True,
+        extra_headers={},
     ):
         self._base_connection = None
 
@@ -127,6 +133,7 @@ def __init__(
         # are passed along, for use in __reduce__ defined elsewhere.
         self._initial_client_info = client_info
         self._initial_client_options = client_options
+        self._extra_headers = extra_headers
 
         kw_args = {"client_info": client_info}
 
@@ -172,7 +179,10 @@ def __init__(
         if no_project:
             self.project = None
 
-        self._connection = Connection(self, **kw_args)
+        # Pass extra_headers to Connection
+        connection = Connection(self, **kw_args)
+        connection.extra_headers = extra_headers
+        self._connection = connection
         self._batch_stack = _LocalStack()
 
     @classmethod
diff --git a/google/cloud/storage/transfer_manager.py b/google/cloud/storage/transfer_manager.py
@@ -1289,6 +1289,7 @@ def _reduce_client(cl):
     _http = None  # Can't carry this over
     client_info = cl._initial_client_info
     client_options = cl._initial_client_options
+    extra_headers = cl._extra_headers
 
     return _LazyClient, (
         client_object_id,
@@ -1297,6 +1298,7 @@ def _reduce_client(cl):
         _http,
         client_info,
         client_options,
+        extra_headers,
     )
 
 
diff --git a/tests/unit/test__http.py b/tests/unit/test__http.py
@@ -71,6 +71,55 @@ def test_extra_headers(self):
             timeout=_DEFAULT_TIMEOUT,
         )
 
+    def test_metadata_op_has_client_custom_headers(self):
+        import requests
+        import google.auth.credentials
+        from google.cloud import _http as base_http
+        from google.cloud.storage import Client
+        from google.cloud.storage.constants import _DEFAULT_TIMEOUT
+
+        custom_headers = {
+            "x-goog-custom-audit-foo": "bar",
+            "x-goog-custom-audit-user": "baz",
+        }
+        http = mock.create_autospec(requests.Session, instance=True)
+        response = requests.Response()
+        response.status_code = 200
+        data = b"brent-spiner"
+        response._content = data
+        http.is_mtls = False
+        http.request.return_value = response
+        credentials = mock.Mock(spec=google.auth.credentials.Credentials)
+        client = Client(
+            project="project",
+            credentials=credentials,
+            _http=http,
+            extra_headers=custom_headers,
+        )
+        req_data = "hey-yoooouuuuu-guuuuuyyssss"
+        with patch.object(
+            _helpers, "_get_invocation_id", return_value=GCCL_INVOCATION_TEST_CONST
+        ):
+            result = client._connection.api_request(
+                "GET", "/rainbow", data=req_data, expect_json=False
+            )
+        self.assertEqual(result, data)
+
+        expected_headers = {
+            **custom_headers,
+            "Accept-Encoding": "gzip",
+            base_http.CLIENT_INFO_HEADER: f"{client._connection.user_agent} {GCCL_INVOCATION_TEST_CONST}",
+            "User-Agent": client._connection.user_agent,
+        }
+        expected_uri = client._connection.build_api_url("/rainbow")
+        http.request.assert_called_once_with(
+            data=req_data,
+            headers=expected_headers,
+            method="GET",
+            url=expected_uri,
+            timeout=_DEFAULT_TIMEOUT,
+        )
+
     def test_build_api_url_no_extra_query_params(self):
         from urllib.parse import parse_qsl
         from urllib.parse import urlsplit
diff --git a/tests/unit/test_blob.py b/tests/unit/test_blob.py
@@ -2246,8 +2246,13 @@ def test__set_metadata_to_none(self):
     def test__get_upload_arguments(self):
         name = "blob-name"
         key = b"[pXw@,p@@AfBfrR3x-2b2SCHR,.?YwRO"
+        custom_headers = {
+            "x-goog-custom-audit-foo": "bar",
+            "x-goog-custom-audit-user": "baz",
+        }
         client = mock.Mock(_connection=_Connection)
         client._connection.user_agent = "testing 1.2.3"
+        client._extra_headers = custom_headers
         blob = self._make_one(name, bucket=None, encryption_key=key)
         blob.content_disposition = "inline"
 
@@ -2271,6 +2276,7 @@ def test__get_upload_arguments(self):
                 "X-Goog-Encryption-Algorithm": "AES256",
                 "X-Goog-Encryption-Key": header_key_value,
                 "X-Goog-Encryption-Key-Sha256": header_key_hash_value,
+                **custom_headers,
             }
         self.assertEqual(
             headers["X-Goog-API-Client"],
@@ -2325,6 +2331,7 @@ def _do_multipart_success(
 
             client = mock.Mock(_http=transport, _connection=_Connection, spec=["_http"])
             client._connection.API_BASE_URL = "https://storage.googleapis.com"
+            client._extra_headers = {}
 
         # Mock get_api_base_url_for_mtls function.
         mtls_url = "https://foo.mtls"
@@ -2424,11 +2431,14 @@ def _do_multipart_success(
         with patch.object(
             _helpers, "_get_invocation_id", return_value=GCCL_INVOCATION_TEST_CONST
         ):
-            headers = _get_default_headers(
-                client._connection.user_agent,
-                b'multipart/related; boundary="==0=="',
-                "application/xml",
-            )
+            headers = {
+                **_get_default_headers(
+                    client._connection.user_agent,
+                    b'multipart/related; boundary="==0=="',
+                    "application/xml",
+                ),
+                **client._extra_headers,
+            }
         client._http.request.assert_called_once_with(
             "POST", upload_url, data=payload, headers=headers, timeout=expected_timeout
         )
@@ -2520,6 +2530,19 @@ def test__do_multipart_upload_with_client(self, mock_get_boundary):
         transport = self._mock_transport(http.client.OK, {})
         client = mock.Mock(_http=transport, _connection=_Connection, spec=["_http"])
         client._connection.API_BASE_URL = "https://storage.googleapis.com"
+        client._extra_headers = {}
+        self._do_multipart_success(mock_get_boundary, client=client)
+
+    @mock.patch("google.resumable_media._upload.get_boundary", return_value=b"==0==")
+    def test__do_multipart_upload_with_client_custom_headers(self, mock_get_boundary):
+        custom_headers = {
+            "x-goog-custom-audit-foo": "bar",
+            "x-goog-custom-audit-user": "baz",
+        }
+        transport = self._mock_transport(http.client.OK, {})
+        client = mock.Mock(_http=transport, _connection=_Connection, spec=["_http"])
+        client._connection.API_BASE_URL = "https://storage.googleapis.com"
+        client._extra_headers = custom_headers
         self._do_multipart_success(mock_get_boundary, client=client)
 
     @mock.patch("google.resumable_media._upload.get_boundary", return_value=b"==0==")
@@ -2597,6 +2620,7 @@ def _initiate_resumable_helper(
             # 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._extra_headers = {}
 
         # Mock get_api_base_url_for_mtls function.
         mtls_url = "https://foo.mtls"
@@ -2677,13 +2701,15 @@ def _initiate_resumable_helper(
             _helpers, "_get_invocation_id", return_value=GCCL_INVOCATION_TEST_CONST
         ):
             if extra_headers is None:
-                self.assertEqual(
-                    upload._headers,
-                    _get_default_headers(client._connection.user_agent, content_type),
-                )
+                expected_headers = {
+                    **_get_default_headers(client._connection.user_agent, content_type),
+                    **client._extra_headers,
+                }
+                self.assertEqual(upload._headers, expected_headers)
             else:
                 expected_headers = {
                     **_get_default_headers(client._connection.user_agent, content_type),
+                    **client._extra_headers,
                     **extra_headers,
                 }
                 self.assertEqual(upload._headers, expected_headers)
@@ -2730,9 +2756,12 @@ def _initiate_resumable_helper(
         with patch.object(
             _helpers, "_get_invocation_id", return_value=GCCL_INVOCATION_TEST_CONST
         ):
-            expected_headers = _get_default_headers(
-                client._connection.user_agent, x_upload_content_type=content_type
-            )
+            expected_headers = {
+                **_get_default_headers(
+                    client._connection.user_agent, x_upload_content_type=content_type
+                ),
+                **client._extra_headers,
+            }
         if size is not None:
             expected_headers["x-upload-content-length"] = str(size)
         if extra_headers is not None:
@@ -2824,6 +2853,21 @@ def test__initiate_resumable_upload_with_client(self):
 
         client = mock.Mock(_http=transport, _connection=_Connection, spec=["_http"])
         client._connection.API_BASE_URL = "https://storage.googleapis.com"
+        client._extra_headers = {}
+        self._initiate_resumable_helper(client=client)
+
+    def test__initiate_resumable_upload_with_client_custom_headers(self):
+        custom_headers = {
+            "x-goog-custom-audit-foo": "bar",
+            "x-goog-custom-audit-user": "baz",
+        }
+        resumable_url = "http://test.invalid?upload_id=hey-you"
+        response_headers = {"location": resumable_url}
+        transport = self._mock_transport(http.client.OK, response_headers)
+
+        client = mock.Mock(_http=transport, _connection=_Connection, spec=["_http"])
+        client._connection.API_BASE_URL = "https://storage.googleapis.com"
+        client._extra_headers = custom_headers
         self._initiate_resumable_helper(client=client)
 
     def _make_resumable_transport(
@@ -3000,6 +3044,7 @@ def _do_resumable_helper(
         client = mock.Mock(_http=transport, _connection=_Connection, spec=["_http"])
         client._connection.API_BASE_URL = "https://storage.googleapis.com"
         client._connection.user_agent = USER_AGENT
+        client._extra_headers = {}
         stream = io.BytesIO(data)
 
         bucket = _Bucket(name="yesterday")
@@ -3612,26 +3657,32 @@ def _create_resumable_upload_session_helper(
         if_metageneration_match=None,
         if_metageneration_not_match=None,
         retry=None,
+        client=None,
     ):
         bucket = _Bucket(name="alex-trebek")
         blob = self._make_one("blob-name", bucket=bucket)
         chunk_size = 99 * blob._CHUNK_SIZE_MULTIPLE
         blob.chunk_size = chunk_size
-
-        # Create mocks to be checked for doing transport.
         resumable_url = "http://test.invalid?upload_id=clean-up-everybody"
-        response_headers = {"location": resumable_url}
-        transport = self._mock_transport(http.client.OK, response_headers)
-        if side_effect is not None:
-            transport.request.side_effect = side_effect
-
-        # Create some mock arguments and call the method under test.
         content_type = "text/plain"
         size = 10000
-        client = mock.Mock(_http=transport, _connection=_Connection, spec=["_http"])
-        client._connection.API_BASE_URL = "https://storage.googleapis.com"
-        client._connection.user_agent = "testing 1.2.3"
+        transport = None
 
+        if not client:
+            # Create mocks to be checked for doing transport.
+            response_headers = {"location": resumable_url}
+            transport = self._mock_transport(http.client.OK, response_headers)
+
+            # 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 = "testing 1.2.3"
+            client._extra_headers = {}
+
+        if transport is None:
+            transport = client._http
+        if side_effect is not None:
+            transport.request.side_effect = side_effect
         if timeout is None:
             expected_timeout = self._get_default_timeout()
             timeout_kwarg = {}
@@ -3689,6 +3740,7 @@ def _create_resumable_upload_session_helper(
                 **_get_default_headers(
                     client._connection.user_agent, x_upload_content_type=content_type
                 ),
+                **client._extra_headers,
                 "x-upload-content-length": str(size),
                 "x-upload-content-type": content_type,
             }
@@ -3750,6 +3802,28 @@ def test_create_resumable_upload_session_with_failure(self):
         self.assertIn(message, exc_info.exception.message)
         self.assertEqual(exc_info.exception.errors, [])
 
+    def test_create_resumable_upload_session_with_client(self):
+        resumable_url = "http://test.invalid?upload_id=clean-up-everybody"
+        response_headers = {"location": resumable_url}
+        transport = self._mock_transport(http.client.OK, response_headers)
+        client = mock.Mock(_http=transport, _connection=_Connection, spec=["_http"])
+        client._connection.API_BASE_URL = "https://storage.googleapis.com"
+        client._extra_headers = {}
+        self._create_resumable_upload_session_helper(client=client)
+
+    def test_create_resumable_upload_session_with_client_custom_headers(self):
+        custom_headers = {
+            "x-goog-custom-audit-foo": "bar",
+            "x-goog-custom-audit-user": "baz",
+        }
+        resumable_url = "http://test.invalid?upload_id=clean-up-everybody"
+        response_headers = {"location": resumable_url}
+        transport = self._mock_transport(http.client.OK, response_headers)
+        client = mock.Mock(_http=transport, _connection=_Connection, spec=["_http"])
+        client._connection.API_BASE_URL = "https://storage.googleapis.com"
+        client._extra_headers = custom_headers
+        self._create_resumable_upload_session_helper(client=client)
+
     def test_get_iam_policy_defaults(self):
         from google.cloud.storage.iam import STORAGE_OWNER_ROLE
         from google.cloud.storage.iam import STORAGE_EDITOR_ROLE
@@ -5815,6 +5889,34 @@ def test_open(self):
         with self.assertRaises(ValueError):
             blob.open("w", ignore_flush=False)
 
+    def test_downloads_w_client_custom_headers(self):
+        import google.auth.credentials
+        from google.cloud.storage import Client
+
+        custom_headers = {
+            "x-goog-custom-audit-foo": "bar",
+            "x-goog-custom-audit-user": "baz",
+        }
+        credentials = mock.Mock(spec=google.auth.credentials.Credentials)
+        client = Client(
+            project="project", credentials=credentials, extra_headers=custom_headers
+        )
+        blob = self._make_one("blob-name", bucket=_Bucket(client))
+        file_obj = io.BytesIO()
+
+        downloads = {
+            client.download_blob_to_file: (blob, file_obj),
+            blob.download_to_file: (file_obj,),
+            blob.download_as_bytes: (),
+        }
+        for method, args in downloads.items():
+            with mock.patch.object(blob, "_do_download"):
+                method(*args)
+                blob._do_download.assert_called()
+                called_headers = blob._do_download.call_args.args[-4]
+                self.assertIsInstance(called_headers, dict)
+                self.assertDictContainsSubset(custom_headers, called_headers)
+
 
 class Test__quote(unittest.TestCase):
     @staticmethod
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
diff --git a/tests/unit/test_transfer_manager.py b/tests/unit/test_transfer_manager.py
