8000 fix: Allow more credentials types for BigQuery tools by copybara-service[bot] · Pull Request #1483 · google/adk-python · GitHub
[go: up one dir, main page]

Skip to content

fix: Allow more credentials types for BigQuery tools #1483

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

Merged
merged 1 commit into from
Jun 19, 2025
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
fix: Allow more credentials types for BigQuery tools
This change accepts the `google.auth.credentials.Credentials` type for `BigQueryCredentialsConfig`, so any subclass of that, including `google.oauth2.credentials.Credentials` would work to integrate with BigQuery service. This opens up a whole range of possibilities, such as using service account credentials to deploy an agent using these tools.

PiperOrigin-RevId: 773190440
  • Loading branch information
google-genai-bot authored and copybara-github committed Jun 19, 2025
commit 2f716ada7fbcf8e03ff5ae16ce26a80ca6fd7bf6
23 changes: 19 additions & 4 deletions contributing/samples/bigquery/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -40,13 +40,28 @@ would set:
### With Application Default Credentials

This mode is useful for quick development when the agent builder is the only
user interacting with the agent. The tools are initialized with the default
credentials present on the machine running the agent.
user interacting with the agent. The tools are run with these credentials.

1. Create application default credentials on the machine where the agent would
be running by following https://cloud.google.com/docs/authentication/provide-credentials-adc.

1. Set `RUN_WITH_ADC=True` in `agent.py` and run the agent
1. Set `CREDENTIALS_TYPE=None` in `agent.py`

1. Run the agent

### With Service Account Keys

This mode is useful for quick development when the agent builder wants to run
the agent with service account credentials. The tools are run with these
credentials.

1. Create service account key by following https://cloud.google.com/iam/docs/service-account-creds#user-managed-keys.

1. Set `CREDENTIALS_TYPE=AuthCredentialTypes.SERVICE_ACCOUNT` in `agent.py`

1. Download the key file and replace `"service_account_key.json"` with the path

1. Run the agent

### With Interactive OAuth

Expand All @@ -72,7 +87,7 @@ type.
Note: don't create a separate .env, instead put it to the same .env file that
stores your Vertex AI or Dev ML credentials

1. Set `RUN_WITH_ADC=False` in `agent.py` and run the agent
1. Set `CREDENTIALS_TYPE=AuthCredentialTypes.OAUTH2` in `agent.py` and run the agent

## Sample prompts

Expand Down
27 changes: 19 additions & 8 deletions contributing/samples/bigquery/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,31 +15,42 @@
import os

from google.adk.agents import llm_agent
from google.adk.auth import AuthCredentialTypes
from google.adk.tools.bigquery import BigQueryCredentialsConfig
from google.adk.tools.bigquery import BigQueryToolset
from google.adk.tools.bigquery.config import BigQueryToolConfig
from google.adk.tools.bigquery.config import WriteMode
import google.auth

RUN_WITH_ADC = False
# Define an appropriate credential type
CREDENTIALS_TYPE = AuthCredentialTypes.OAUTH2


# Define BigQuery tool config
tool_config = BigQueryToolConfig(write_mode=WriteMode.ALLOWED)

if RUN_WITH_ADC:< 8000 /span>
# Initialize the tools to use the application default credentials.
application_default_credentials, _ = google.auth.default()
credentials_config = BigQueryCredentialsConfig(
credentials=application_default_credentials
)
else:
if CREDENTIALS_TYPE == AuthCredentialTypes.OAUTH2:
# Initiaze the tools to do interactive OAuth
# The environment variables OAUTH_CLIENT_ID and OAUTH_CLIENT_SECRET
# must be set
credentials_config = BigQueryCredentialsConfig(
client_id=os.getenv("OAUTH_CLIENT_ID"),
client_secret=os.getenv("OAUTH_CLIENT_SECRET"),
)
elif CREDENTIALS_TYPE == AuthCredentialTypes.SERVICE_ACCOUNT:
# Initialize the tools to use the credentials in the service account key.
# If this flow is enabled, make sure to replace the file path with your own
# service account key file
# https://cloud.google.com/iam/docs/service-account-creds#user-managed-keys
creds, _ = google.auth.load_credentials_from_file("service_account_key.json")
credentials_config = BigQueryCredentialsConfig(credentials=creds)
else:
# Initialize the tools to use the application default credentials.
# https://cloud.google.com/docs/authentication/provide-credentials-adc
application_default_credentials, _ = google.auth.default()
credentials_config = BigQueryCredentialsConfig(
credentials=application_default_credentials
)

