-
Notifications
You must be signed in to change notification settings - Fork 63
perf: avoid re-authenticating if credentials have already been fetched #2058
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
Changes from all commits
a91fbe0
14f4199
34148f8
e0544d3
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 |
|---|---|---|
| @@ -0,0 +1,57 @@ | ||
| # Copyright 2025 Google LLC | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); | ||
| # you may not use this file except in compliance with the License. | ||
| # You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, software | ||
| # distributed under the License is distributed on an "AS IS" BASIS, | ||
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| import threading | ||
| from typing import Optional | ||
|
|
||
| import google.auth.credentials | ||
| import google.auth.transport.requests | ||
| import pydata_google_auth | ||
|
|
||
| _SCOPES = ["https://www.googleapis.com/auth/cloud-platform"] | ||
|
|
||
| # Put the lock here rather than in BigQueryOptions so that BigQueryOptions | ||
| # remains deepcopy-able. | ||
| _AUTH_LOCK = threading.Lock() | ||
| _cached_credentials: Optional[google.auth.credentials.Credentials] = None | ||
| _cached_project_default: Optional[str] = None | ||
|
|
||
|
|
||
| def get_default_credentials_with_project() -> tuple[ | ||
| google.auth.credentials.Credentials, Optional[str] | ||
| ]: | ||
| global _AUTH_LOCK, _cached_credentials, _cached_project_default | ||
|
|
||
| with _AUTH_LOCK: | ||
| if _cached_credentials is not None: | ||
| return _cached_credentials, _cached_project_default | ||
|
|
||
| _cached_credentials, _cached_project_default = pydata_google_auth.default( | ||
| scopes=_SCOPES, use_local_webserver=False | ||
| ) | ||
|
|
||
| # Ensure an access token is available. | ||
| _cached_credentials.refresh(google.auth.transport.requests.Request()) | ||
|
|
||
| return _cached_credentials, _cached_project_default | ||
|
|
||
|
|
||
| def reset_default_credentials_and_project(): | ||
| global _AUTH_LOCK, _cached_credentials, _cached_project_default | ||
|
|
||
| with _AUTH_LOCK: | ||
| _cached_credentials = None | ||
| _cached_project_default = None | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -22,6 +22,7 @@ | |
| import google.auth.credentials | ||
| import requests.adapters | ||
|
|
||
| import bigframes._config.auth | ||
| import bigframes._importing | ||
| import bigframes.enums | ||
| import bigframes.exceptions as bfe | ||
|
|
@@ -37,6 +38,7 @@ | |
|
|
||
| def _get_validated_location(value: Optional[str]) -> Optional[str]: | ||
| import bigframes._tools.strings | ||
| import bigframes.constants | ||
|
|
||
| if value is None or value in bigframes.constants.ALL_BIGQUERY_LOCATIONS: | ||
| return value | ||
|
|
@@ -141,20 +143,52 @@ def application_name(self, value: Optional[str]): | |
| ) | ||
| self._application_name = value | ||
|
|
||
| def _try_set_default_credentials_and_project( | ||
| self, | ||
| ) -> tuple[google.auth.credentials.Credentials, Optional[str]]: | ||
| # Don't fetch credentials or project if credentials is already set. | ||
| # If it's set, we've already authenticated, so if the user wants to | ||
| # re-auth, they should explicitly reset the credentials. | ||
| if self._credentials is not None: | ||
| return self._credentials, self._project | ||
|
|
||
| ( | ||
| credentials, | ||
| credentials_project, | ||
| ) = bigframes._config.auth.get_default_credentials_with_project() | ||
| self._credentials = credentials | ||
|
|
||
| # Avoid overriding an explicitly set project with a default value. | ||
| if self._project is None: | ||
| self._project = credentials_project | ||
|
|
||
| return credentials, self._project | ||
|
|
||
| @property | ||
| def credentials(self) -> Optional[google.auth.credentials.Credentials]: | ||
| def credentials(self) -> google.auth.credentials.Credentials: | ||
| """The OAuth2 credentials to use for this client. | ||
|
|
||
| Set to None to force re-authentication. | ||
|
|
||
| Returns: | ||
| None or google.auth.credentials.Credentials: | ||
| google.auth.credentials.Credentials if exists; otherwise None. | ||
| """ | ||
| return self._credentials | ||
| if self._credentials: | ||
| return self._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. I am not sure how the bigquery recycles the credentials. If it's recycled by the expired time, do we need to refresh the credentials whenever we are fetching them?
Collaborator
Author
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. google-auth will automatically refresh the credentials when they are expired. |
||
|
|
||
| credentials, _ = self._try_set_default_credentials_and_project() | ||
| return credentials | ||
|
|
||
| @credentials.setter | ||
| def credentials(self, value: Optional[google.auth.credentials.Credentials]): | ||
| if self._session_started and self._credentials is not value: | ||
| raise ValueError(SESSION_STARTED_MESSAGE.format(attribute="credentials")) | ||
|
|
||
| if value is None: | ||
| # The user has _explicitly_ asked that we re-authenticate. | ||
| bigframes._config.auth.reset_default_credentials_and_project() | ||
|
|
||
| self._credentials = value | ||
|
|
||
| @property | ||
|
|
@@ -183,7 +217,11 @@ def project(self) -> Optional[str]: | |
| None or str: | ||
| Google Cloud project ID as a string; otherwise None. | ||
| """ | ||
| return self._project | ||
| if self._project: | ||
| return self._project | ||
|
|
||
| _, project = self._try_set_default_credentials_and_project() | ||
| return project | ||
|
|
||
| @project.setter | ||
| def project(self, value: Optional[str]): | ||
|
|
||
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.
If the cached credentials is failed to refresh (e.g. expired), should we reload another new credentials?
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.
Reloading generally wouldn't help. Refreshing is just exchanging a refresh token (generally long lived) for an access token (short-lived, default 1 hour, I think).
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 step shouldn't actually be necessary. I think it was added for the sole reason of having tests that check if a credential is "valid", which I don't think is populated until after the first refresh.