-
Notifications
You must be signed in to change notification settings - Fork 348
feat: Introduce a way to provide custom headers #1811
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||
|---|---|---|---|---|
|
|
@@ -69,6 +69,7 @@ def __init__(self): | |||
|
|
||||
| self._use_non_blocking_refresh = False | ||||
| self._refresh_worker = RefreshThreadManager() | ||||
| self._custom_headers = {} | ||||
|
|
||||
| @property | ||||
| def expired(self): | ||||
|
|
@@ -185,6 +186,7 @@ def apply(self, headers, token=None): | |||
| self._apply(headers, token) | ||||
| if self.quota_project_id: | ||||
| headers["x-goog-user-project"] = self.quota_project_id | ||||
| headers.update(self._custom_headers) | ||||
|
|
||||
| def _blocking_refresh(self, request): | ||||
| if not self.valid: | ||||
|
|
@@ -233,6 +235,38 @@ def before_request(self, request, method, url, headers): | |||
| def with_non_blocking_refresh(self): | ||||
| self._use_non_blocking_refresh = True | ||||
|
|
||||
| def with_headers(self, headers): | ||||
| """Returns a copy of these credentials with additional custom headers. | ||||
|
|
||||
| Args: | ||||
| headers (Mapping[str, str]): The custom headers to add. | ||||
|
|
||||
| Returns: | ||||
| google.auth.credentials.Credentials: A new credentials instance. | ||||
|
|
||||
| Raises: | ||||
| ValueError: If a protected header is included in the input headers. | ||||
| """ | ||||
| import copy | ||||
|
|
||||
| PROTECTED_HEADERS = { | ||||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Is there a way to capture this in a more central way? I worry that we are missing some now or may miss some in the future if new ones get added for example
|
||||
| "authorization", | ||||
| "x-goog-user-project", | ||||
| "x-goog-api-client", | ||||
| "x-allowed-locations", | ||||
| } | ||||
|
|
||||
| for key in headers: | ||||
| if key.lower() in PROTECTED_HEADERS: | ||||
| raise ValueError( | ||||
| f"Header '{key}' is protected and cannot be set with with_headers. " | ||||
| "These headers are managed by the library." | ||||
| ) | ||||
|
|
||||
| new_creds = copy.deepcopy(self) | ||||
| new_creds._custom_headers.update(headers) | ||||
| return new_creds | ||||
|
|
||||
|
|
||||
| class CredentialsWithQuotaProject(Credentials): | ||||
| """Abstract base for credentials supporting ``with_quota_project`` factory""" | ||||
|
|
||||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -72,6 +72,48 @@ def test_with_non_blocking_refresh(): | |
| assert c._use_non_blocking_refresh | ||
|
|
||
|
|
||
| def test_with_headers(): | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit: this test is testing multiple scenarios. could you break it up so that we are testing ~1 thing per test? https://engdoc.corp.google.com/eng/doc/tott/episodes/726.md?cl=head |
||
| credentials = CredentialsImpl() | ||
| request = mock.Mock() | ||
|
|
||
| # 1. Add a new custom header | ||
| creds_with_header = credentials.with_headers({"X-Custom-Header": "value1"}) | ||
| headers = {} | ||
| creds_with_header.before_request(request, "http://example.com", "GET", headers) | ||
| assert headers["X-Custom-Header"] == "value1" | ||
| assert "authorization" in headers # Ensure base apply logic ran | ||
| assert creds_with_header is not credentials | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Not sure what you're trying to test here but as I understand python, this will never be true. i.e. if you did: |
||
| assert not hasattr(credentials, "_custom_headers") or not credentials._custom_headers | ||
|
|
||
| # 2. Update an existing custom header | ||
| creds_updated = creds_with_header.with_headers({"X-Custom-Header": "value2"}) | ||
| headers = {} | ||
| creds_updated.before_request(request, "http://example.com", "GET", headers) | ||
| assert headers["X-Custom-Header"] == "value2" | ||
|
|
||
| # 3. Chaining with_headers calls | ||
| creds_chained = credentials.with_headers({"X-Header-1": "v1"}).with_headers( | ||
| {"X-Header-2": "v2"} | ||
| ) | ||
| headers = {} | ||
| creds_chained.before_request(request, "http://example.com", "GET", headers) | ||
| assert headers["X-Header-1"] == "v1" | ||
| assert headers["X-Header-2"] == "v2" | ||
|
|
||
| # 4. Ensure protected headers cannot be set | ||
| with pytest.raises(ValueError): | ||
| credentials.with_headers({"Authorization": "Bearer token"}) | ||
| with pytest.raises(ValueError): | ||
| credentials.with_headers({"X-Goog-User-Project": "test"}) | ||
| with pytest.raises(ValueError): | ||
| credentials.with_headers({"authorization": "Bearer token"}) # Case-insensitive | ||
|
|
||
| # 5. Check original credentials are not modified | ||
| headers = {} | ||
| credentials.before_request(request, "http://example.com", "GET", headers) | ||
| assert "X-Custom-Header" not in headers | ||
|
|
||
|
|
||
| def test_expired_and_valid(): | ||
| credentials = CredentialsImpl() | ||
| credentials.token = "token" | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is not used?