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

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=mock_creds)
+    config = BigQueryCredentialsConfig(credentials=auth_creds)
 
     # Verify that the credentials are properly stored and attributes are extracted
-    assert config.credentials == mock_creds
+    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=oauth2_creds)
+
+    # Verify that the credentials are properly stored and attributes are extracted
+    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"]

tests/unittests/tools/bigquery/test_bigquery_credentials_manager.py

@@ -22,9 +22,10 @@ from google.adk.tools import ToolContext
 from google.adk.tools.bigquery.bigquery_credentials import BIGQUERY_TOKEN_CACHE_KEY
 from google.adk.tools.bigquery.bigquery_credentials import BigQueryCredentialsConfig
 from google.adk.tools.bigquery.bigquery_credentials import BigQueryCredentialsManager
+from google.auth.credentials import Credentials as AuthCredentials
 from google.auth.exceptions import RefreshError
 # Mock the Google OAuth and API dependencies
-from google.oauth2.credentials import Credentials
+from google.oauth2.credentials import Credentials as OAuthCredentials
 import pytest
 
 
@@ -64,9 +65,16 @@ class TestBigQueryCredentialsManager:
     """Create a credentials manager instance for testing."""
     return BigQueryCredentialsManager(credentials_config)
 
+  @pytest.mark.parametrize(
+      ("credentials_class",),
+      [
+          pytest.param(OAuthCredentials, id="oauth"),
+          pytest.param(AuthCredentials, id="auth"),
+      ],
+  )
   @pytest.mark.asyncio
   async def test_get_valid_credentials_with_valid_existing_creds(
-      self, manager, mock_tool_context
+      self, manager, mock_tool_context, credentials_class
   ):
     """Test that valid existing credentials are returned immediately.
 
@@ -74,7 +82,7 @@ class TestBigQueryCredentialsManager:
     should be needed. This is the optimal happy path scenario.
     """
     # Create mock credentials that are already valid
-    mock_creds = Mock(spec=Credentials)
+    mock_creds = Mock(spec=credentials_class)
     mock_creds.valid = True
     manager.credentials_config.credentials = mock_creds
 
@@ -85,6 +93,34 @@ class TestBigQueryCredentialsManager:
     mock_tool_context.get_auth_response.assert_not_called()
     mock_tool_context.request_credential.assert_not_called()
 
