feat: add Spanner first-party toolset (breaking change to BigQueryTool, consolidating into generic GoogleTool)

Spanner toolset support basic operations to interact with Spanner table metadata and query results.

Consolidate BigQueryTool into generic GoogleTool, so that BigQueryToolset and SpannerToolset can share.

PiperOrigin-RevId: 794259782

tests/unittests/tools/spanner/__init__

@@ -0,0 +1,13 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tests/unittests/tools/spanner/test_spanner_client.py

@@ -0,0 +1,142 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import os
+import re
+from unittest import mock
+
+from google.adk.tools.spanner.client import get_spanner_client
+from google.auth.exceptions import DefaultCredentialsError
+from google.oauth2.credentials import Credentials
+import pytest
+
+
+def test_spanner_client_project():
+  """Test spanner client project."""
+  # Trigger the spanner client creation
+  client = get_spanner_client(
+      project="test-gcp-project",
+      credentials=mock.create_autospec(Credentials, instance=True),
+  )
+
+  # Verify that the client has the desired project set
+  assert client.project == "test-gcp-project"
+
+
+def test_spanner_client_project_set_explicit():
+  """Test spanner client creation does not invoke default auth."""
+  # Let's simulate that no environment variables are set, so that any project
+  # set in there does not interfere with this test
+  with mock.patch.dict(os.environ, {}, clear=True):
+    with mock.patch("google.auth.default", autospec=True) as mock_default_auth:
+      # Simulate exception from default auth
+      mock_default_auth.side_effect = DefaultCredentialsError(
+          "Your default credentials were not found"
+      )
+
+      # Trigger the spanner client creation
+      client = get_spanner_client(
+          project="test-gcp-project",
+          credentials=mock.create_autospec(Credentials, instance=True),
+      )
+
+      # If we are here that already means client creation did not call default
+      # auth (otherwise we would have run into DefaultCredentialsError set
+      # above). For the sake of explicitness, trivially assert that the default
+      # auth was not called, and yet the project was set correctly
+      mock_default_auth.assert_not_called()
+      assert client.project == "test-gcp-project"
+
+
+def test_spanner_client_project_set_with_default_auth():
+  """Test spanner client creation invokes default auth to set the project."""
+  # Let's simulate that no environment variables are set, so that any project
+  # set in there does not interfere with this test
+  with mock.patch.dict(os.environ, {}, clear=True):
+    with mock.patch("google.auth.default", autospec=True) as mock_default_auth:
+      # Simulate credentials
+      mock_creds = mock.create_autospec(Credentials, instance=True)
+
+      # Simulate output of the default auth
+      mock_default_auth.return_value = (mock_creds, "test-gcp-project")
+
+      # Trigger the spanner client creation
+      client = get_spanner_client(
+          project=None,
+          credentials=mock_creds,
+      )
+
+      # Verify that default auth was called once to set the client project
+      mock_default_auth.assert_called_once()
+      assert client.project == "test-gcp-project"
+
+
+def test_spanner_client_project_set_with_env():
+  """Test spanner client creation sets the project from environment variable."""
+  # Let's simulate the project set in environment variables
+  with mock.patch.dict(
+      os.environ, {"GOOGLE_CLOUD_PROJECT": "test-gcp-project"}, clear=True
+  ):
+    with mock.patch("google.auth.default", autospec=True) as mock_default_auth:
+      # Simulate exception from default auth
+      mock_default_auth.side_effect = DefaultCredentialsError(
+          "Your default credentials were not found"
+      )
+
+      # Trigger the spanner client creation
+      client = get_spanner_client(
+          project=None,
+          credentials=mock.create_autospec(Credentials, instance=True),
+      )
+
+      # If we are here that already means client creation did not call default
+      # auth (otherwise we would have run into DefaultCredentialsError set
+      # above). For the sake of explicitness, trivially assert that the default
+      # auth was not called, and yet the project was set correctly
+      mock_default_auth.assert_not_called()
+      assert client.project == "test-gcp-project"
+
+
+def test_spanner_client_user_agent():
+  """Test spanner client user agent."""
+  # Patch the Client constructor
+  with mock.patch(
+      "google.cloud.spanner.Client", autospec=True
+  ) as mock_client_class:
+    # The mock instance that will be returned by spanner.Client()
+    mock_instance = mock_client_class.return_value
+    # The real spanner.Client instance has a `_client_info` attribute.
+    # We need to add it to our mock instance so that the user_agent can be set.
+    mock_instance._client_info = mock.Mock()
+
+    # Call the function that creates the client
+    client = get_spanner_client(
+        project="test-gcp-project",
+        credentials=mock.create_autospec(Credentials, instance=True),
+    )
+
+    # Verify that the Spanner Client was instantiated.
+    mock_client_class.assert_called_once_with(
+        project="test-gcp-project",
+        credentials=mock.ANY,
+    )
+
+    # Verify that the user_agent was set on the client instance.
+    # The client returned by get_spanner_client is the mock instance.
+    assert re.search(
+        r"adk-spanner-tool google-adk/([0-9A-Za-z._\-+/]+)",
+        client._client_info.user_agent,
+    )