bigquery_toolset = BigQueryToolset(
credentials_config=credentials_config, bigquery_tool_config=tool_config
Expand Down
39 changes: 28 additions & 11 deletions src/google/adk/tools/bigquery/bigquery_credentials.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,10 @@
from fastapi.openapi.models import OAuth2
from fastapi.openapi.models import OAuthFlowAuthorizationCode
from fastapi.openapi.models import OAuthFlows
import google.auth.credentials
from google.auth.exceptions import RefreshError
from google.auth.transport.requests import Request
from google.oauth2.credentials import Credentials
import google.oauth2.credentials
from pydantic import BaseModel
from pydantic import model_validator

Expand All @@ -40,26 +41,35 @@

@experimental
class BigQueryCredentialsConfig(BaseModel):
"""Configuration for Google API tools. (Experimental)"""
"""Configuration for Google API tools (Experimental).

Please do not use this in production, as it may be deprecated later.
"""

# Configure the model to allow arbitrary types like Credentials
model_config = {"arbitrary_types_allowed": True}

credentials: Optional[Credentials] = None
"""the existing oauth credentials to use. If set,this credential will be used
credentials: Optional[google.auth.credentials.Credentials] = None
"""The existing auth credentials to use. If set, this credential will be used
for every end user, end users don't need to be involved in the oauthflow. This
field is mutually exclusive with client_id, client_secret and scopes.
Don't set this field unless you are sure this credential has the permission to
access every end user's data.

Example usage: when the agent is deployed in Google Cloud environment and
Example usage 1: When the agent is deployed in Google Cloud environment and
the service account (used as application default credentials) has access to
all the required BigQuery resource. Setting this credential to allow user to
access the BigQuery resource without end users going through oauth flow.

To get application default credential: `google.auth.default(...)`. See more
To get application default credential, use: `google.auth.default(...)`. See more
details in https://cloud.google.com/docs/authentication/application-default-credentials.

Example usage 2: When the agent wants to access the user's BigQuery resources
using the service account key credentials.

To load service account key credentials, use: `google.auth.load_credentials_from_file(...)`.
See more details in https://cloud.google.com/iam/docs/service-account-creds#user-managed-keys.

When the deployed environment cannot provide a pre-existing credential,
consider setting below client_id, client_secret and scope for end users to go
through oauth flow, so that agent can access the user data.
Expand All @@ -86,7 +96,9 @@ def __post_init__(self) -> BigQueryCredentialsConfig:
" client_id/client_secret/scopes."
)

if self.credentials:
if self.credentials and isinstance(
self.credentials, google.oauth2.credentials.Credentials
):
self.client_id = self.credentials.client_id
self.client_secret = self.credentials.client_secret
self.scopes = self.credentials.scopes
Expand Down Expand Up @@ -115,7 +127,7 @@ def __init__(self, credentials_config: BigQueryCredentialsConfig):

async def get_valid_credentials(
self, tool_context: ToolContext
) -> Optional[Credentials]:
) -> Optional[google.auth.credentials.Credentials]:
"""Get valid credentials, handling refresh and OAuth flow as needed.

Args:
Expand All @@ -127,7 +139,7 @@ async def get_valid_credentials(
# First, try to get credentials from the tool context
creds_json = tool_context.state.get(BIGQUERY_TOKEN_CACHE_KEY, None)
creds = (
Credentials.from_authorized_user_info(
google.oauth2.credentials.Credentials.from_authorized_user_info(
json.loads(creds_json), self.credentials_config.scopes
)
if creds_json
Expand All @@ -138,6 +150,11 @@ async def get_valid_credentials(
if not creds:
creds = self.credentials_config.credentials

# If non-oauth credentials are provided then use them as is. This helps
# in flows such as service account keys
if creds and not isinstance(creds, google.oauth2.credentials.Credentials):
return creds

# Check if we have valid credentials
if creds and creds.valid:
return creds
Expand All @@ -159,7 +176,7 @@ async def get_valid_credentials(

async def _perform_oauth_flow(
self, tool_context: ToolContext
) -> Optional[Credentials]:
) -> Optional[google.oauth2.credentials.Credentials]:
"""Perform OAuth flow to get new credentials.

