feat: use bigquery scope by default in bigquery credentials.

Right now the agent builder has to specify the bigquery scope explicitly for bigquery tools, which is somewhat unnecessary. With this change the user does not have to specify scopes, although they would still have the ability to overrides scopes if necessary.

PiperOrigin-RevId: 765354010
This commit is contained in:
Google Team Member 2025-05-30 15:12:20 -07:00 committed by Copybara-Service
parent 8759a25251
commit ba5b80d5d7
3 changed files with 56 additions and 11 deletions

View File

@ -36,7 +36,6 @@ else:
credentials_config = BigQueryCredentialsConfig( credentials_config = BigQueryCredentialsConfig(
client_id=os.getenv("OAUTH_CLIENT_ID"), client_id=os.getenv("OAUTH_CLIENT_ID"),
client_secret=os.getenv("OAUTH_CLIENT_SECRET"), client_secret=os.getenv("OAUTH_CLIENT_SECRET"),
scopes=["https://www.googleapis.com/auth/bigquery"],
) )
bigquery_toolset = BigQueryToolset(credentials_config=credentials_config) bigquery_toolset = BigQueryToolset(credentials_config=credentials_config)

View File

@ -34,6 +34,7 @@ from ...auth import OAuth2Auth
from ..tool_context import ToolContext from ..tool_context import ToolContext
BIGQUERY_TOKEN_CACHE_KEY = "bigquery_token_cache" BIGQUERY_TOKEN_CACHE_KEY = "bigquery_token_cache"
BIGQUERY_DEFAULT_SCOPE = ["https://www.googleapis.com/auth/bigquery"]
class BigQueryCredentialsConfig(BaseModel): class BigQueryCredentialsConfig(BaseModel):
@ -66,15 +67,14 @@ class BigQueryCredentialsConfig(BaseModel):
client_secret: Optional[str] = None client_secret: Optional[str] = None
"""the oauth client secret to use.""" """the oauth client secret to use."""
scopes: Optional[List[str]] = None scopes: Optional[List[str]] = None
"""the scopes to use. """the scopes to use."""
"""
@model_validator(mode="after") @model_validator(mode="after")
def __post_init__(self) -> BigQueryCredentialsConfig: def __post_init__(self) -> BigQueryCredentialsConfig:
"""Validate that either credentials or client ID/secret are provided.""" """Validate that either credentials or client ID/secret are provided."""
if not self.credentials and (not self.client_id or not self.client_secret): if not self.credentials and (not self.client_id or not self.client_secret):
raise ValueError( raise ValueError(
"Must provide either credentials or client_id abd client_secret pair." "Must provide either credentials or client_id and client_secret pair."
) )
if self.credentials and ( if self.credentials and (
self.client_id or self.client_secret or self.scopes self.client_id or self.client_secret or self.scopes
@ -88,6 +88,10 @@ class BigQueryCredentialsConfig(BaseModel):
self.client_id = self.credentials.client_id self.client_id = self.credentials.client_id
self.client_secret = self.credentials.client_secret self.client_secret = self.credentials.client_secret
self.scopes = self.credentials.scopes self.scopes = self.credentials.scopes
if not self.scopes:
self.scopes = BIGQUERY_DEFAULT_SCOPE
return self return self

View File

@ -47,16 +47,58 @@ class TestBigQueryCredentials:
assert config.client_secret == "test_client_secret" assert config.client_secret == "test_client_secret"
assert config.scopes == ["https://www.googleapis.com/auth/calendar"] assert config.scopes == ["https://www.googleapis.com/auth/calendar"]
def test_valid_client_id_secret_pair(self): def test_valid_client_id_secret_pair_default_scope(self):
"""Test that providing client ID and secret without credentials works. """Test that providing client ID and secret with default scope works.
This tests the scenario where users want to create new OAuth credentials This tests the scenario where users want to create new OAuth credentials
from scratch using their application's client ID and secret. from scratch using their application's client ID and secret and does not
specify the scopes explicitly.
""" """
config = BigQueryCredentialsConfig( config = BigQueryCredentialsConfig(
client_id="test_client_id", client_id="test_client_id",
client_secret="test_client_secret", client_secret="test_client_secret",
scopes=["https://www.googleapis.com/auth/bigquery"], )
assert config.credentials is None
assert config.client_id == "test_client_id"
assert config.client_secret == "test_client_secret"
assert config.scopes == ["https://www.googleapis.com/auth/bigquery"]
def test_valid_client_id_secret_pair_w_scope(self):
"""Test that providing client ID and secret with explicit scopes works.
This tests the scenario where users want to create new OAuth credentials
from scratch using their application's client ID and secret and does specify
the scopes explicitly.
"""
config = BigQueryCredentialsConfig(
client_id="test_client_id",
client_secret="test_client_secret",
scopes=[
"https://www.googleapis.com/auth/bigquery",
"https://www.googleapis.com/auth/drive",
],
)
assert config.credentials is None
assert config.client_id == "test_client_id"
assert config.client_secret == "test_client_secret"
assert config.scopes == [
"https://www.googleapis.com/auth/bigquery",
"https://www.googleapis.com/auth/drive",
]
def test_valid_client_id_secret_pair_w_empty_scope(self):
"""Test that providing client ID and secret with empty scope works.
This tests the corner case scenario where users want to create new OAuth
credentials from scratch using their application's client ID and secret but
specifies empty scope, in which case the default BQ scope is used.
"""
config = BigQueryCredentialsConfig(
client_id="test_client_id",
client_secret="test_client_secret",
scopes=[],
) )
assert config.credentials is None assert config.credentials is None
@ -73,7 +115,7 @@ class TestBigQueryCredentials:
with pytest.raises( with pytest.raises(
ValueError, ValueError,
match=( match=(
"Must provide either credentials or client_id abd client_secret" "Must provide either credentials or client_id and client_secret"
" pair" " pair"
), ),
): ):
@ -84,7 +126,7 @@ class TestBigQueryCredentials:
with pytest.raises( with pytest.raises(
ValueError, ValueError,
match=( match=(
"Must provide either credentials or client_id abd client_secret" "Must provide either credentials or client_id and client_secret"
" pair" " pair"
), ),
): ):
@ -99,7 +141,7 @@ class TestBigQueryCredentials:
with pytest.raises( with pytest.raises(
ValueError, ValueError,
match=( match=(
"Must provide either credentials or client_id abd client_secret" "Must provide either credentials or client_id and client_secret"
" pair" " pair"
), ),
): ):