-
Notifications
You must be signed in to change notification settings - Fork 336
feat: Introduce a way to provide custom headers #1811
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
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 = { | ||||
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 |
---|---|---|
|
@@ -71,6 +71,48 @@ def test_with_non_blocking_refresh(): | |
assert c._use_non_blocking_refresh | ||
|
||
|
||
def test_with_headers(): | ||
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 | ||
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?