Args:
Expand Down Expand Up @@ -199,7 +216,7 @@ async def _perform_oauth_flow(

if auth_response:
# OAuth flow completed, create credentials
creds = Credentials(
creds = google.oauth2.credentials.Credentials(
token=auth_response.oauth2.access_token,
refresh_token=auth_response.oauth2.refresh_token,
token_uri=auth_scheme.flows.authorizationCode.tokenUrl,
Expand Down
2 changes: 1 addition & 1 deletion src/google/adk/tools/bigquery/bigquery_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
from typing import Callable
from typing import Optional

from google.oauth2.credentials import Credentials
from google.auth.credentials import Credentials
from typing_extensions import override

from ...utils.feature_decorator import experimental
Expand Down
2 changes: 1 addition & 1 deletion src/google/adk/tools/bigquery/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,8 @@
from __future__ import annotations

import google.api_core.client_info
from google.auth.credentials import Credentials
from google.cloud import bigquery
from google.oauth2.credentials import Credentials

from ... import version

Expand Down
2 changes: 1 addition & 1 deletion src/google/adk/tools/bigquery/metadata_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.

from google.auth.credentials import Credentials
from google.cloud import bigquery
from google.oauth2.credentials import Credentials

from . import client

Expand Down
2 changes: 1 addition & 1 deletion src/google/adk/tools/bigquery/query_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,8 @@
import types
from typing import Callable

from google.auth.credentials import Credentials
from google.cloud import bigquery
from google.oauth2.credentials import Credentials

from . import client
from .config import BigQueryToolConfig
Expand Down
47 changes: 36 additions & 11 deletions tests/unittests/tools/bigquery/test_bigquery_credentials.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,11 +12,12 @@
# See the License for the specific language governing permissions and
# limitations under the License.

from unittest.mock import Mock
from unittest import mock

from google.adk.tools.bigquery.bigquery_credentials import BigQueryCredentialsConfig
# Mock the Google OAuth and API dependencies
from google.oauth2.credentials import Credentials
import google.auth.credentials
import google.oauth2.credentials
import pytest


Expand All @@ -27,22 +28,46 @@ class TestBigQueryCredentials:
either existing credentials or client ID/secret pairs are provided.
"""

def test_valid_credentials_object(self):
"""Test that providing valid Credentials object works correctly.
def test_valid_credentials_object_auth_credentials(self):
"""Test that providing valid Credentials object works correctly with
google.auth.credentials.Credentials.
When a user already has valid OAuth credentials, they should be able
to pass them directly without needing to provide client ID/secret.
"""
# Create a mock credentials object with the expected attributes
mock_creds = Mock(spec=Credentials)
mock_creds.client_id = "test_client_id"
mock_creds.client_secret = "test_client_secret"
mock_creds.scopes = ["https://www.googleapis.com/auth/calendar"]
# Create a mock auth credentials object
# auth_creds = google.auth.credentials.Credentials()
auth_creds = mock.create_autospec(
google.auth.credentials.Credentials, instance=True
)

config = BigQueryCredentialsConfig(credentials=auth_creds)

# Verify that the credentials are properly stored and attributes are extracted
assert config.credentials == auth_creds
assert config.client_id is None
assert config.client_secret is None
assert config.scopes == ["https://www.googleapis.com/auth/bigquery"]

def test_valid_credentials_object_oauth2_credentials(self):
"""Test that providing valid Credentials object works correctly with
google.oauth2.credentials.Credentials.
When a user already has valid OAuth credentials, they should be able
to pass them directly without needing to provide client ID/secret.
"""
# Create a mock oauth2 credentials object
oauth2_creds = google.oauth2.credentials.Credentials(
"test_token",
client_id="test_client_id",
client_secret="test_client_secret",
scopes=["https://www.googleapis.com/auth/calendar"],
)

config = BigQueryCredentialsConfig(credentials=mock_creds)
config = BigQueryCredentialsConfig(credentials=oauth2_creds)

# Verify that the credentials are properly stored and attributes are extracted
assert config.credentials == mock_creds
assert config.credentials == oauth2_creds
assert config.client_id == "test_client_id"
assert config.client_secret == "test_client_secret"
assert config.scopes == ["https://www.googleapis.com/auth/calendar"]
Expand Down
Loading
Loading
0