feat: add Bigtable tools

These tools support basic operations to interact with Bigtable table metadata and query results.

PiperOrigin-RevId: 796571736

tests/unittests/tools/bigtable/__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/bigtable/test_bigtable_credentials.py

@@ -0,0 +1,91 @@
+# 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 unittest import mock
+
+from google.adk.tools.bigtable.bigtable_credentials import BIGTABLE_DEFAULT_SCOPE
+from google.adk.tools.bigtable.bigtable_credentials import BigtableCredentialsConfig
+from google.auth.credentials import Credentials
+import google.oauth2.credentials
+import pytest
+
+
+class TestBigtableCredentials:
+  """Test suite for Bigtable credentials configuration validation.
+
+  This class tests the credential configuration logic that ensures
+  either existing credentials or client ID/secret pairs are provided.
+  """
+
+  def test_bigtable_credentials_config_client_id_secret(self):
+    """Test BigtableCredentialsConfig with client_id and client_secret.
+
+    Ensures that when client_id and client_secret are provided, the config
+    object is created with the correct attributes.
+    """
+    config = BigtableCredentialsConfig(client_id="abc", client_secret="def")
+    assert config.client_id == "abc"
+    assert config.client_secret == "def"
+    assert config.scopes == BIGTABLE_DEFAULT_SCOPE
+    assert config.credentials is None
+
+  def test_bigtable_credentials_config_existing_creds(self):
+    """Test BigtableCredentialsConfig with existing generic credentials.
+
+    Ensures that when a generic Credentials object is provided, it is
+    stored correctly.
+    """
+    mock_creds = mock.create_autospec(Credentials, instance=True)
+    config = BigtableCredentialsConfig(credentials=mock_creds)
+    assert config.credentials == mock_creds
+    assert config.client_id is None
+    assert config.client_secret is None
+
+  def test_bigtable_credentials_config_oauth2_creds(self):
+    """Test BigtableCredentialsConfig with existing OAuth2 credentials.
+
+    Ensures that when a google.oauth2.credentials.Credentials object is
+    provided, the client_id, client_secret, and scopes are extracted
+    from the credentials object.
+    """
+    mock_creds = mock.create_autospec(
+        google.oauth2.credentials.Credentials, instance=True
+    )
+    mock_creds.client_id = "oauth_client_id"
+    mock_creds.client_secret = "oauth_client_secret"
+    mock_creds.scopes = ["fake_scope"]
+    config = BigtableCredentialsConfig(credentials=mock_creds)
+    assert config.client_id == "oauth_client_id"
+    assert config.client_secret == "oauth_client_secret"
+    assert config.scopes == ["fake_scope"]
+
+  def test_bigtable_credentials_config_validation_errors(self):
+    """Test BigtableCredentialsConfig validation errors.
+
+    Ensures that ValueError is raised under the following conditions:
+    - No arguments are provided.
+    - Only client_id is provided.
+    - Both credentials and client_id/client_secret are provided.
+    """
+    with pytest.raises(ValueError):
+      BigtableCredentialsConfig()
+
+    with pytest.raises(ValueError):
+      BigtableCredentialsConfig(client_id="abc")
+
+    mock_creds = mock.create_autospec(Credentials, instance=True)
+    with pytest.raises(ValueError):
+      BigtableCredentialsConfig(
+          credentials=mock_creds, client_id="abc", client_secret="def"
+      )

tests/unittests/tools/bigtable/test_bigtable_metadata_tool.py

