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

Author: google-genai-bot

contributing/samples/bigquery/README.md

@@ -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
 
@@ -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
 

contributing/samples/bigquery/agent.py

@@ -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:
- # 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

src/google/adk/tools/bigquery/bigquery_credentials.py

@@ -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
 
@@ -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.
@@ -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
@@ -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:
@@ -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
@@ -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
@@ -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:
@@ -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,

src/google/adk/tools/bigquery/bigquery_tool.py

@@ -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

src/google/adk/tools/bigquery/client.py

@@ -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
 

src/google/adk/tools/bigquery/metadata_tool.py

@@ -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
 

src/google/adk/tools/bigquery/query_tool.py

@@ -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

tests/unittests/tools/bigquery/test_bigquery_credentials.py

@@ -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
 
 
@@ -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"]