+  @pytest.mark.parametrize(
+      ("valid",),
+      [
+          pytest.param(False, id="invalid"),
+          pytest.param(True, id="valid"),
+      ],
+  )
+  @pytest.mark.asyncio
+  async def test_get_valid_credentials_with_existing_non_oauth_creds(
+      self, manager, mock_tool_context, valid
+  ):
+    """Test that existing non-oauth credentials are returned immediately.
+
+    When credentials are of non-oauth type, no refresh or OAuth flow
+    is triggered irrespective of whether it is valid or not.
+    """
+    # Create mock credentials that are already valid
+    mock_creds = Mock(spec=AuthCredentials)
+    mock_creds.valid = valid
+    manager.credentials_config.credentials = mock_creds
+
+    result = await manager.get_valid_credentials(mock_tool_context)
+
+    assert result == mock_creds
+    # Verify no OAuth flow was triggered
+    mock_tool_context.get_auth_response.assert_not_called()
+    mock_tool_context.request_credential.assert_not_called()
+
   @pytest.mark.asyncio
   async def test_get_credentials_from_cache_when_none_in_manager(
       self, manager, mock_tool_context
@@ -113,7 +149,7 @@ class TestBigQueryCredentialsManager:
     with patch(
         "google.oauth2.credentials.Credentials.from_authorized_user_info"
     ) as mock_from_json:
-      mock_creds = Mock(spec=Credentials)
+      mock_creds = Mock(spec=OAuthCredentials)
       mock_creds.valid = True
       mock_from_json.return_value = mock_creds
 
@@ -179,7 +215,7 @@ class TestBigQueryCredentialsManager:
     mock_tool_context.state[BIGQUERY_TOKEN_CACHE_KEY] = mock_cached_creds_json
 
     # Create expired cached credentials with refresh token
-    mock_cached_creds = Mock(spec=Credentials)
+    mock_cached_creds = Mock(spec=OAuthCredentials)
     mock_cached_creds.valid = False
     mock_cached_creds.expired = True
     mock_cached_creds.refresh_token = "valid_refresh_token"
@@ -227,7 +263,7 @@ class TestBigQueryCredentialsManager:
     users from having to re-authenticate for every expired token.
     """
     # Create expired credentials with refresh token
-    mock_creds = Mock(spec=Credentials)
+    mock_creds = Mock(spec=OAuthCredentials)
     mock_creds.valid = False
     mock_creds.expired = True
     mock_creds.refresh_token = "refresh_token"
@@ -257,7 +293,7 @@ class TestBigQueryCredentialsManager:
     gracefully fall back to requesting a new OAuth flow.
     """
     # Create expired credentials that fail to refresh
-    mock_creds = Mock(spec=Credentials)
+    mock_creds = Mock(spec=OAuthCredentials)
     mock_creds.valid = False
     mock_creds.expired = True
     mock_creds.refresh_token = "expired_refresh_token"
@@ -287,7 +323,7 @@ class TestBigQueryCredentialsManager:
     mock_tool_context.get_auth_response.return_value = mock_auth_response
 
     # Create a mock credentials instance that will represent our created credentials
-    mock_creds = Mock(spec=Credentials)
+    mock_creds = Mock(spec=OAuthCredentials)
     # Make the JSON match what a real Credentials object would produce
     mock_creds_json = (
         '{"token": "new_access_token", "refresh_token": "new_refresh_token",'
@@ -300,7 +336,7 @@ class TestBigQueryCredentialsManager:
 
     # Use the full module path as it appears in the project structure
     with patch(
-        "google.adk.tools.bigquery.bigquery_credentials.Credentials",
+        "google.adk.tools.bigquery.bigquery_credentials.google.oauth2.credentials.Credentials",
         return_value=mock_creds,
     ) as mock_credentials_class:
       result = await manager.get_valid_credentials(mock_tool_context)
@@ -361,7 +397,7 @@ class TestBigQueryCredentialsManager:
     mock_tool_context.get_auth_response.return_value = mock_auth_response
 
     # Create the mock credentials instance that will be returned by the constructor
-    mock_creds = Mock(spec=Credentials)
+    mock_creds = Mock(spec=OAuthCredentials)
     # Make sure our mock JSON matches the structure that real Credentials objects produce
     mock_creds_json = (
         '{"token": "cached_access_token", "refresh_token":'
@@ -376,7 +412,7 @@ class TestBigQueryCredentialsManager:
 
     # Use the correct module path - without the 'src.' prefix
     with patch(
-        "google.adk.tools.bigquery.bigquery_credentials.Credentials",
+        "google.adk.tools.bigquery.bigquery_credentials.google.oauth2.credentials.Credentials",
         return_value=mock_creds,
     ) as mock_credentials_class:
       # Complete OAuth flow with first manager
@@ -396,9 +432,9 @@ class TestBigQueryCredentialsManager:
 
     # Mock the from_authorized_user_info method for the second manager
     with patch(
-        "google.adk.tools.bigquery.bigquery_credentials.Credentials.from_authorized_user_info"
+        "google.adk.tools.bigquery.bigquery_credentials.google.oauth2.credentials.Credentials.from_authorized_user_info"
     ) as mock_from_json:
-      mock_cached_creds = Mock(spec=Credentials)
+      mock_cached_creds = Mock(spec=OAuthCredentials)
       mock_cached_creds.valid = True
       mock_from_json.return_value = mock_cached_creds