@@ -0,0 +1,137 @@
+# 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.
+
+import logging
+from unittest import mock
+
+from google.adk.tools.bigtable import metadata_tool
+from google.auth.credentials import Credentials
+
+
+def test_list_instances():
+  """Test list_instances function."""
+  with mock.patch(
+      "google.adk.tools.bigtable.client.get_bigtable_admin_client"
+  ) as mock_get_client:
+    mock_client = mock.MagicMock()
+    mock_get_client.return_value = mock_client
+    mock_instance = mock.MagicMock()
+    mock_instance.instance_id = "test-instance"
+    mock_client.list_instances.return_value = ([mock_instance], [])
+
+    creds = mock.create_autospec(Credentials, instance=True)
+    result = metadata_tool.list_instances("test-project", creds)
+    assert result == {"status": "SUCCESS", "results": ["test-instance"]}
+
+
+def test_list_instances_failed_locations():
+  """Test list_instances function when some locations fail."""
+  with mock.patch(
+      "google.adk.tools.bigtable.client.get_bigtable_admin_client"
+  ) as mock_get_client:
+    with mock.patch.object(logging, "warning") as mock_warning:
+      mock_client = mock.MagicMock()
+      mock_get_client.return_value = mock_client
+      mock_instance = mock.MagicMock()
+      mock_instance.instance_id = "test-instance"
+      failed_locations = ["us-west1-a"]
+      mock_client.list_instances.return_value = (
+          [mock_instance],
+          failed_locations,
+      )
+
+      creds = mock.create_autospec(Credentials, instance=True)
+      result = metadata_tool.list_instances("test-project", creds)
+      assert result == {"status": "SUCCESS", "results": ["test-instance"]}
+      mock_warning.assert_called_once_with(
+          "Failed to list instances from the following locations: %s",
+          failed_locations,
+      )
+
+
+def test_get_instance_info():
+  """Test get_instance_info function."""
+  with mock.patch(
+      "google.adk.tools.bigtable.client.get_bigtable_admin_client"
+  ) as mock_get_client:
+    mock_client = mock.MagicMock()
+    mock_get_client.return_value = mock_client
+    mock_instance = mock.MagicMock()
+    mock_client.instance.return_value = mock_instance
+    mock_instance.instance_id = "test-instance"
+    mock_instance.display_name = "Test Instance"
+    mock_instance.state = "READY"
+    mock_instance.type_ = "PRODUCTION"
+    mock_instance.labels = {"env": "test"}
+
+    creds = mock.create_autospec(Credentials, instance=True)
+    result = metadata_tool.get_instance_info(
+        "test-project", "test-instance", creds
+    )
+    expected_result = {
+        "project_id": "test-project",
+        "instance_id": "test-instance",
+        "display_name": "Test Instance",
+        "state": "READY",
+        "type": "PRODUCTION",
+        "labels": {"env": "test"},
+    }
+    assert result == {"status": "SUCCESS", "results": expected_result}
+    mock_instance.reload.assert_called_once()
+
+
+def test_list_tables():
+  """Test list_tables function."""
+  with mock.patch(
+      "google.adk.tools.bigtable.client.get_bigtable_admin_client"
+  ) as mock_get_client:
+    mock_client = mock.MagicMock()
+    mock_get_client.return_value = mock_client
+    mock_instance = mock.MagicMock()
+    mock_client.instance.return_value = mock_instance
+    mock_table = mock.MagicMock()
+    mock_table.table_id = "test-table"
+    mock_instance.list_tables.return_value = [mock_table]
+
+    creds = mock.create_autospec(Credentials, instance=True)
+    result = metadata_tool.list_tables("test-project", "test-instance", creds)
+    assert result == {"status": "SUCCESS", "results": ["test-table"]}
+
+
+def test_get_table_info():
+  """Test get_table_info function."""
+  with mock.patch(
+      "google.adk.tools.bigtable.client.get_bigtable_admin_client"
+  ) as mock_get_client:
+    mock_client = mock.MagicMock()
+    mock_get_client.return_value = mock_client
+    mock_instance = mock.MagicMock()
+    mock_client.instance.return_value = mock_instance
+    mock_table = mock.MagicMock()
+    mock_instance.table.return_value = mock_table
+    mock_table.table_id = "test-table"
+    mock_instance.instance_id = "test-instance"
+    mock_table.list_column_families.return_value = {"cf1": mock.MagicMock()}
+
+    creds = mock.create_autospec(Credentials, instance=True)
+    result = metadata_tool.get_table_info(
+        "test-project", "test-instance", "test-table", creds
+    )
+    expected_result = {
+        "project_id": "test-project",
+        "instance_id": "test-instance",
+        "table_id": "test-table",
+        "column_families": ["cf1"],
+    }
+    assert result == {"status": "SUCCESS", "results": expected_result}

tests/unittests/tools/bigtable/test_bigtable_query_tool.py

