mirror of
https://github.com/EvolutionAPI/adk-python.git
synced 2025-07-14 01:41:25 -06:00
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:
parent
8759a25251
commit
ba5b80d5d7
@ -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)
|
||||||
|
@ -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
|
||||||
|
|
||||||
|
|
||||||
|
@ -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"
|
||||||
),
|
),
|
||||||
):
|
):
|
||||||
|
Loading…
Reference in New Issue
Block a user