tests/unittests/tools/spanner/test_spanner_credentials.py

@@ -0,0 +1,54 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.adk.tools.spanner.spanner_credentials import SpannerCredentialsConfig
+# Mock the Google OAuth and API dependencies
+import google.auth.credentials
+import google.oauth2.credentials
+import pytest
+
+
+class TestSpannerCredentials:
+  """Test suite for Spanner credentials configuration validation.
+
+  This class tests the credential configuration logic that ensures
+  either existing credentials or client ID/secret pairs are provided.
+  """
+
+  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=[],
+    )
+
+    config = SpannerCredentialsConfig(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/spanner.data",
+    ]
+
+    assert config._token_cache_key == "spanner_token_cache"  # pylint: disable=protected-access

tests/unittests/tools/spanner/test_spanner_tool_settings.py

@@ -0,0 +1,27 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from google.adk.tools.spanner.settings import SpannerToolSettings
+import pytest
+
+
+def test_spanner_tool_settings_experimental_warning():
+  """Test SpannerToolSettings experimental warning."""
+  with pytest.warns(
+      UserWarning,
+      match="Tool settings defaults may have breaking change in the future.",
+  ):
+    SpannerToolSettings()

tests/unittests/tools/spanner/test_spanner_toolset.py

@@ -0,0 +1,185 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from google.adk.tools.google_tool import GoogleTool
+from google.adk.tools.spanner import SpannerCredentialsConfig
+from google.adk.tools.spanner import SpannerToolset
+from google.adk.tools.spanner.settings import SpannerToolSettings
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_spanner_toolset_tools_default():
+  """Test default Spanner toolset.
+
+  This test verifies the behavior of the Spanner toolset when no filter is
+  specified.
+  """
+  credentials_config = SpannerCredentialsConfig(
+      client_id="abc", client_secret="def"
+  )
+  toolset = SpannerToolset(credentials_config=credentials_config)
+  assert isinstance(toolset._tool_settings, SpannerToolSettings)  # pylint: disable=protected-access
+  assert toolset._tool_settings.__dict__ == SpannerToolSettings().__dict__  # pylint: disable=protected-access
+  tools = await toolset.get_tools()
+  assert tools is not None
+
+  assert len(tools) == 6
+  assert all([isinstance(tool, GoogleTool) for tool in tools])
+
+  expected_tool_names = set([
+      "list_table_names",
+      "list_table_indexes",
+      "list_table_index_columns",
+      "list_named_schemas",
+      "get_table_schema",
+      "execute_sql",
+  ])
+  actual_tool_names = set([tool.name for tool in tools])
+  assert actual_tool_names == expected_tool_names
+
+
+@pytest.mark.parametrize(
+    "selected_tools",
+    [
+        pytest.param([], id="None"),
+        pytest.param(
+            ["list_table_names", "get_table_schema"],
+            id="table-metadata",
+        ),
+        pytest.param(["execute_sql"], id="query"),
+    ],
+)
+@pytest.mark.asyncio
+async def test_spanner_toolset_selective(selected_tools):
+  """Test selective Spanner toolset.
+
+  This test verifies the behavior of the Spanner toolset when a filter is
+  specified.
+
+  Args:
+      selected_tools: A list of tool names to filter.
+  """
+  credentials_config = SpannerCredentialsConfig(
+      client_id="abc", client_secret="def"
+  )
+  toolset = SpannerToolset(
+      credentials_config=credentials_config,
+      tool_filter=selected_tools,
+      spanner_tool_settings=SpannerToolSettings(),
+  )
+  tools = await toolset.get_tools()
+  assert tools is not None
+
+  assert len(tools) == len(selected_tools)
+  assert all([isinstance(tool, GoogleTool) for tool in tools])
+
+  expected_tool_names = set(selected_tools)
+  actual_tool_names = set([tool.name for tool in tools])
+  assert actual_tool_names == expected_tool_names
+
+
+@pytest.mark.parametrize(
+    ("selected_tools", "returned_tools"),
+    [
+        pytest.param(["unknown"], [], id="all-unknown"),
+        pytest.param(
+            ["unknown", "execute_sql"],
+            ["execute_sql"],
+            id="mixed-known-unknown",
+        ),
+    ],
+)
+@pytest.mark.asyncio
+async def test_spanner_toolset_unknown_tool(selected_tools, returned_tools):
+  """Test Spanner toolset with unknown tools.
+
+  This test verifies the behavior of the Spanner toolset when unknown tools are
+  specified in the filter.
+
+  Args:
+      selected_tools: A list of tool names to filter, including unknown ones.
+      returned_tools: A list of tool names that are expected to be returned.
+  """
+  credentials_config = SpannerCredentialsConfig(
+      client_id="abc", client_secret="def"
+  )
+
+  toolset = SpannerToolset(
+      credentials_config=credentials_config,
+      tool_filter=selected_tools,
+      spanner_tool_settings=SpannerToolSettings(),
+  )
+
+  tools = await toolset.get_tools()
+  assert tools is not None
+
+  assert len(tools) == len(returned_tools)
+  assert all([isinstance(tool, GoogleTool) for tool in tools])
+
+  expected_tool_names = set(returned_tools)
+  actual_tool_names = set([tool.name for tool in tools])
+  assert actual_tool_names == expected_tool_names
+
+
+@pytest.mark.parametrize(
+    ("selected_tools", "returned_tools"),
+    [
+        pytest.param(
+            ["execute_sql", "list_table_names"],
+            ["list_table_names"],
+            id="read-not-added",
+        ),
+        pytest.param(
+            ["list_table_names", "list_table_indexes"],
+            ["list_table_names", "list_table_indexes"],
+            id="no-effect",
+        ),
+    ],
+)
+@pytest.mark.asyncio
+async def test_spanner_toolset_without_read_capability(
+    selected_tools, returned_tools
+):
+  """Test Spanner toolset without read capability.
+
+  This test verifies the behavior of the Spanner toolset when read capability is
+  not enabled.
+
+  Args:
+      selected_tools: A list of tool names to filter.
+      returned_tools: A list of tool names that are expected to be returned.
+  """
+  credentials_config = SpannerCredentialsConfig(
+      client_id="abc", client_secret="def"
+  )
+
+  spanner_tool_settings = SpannerToolSettings(capabilities=[])
+  toolset = SpannerToolset(
+      credentials_config=credentials_config,
+      tool_filter=selected_tools,
+      spanner_tool_settings=spanner_tool_settings,
+  )
+
+  tools = await toolset.get_tools()
+  assert tools is not None
+
+  assert len(tools) == len(returned_tools)
+  assert all([isinstance(tool, GoogleTool) for tool in tools])
+
+  expected_tool_names = set(returned_tools)
+  actual_tool_names = set([tool.name for tool in tools])
+  assert actual_tool_names == expected_tool_names