@@ -0,0 +1,137 @@
+# 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 typing import Optional
+from unittest import mock
+
+from google.adk.tools.base_tool import BaseTool
+from google.adk.tools.bigtable import BigtableCredentialsConfig
+from google.adk.tools.bigtable.bigtable_toolset import BigtableToolset
+from google.adk.tools.bigtable.query_tool import execute_sql
+from google.adk.tools.bigtable.settings import BigtableToolSettings
+from google.adk.tools.tool_context import ToolContext
+from google.auth.credentials import Credentials
+from google.cloud import bigtable
+from google.cloud.bigtable.data.execute_query import ExecuteQueryIterator
+import pytest
+
+
+def test_execute_sql_basic():
+  """Test execute_sql tool basic functionality."""
+  project = "my_project"
+  instance_id = "my_instance"
+  query = "SELECT * FROM my_table"
+  credentials = mock.create_autospec(Credentials, instance=True)
+  tool_context = mock.create_autospec(ToolContext, instance=True)
+
+  with mock.patch(
+      "google.adk.tools.bigtable.client.get_bigtable_data_client"
+  ) as mock_get_client:
+    mock_client = mock.MagicMock()
+    mock_get_client.return_value = mock_client
+    mock_iterator = mock.create_autospec(ExecuteQueryIterator, instance=True)
+    mock_client.execute_query.return_value = mock_iterator
+
+    # Mock row data
+    mock_row = mock.MagicMock()
+    mock_row.fields = {"col1": "val1", "col2": 123}
+    mock_iterator.__iter__.return_value = [mock_row]
+
+    result = execute_sql(
+        project_id=project,
+        instance_id=instance_id,
+        credentials=credentials,
+        query=query,
+        settings=BigtableToolSettings(),
+        tool_context=tool_context,
+    )
+
+    expected_rows = [{"col1": "val1", "col2": 123}]
+    assert result == {"status": "SUCCESS", "rows": expected_rows}
+    mock_client.execute_query.assert_called_once_with(
+        query=query, instance_id=instance_id
+    )
+    mock_iterator.close.assert_called_once()
+
+
+def test_execute_sql_truncated():
+  """Test execute_sql tool truncation functionality."""
+  project = "my_project"
+  instance_id = "my_instance"
+  query = "SELECT * FROM my_table"
+  credentials = mock.create_autospec(Credentials, instance=True)
+  tool_context = mock.create_autospec(ToolContext, instance=True)
+
+  with mock.patch(
+      "google.adk.tools.bigtable.client.get_bigtable_data_client"
+  ) as mock_get_client:
+    mock_client = mock.MagicMock()
+    mock_get_client.return_value = mock_client
+    mock_iterator = mock.create_autospec(ExecuteQueryIterator, instance=True)
+    mock_client.execute_query.return_value = mock_iterator
+
+    # Mock row data
+    mock_row1 = mock.MagicMock()
+    mock_row1.fields = {"col1": "val1"}
+    mock_row2 = mock.MagicMock()
+    mock_row2.fields = {"col1": "val2"}
+    mock_iterator.__iter__.return_value = [mock_row1, mock_row2]
+
+    result = execute_sql(
+        project_id=project,
+        instance_id=instance_id,
+        credentials=credentials,
+        query=query,
+        settings=BigtableToolSettings(max_query_result_rows=1),
+        tool_context=tool_context,
+    )
+
+    expected_rows = [{"col1": "val1"}]
+    assert result == {
+        "status": "SUCCESS",
+        "rows": expected_rows,
+        "result_is_likely_truncated": True,
+    }
+    mock_client.execute_query.assert_called_once_with(
+        query=query, instance_id=instance_id
+    )
+    mock_iterator.close.assert_called_once()
+
+
+def test_execute_sql_error():
+  """Test execute_sql tool error handling."""
+  project = "my_project"
+  instance_id = "my_instance"
+  query = "SELECT * FROM my_table"
+  credentials = mock.create_autospec(Credentials, instance=True)
+  tool_context = mock.create_autospec(ToolContext, instance=True)
+
+  with mock.patch(
+      "google.adk.tools.bigtable.client.get_bigtable_data_client"
+  ) as mock_get_client:
+    mock_client = mock.MagicMock()
+    mock_get_client.return_value = mock_client
+    mock_client.execute_query.side_effect = Exception("Test error")
+
+    result = execute_sql(
+        project_id=project,
+        instance_id=instance_id,
+        credentials=credentials,
+        query=query,
+        settings=BigtableToolSettings(),
+        tool_context=tool_context,
+    )
+    assert result == {"status": "ERROR", "error_details": "Test error"}

