Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 3 additions & 41 deletions bigframes/_config/bigquery_options.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@
import google.auth.credentials
import requests.adapters

import bigframes._config.auth
import bigframes._importing
import bigframes.enums
import bigframes.exceptions as bfe
Expand All @@ -38,7 +37,6 @@

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
Expand Down Expand Up @@ -143,52 +141,20 @@ 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) -> google.auth.credentials.Credentials:
def credentials(self) -> Optional[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.
"""
if self._credentials:
return self._credentials

credentials, _ = self._try_set_default_credentials_and_project()
return credentials
return self._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
Expand Down Expand Up @@ -217,11 +183,7 @@ def project(self) -> Optional[str]:
None or str:
Google Cloud project ID as a string; otherwise None.
"""
if self._project:
return self._project

_, project = self._try_set_default_credentials_and_project()
return project
return self._project

@project.setter
def project(self, value: Optional[str]):
Expand Down
9 changes: 6 additions & 3 deletions bigframes/session/clients.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@
import google.cloud.storage # type: ignore
import requests

import bigframes._config
import bigframes._config.auth
import bigframes.constants
import bigframes.version

Expand All @@ -50,6 +50,10 @@
_BIGQUERYSTORAGE_REGIONAL_ENDPOINT = "bigquerystorage.{location}.rep.googleapis.com"


def _get_default_credentials_with_project():
return bigframes._config.auth.get_default_credentials_with_project()


def _get_application_names():
apps = [_APPLICATION_NAME]

Expand Down Expand Up @@ -84,8 +88,7 @@ def __init__(
):
credentials_project = None
if credentials is None:
credentials = bigframes._config.options.bigquery.credentials
credentials_project = bigframes._config.options.bigquery.project
credentials, credentials_project = _get_default_credentials_with_project()

# Prefer the project in this order:
# 1. Project explicitly specified by the user
Expand Down
5 changes: 5 additions & 0 deletions tests/unit/_config/test_bigquery_options.py
Original file line number Diff line number Diff line change
Expand Up @@ -203,3 +203,8 @@ def test_default_options():

assert options.allow_large_results is False
assert options.ordering_mode == "strict"

# We should default to None as an indicator that the user hasn't set these
# explicitly. See internal issue b/445731915.
assert options.credentials is None
assert options.project is None
5 changes: 4 additions & 1 deletion tests/unit/pandas/io/test_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
import google.cloud.bigquery
import pytest

import bigframes._config.auth
import bigframes.dataframe
import bigframes.pandas
import bigframes.pandas.io.api as bf_io_api
Expand Down Expand Up @@ -50,7 +51,7 @@ def test_read_gbq_colab_dry_run_doesnt_call_set_location(
mock_set_location.assert_not_called()


@mock.patch("bigframes._config.auth.get_default_credentials_with_project")
@mock.patch("bigframes._config.auth.pydata_google_auth.default")
@mock.patch("bigframes.core.global_session.with_default_session")
def test_read_gbq_colab_dry_run_doesnt_authenticate_multiple_times(
mock_with_default_session, mock_get_credentials, monkeypatch
Expand All @@ -77,12 +78,14 @@ def test_read_gbq_colab_dry_run_doesnt_authenticate_multiple_times(
mock_df = mock.create_autospec(bigframes.dataframe.DataFrame)
mock_with_default_session.return_value = mock_df

bigframes._config.auth._cached_credentials = None
query_or_table = "SELECT {param1} AS param1"
sample_pyformat_args = {"param1": "value1"}
bf_io_api._read_gbq_colab(
query_or_table, pyformat_args=sample_pyformat_args, dry_run=True
)

mock_get_credentials.assert_called()
mock_with_default_session.assert_not_called()
mock_get_credentials.reset_mock()

Expand Down