tests/unittests/tools/bigtable/test_bigtable_toolset.py

@@ -0,0 +1,133 @@
+# 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 unittest import mock
+
+from google.adk.tools.bigtable import BigtableCredentialsConfig
+from google.adk.tools.bigtable import metadata_tool
+from google.adk.tools.bigtable import query_tool
+from google.adk.tools.bigtable.bigtable_toolset import BigtableToolset
+from google.adk.tools.bigtable.bigtable_toolset import DEFAULT_BIGTABLE_TOOL_NAME_PREFIX
+from google.adk.tools.google_tool import GoogleTool
+import pytest
+
+
+def test_bigtable_toolset_name_prefix():
+  """Test Bigtable toolset name prefix."""
+  credentials_config = BigtableCredentialsConfig(
+      client_id="abc", client_secret="def"
+  )
+  toolset = BigtableToolset(credentials_config=credentials_config)
+  assert toolset.tool_name_prefix == DEFAULT_BIGTABLE_TOOL_NAME_PREFIX
+
+
+@pytest.mark.asyncio
+async def test_bigtable_toolset_tools_default():
+  """Test default Bigtable toolset."""
+  credentials_config = BigtableCredentialsConfig(
+      client_id="abc", client_secret="def"
+  )
+  toolset = BigtableToolset(credentials_config=credentials_config)
+
+  tools = await toolset.get_tools()
+  assert tools is not None
+
+  assert len(tools) == 5
+  assert all([isinstance(tool, GoogleTool) for tool in tools])
+
+  expected_tool_names = set([
+      "list_instances",
+      "get_instance_info",
+      "list_tables",
+      "get_table_info",
+      "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_instances", "get_instance_info"], id="instance-metadata"
+        ),
+        pytest.param(["list_tables", "get_table_info"], id="table-metadata"),
+        pytest.param(["execute_sql"], id="query"),
+    ],
+)
+@pytest.mark.asyncio
+async def test_bigtable_toolset_tools_selective(selected_tools):
+  """Test Bigtable toolset with filter.
+
+  This test verifies the behavior of the Bigtable toolset when filter is
+  specified. A use case for this would be when the agent builder wants to
+  use only a subset of the tools provided by the toolset.
+  """
+  credentials_config = BigtableCredentialsConfig(
+      client_id="abc", client_secret="def"
+  )
+  toolset = BigtableToolset(
+      credentials_config=credentials_config, tool_filter=selected_tools
+  )
+
+  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_bigtable_toolset_unknown_tool(selected_tools, returned_tools):
+  """Test Bigtable toolset with filter.
+
+  This test verifies the behavior of the Bigtable toolset when filter is
+  specified with an unknown tool.
+  """
+  credentials_config = BigtableCredentialsConfig(
+      client_id="abc", client_secret="def"
+  )
+
+  toolset = BigtableToolset(
+      credentials_config=credentials_config, tool_filter=selected_tools
+  )
+
+  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

tests/unittests/tools/bigtable/test_client.py

@@ -0,0 +1,50 @@
+# 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 unittest import mock
+
+from google.adk.tools.bigtable import client
+from google.auth.credentials import Credentials
+
+
+def test_get_bigtable_data_client():
+  """Test get_bigtable_client function."""
+  with mock.patch(
+      "google.cloud.bigtable.data.BigtableDataClient"
+  ) as MockBigtableDataClient:
+    mock_creds = mock.create_autospec(Credentials, instance=True)
+    client.get_bigtable_data_client(
+        project="test-project", credentials=mock_creds
+    )
+    MockBigtableDataClient.assert_called_once_with(
+        project="test-project",
+        credentials=mock_creds,
+        client_info=mock.ANY,
+    )
+
+
+def test_get_bigtable_admin_client():
+  """Test get_bigtable_admin_client function."""
+  with mock.patch("google.cloud.bigtable.Client") as BigtableDataClient:
+    mock_creds = mock.create_autospec(Credentials, instance=True)
+    client.get_bigtable_admin_client(
+        project="test-project", credentials=mock_creds
+    )
+    # Admin client is a BigtableDataClient created with admin=True.
+    BigtableDataClient.assert_called_once_with(
+        project="test-project",
+        admin=True,
+        credentials=mock_creds,
+        client_info=mock.ANY,
+    )