From 4bbd78370562cf6e3bee66bc4508cc0a6d12dfa4 Mon Sep 17 00:00:00 2001 From: Anwesha Das Date: Mon, 15 Jun 2026 13:24:28 -0400 Subject: [PATCH 1/2] feat(bigtable): Support parameterized views with secure parameter injection Expose a parameterized query tool execute_sql_parameterized that automatically maps and injects secure parameters (like user_id) from the tool context to Bigtable's view_parameters. --- pyproject.toml | 6 +- .../adk/tools/bigtable/bigtable_toolset.py | 86 +++++++++++++ src/google/adk/tools/bigtable/query_tool.py | 3 + .../bigtable/test_bigtable_query_tool.py | 38 ++++++ .../tools/bigtable/test_bigtable_toolset.py | 116 +++++++++++++++++- 5 files changed, 245 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index fe33e65229..1f3be4a2fb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,7 +73,7 @@ optional-dependencies.all = [ "google-cloud-aiplatform[agent-engines]>=1.148.1,<2", "google-cloud-bigquery>=2.2", "google-cloud-bigquery-storage>=2", - "google-cloud-bigtable>=2.32", + "google-cloud-bigtable>=2.38.0", "google-cloud-dataplex>=1.7,<3", "google-cloud-discoveryengine>=0.13.12,<0.14", "google-cloud-pubsub>=2,<3", @@ -153,7 +153,7 @@ optional-dependencies.gcp = [ "google-cloud-aiplatform[agent-engines]>=1.148.1,<2", "google-cloud-bigquery>=2.2", "google-cloud-bigquery-storage>=2", - "google-cloud-bigtable>=2.32", + "google-cloud-bigtable>=2.38.0", "google-cloud-dataplex>=1.7,<3", "google-cloud-discoveryengine>=0.13.12,<0.14", "google-cloud-pubsub>=2,<3", @@ -192,7 +192,7 @@ optional-dependencies.test = [ "google-cloud-aiplatform[agent-engines,evaluation]>=1.148.1,<2", "google-cloud-bigquery>=2.2", "google-cloud-bigquery-storage>=2", - "google-cloud-bigtable>=2.32", + "google-cloud-bigtable>=2.38.0", "google-cloud-dataplex>=1.7,<3", "google-cloud-discoveryengine>=0.13.12,<0.14", "google-cloud-firestore>=2.11,<3", diff --git a/src/google/adk/tools/bigtable/bigtable_toolset.py b/src/google/adk/tools/bigtable/bigtable_toolset.py index 97fc2eb0b6..c06edafb1c 100644 --- a/src/google/adk/tools/bigtable/bigtable_toolset.py +++ b/src/google/adk/tools/bigtable/bigtable_toolset.py @@ -14,11 +14,16 @@ from __future__ import annotations +import inspect +from typing import Any +from typing import Callable from typing import List from typing import Optional from typing import Union from google.adk.agents.readonly_context import ReadonlyContext +from google.auth.credentials import Credentials +from pydantic import BaseModel from typing_extensions import override from . import metadata_tool @@ -29,9 +34,81 @@ from ...tools.base_toolset import BaseToolset from ...tools.base_toolset import ToolPredicate from ...tools.google_tool import GoogleTool +from ..tool_context import ToolContext from .bigtable_credentials import BigtableCredentialsConfig from .settings import BigtableToolSettings + +class BigtableParameterizedViewTool(GoogleTool): + """Wrapper FunctionTool for Bigtable execute_sql query tool that passes view parameters. + + This tool wraps the Bigtable query tool to automatically resolve and inject a + parameter from the ToolContext (e.g. user_id) into the query's + view_parameters. The parameter name to resolve is configured via + view_parameter_name. + + Example: + If a parameterized view `purchase_history_pv` was created with the query: + `SELECT * FROM purchases WHERE user_id = VIEW_PARAMETERS('user_id')` + + By configuring `view_parameter_name="user_id"`, the wrapper will resolve the + `user_id` value from the `tool_context.user_id` at runtime and pass it as + `view_parameters={"user_id": user_id}`. This securely restricts query execution to the + logged-in user's data without exposing the `user_id` parameter to the LLM. + """ + + def __init__( + self, + func: Callable[..., Any], + *, + credentials_config: Optional[BigtableCredentialsConfig] = None, + tool_settings: Optional[BigtableToolSettings] = None, + view_parameter_name: Optional[str] = None, + ): + """Initializes the BigtableParameterizedViewTool. + + Args: + func: The Bigtable query function to wrap. + credentials_config: The credentials configuration. + tool_settings: The tool settings. + view_parameter_name: The name of the parameter to resolve from + tool_context and pass into view_parameters. This is typically + configured on the toolset (BigtableToolset) and forwarded here. + """ + super().__init__( + func=func, + credentials_config=credentials_config, + tool_settings=tool_settings, + ) + self.name = "execute_sql_parameterized" + self.description = ( + "Execute a GoogleSQL query from a Bigtable table using parameterized views " + "to securely check permissions." + ) + self._view_parameter_name = view_parameter_name + # Exclude from being parsed and exposed to the LLM when generating tool schemas + self._ignore_params.append("view_parameters") + + @override + async def _run_async_with_credential( + self, + credentials: Credentials, + tool_settings: BaseModel, + args: dict[str, Any], + tool_context: ToolContext, + ) -> Any: + args_to_call = args.copy() + signature = inspect.signature(self.func) + if "view_parameters" in signature.parameters: + view_params = {} + if self._view_parameter_name and hasattr(tool_context, self._view_parameter_name): + view_params[self._view_parameter_name] = getattr(tool_context, self._view_parameter_name) + args_to_call["view_parameters"] = view_params + return await super()._run_async_with_credential( + credentials, tool_settings, args_to_call, tool_context + ) + + DEFAULT_BIGTABLE_TOOL_NAME_PREFIX = "bigtable" @@ -55,6 +132,7 @@ def __init__( tool_filter: Optional[Union[ToolPredicate, List[str]]] = None, credentials_config: Optional[BigtableCredentialsConfig] = None, bigtable_tool_settings: Optional[BigtableToolSettings] = None, + view_parameter_name: Optional[str] = None, ): super().__init__( tool_filter=tool_filter, @@ -66,6 +144,7 @@ def __init__( if bigtable_tool_settings else BigtableToolSettings() ) + self._view_parameter_name = view_parameter_name def _is_tool_selected( self, tool: BaseTool, readonly_context: ReadonlyContext @@ -101,6 +180,13 @@ async def get_tools( metadata_tool.get_cluster_info, query_tool.execute_sql, ] + ] + [ + BigtableParameterizedViewTool( + func=query_tool.execute_sql, + credentials_config=self._credentials_config, + tool_settings=self._tool_settings, + view_parameter_name=self._view_parameter_name, + ) ] return [ tool diff --git a/src/google/adk/tools/bigtable/query_tool.py b/src/google/adk/tools/bigtable/query_tool.py index bf64b282a1..c810c62cc1 100644 --- a/src/google/adk/tools/bigtable/query_tool.py +++ b/src/google/adk/tools/bigtable/query_tool.py @@ -42,6 +42,7 @@ async def execute_sql( tool_context: ToolContext, parameters: Dict[str, Any] | None = None, parameter_types: Dict[str, Any] | None = None, + view_parameters: Dict[str, Any] | None = None, ) -> dict: """Execute a GoogleSQL query from a Bigtable table. @@ -56,6 +57,7 @@ async def execute_sql( parameters (dict): properties for parameter replacement. Keys must match the names used in ``query``. parameter_types (dict): maps explicit types for one or more param values. + view_parameters (dict): maps properties for parameterized authorized views. Returns: dict: Dictionary containing the status and the rows read. @@ -91,6 +93,7 @@ def _execute_sql(): instance_id=instance_id, parameters=parameters, parameter_types=parameter_types, + view_parameters=view_parameters, ) rows: List[Dict[str, Any]] = [] diff --git a/tests/unittests/tools/bigtable/test_bigtable_query_tool.py b/tests/unittests/tools/bigtable/test_bigtable_query_tool.py index 46b65a3a07..da21bc58f9 100644 --- a/tests/unittests/tools/bigtable/test_bigtable_query_tool.py +++ b/tests/unittests/tools/bigtable/test_bigtable_query_tool.py @@ -193,6 +193,7 @@ def raise_error(): instance_id=instance_id, parameters=parameters, parameter_types=parameter_types, + view_parameters=None, ) mock_iterator.close.assert_called_once() @@ -228,3 +229,40 @@ async def test_execute_sql_row_value_circular_reference_fallback(): assert result["status"] == "SUCCESS" assert result["rows"][0]["col1"] == str(circular_value) + + +@pytest.mark.asyncio +async def test_execute_sql_with_view_parameters(): + """Test execute_sql with view_parameters passed.""" + 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) + view_parameters = {"user_id": "test-user-123"} + + with mock.patch.object(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_iterator.__iter__.return_value = [] + + result = await execute_sql( + project_id=project, + instance_id=instance_id, + credentials=credentials, + query=query, + settings=BigtableToolSettings(), + tool_context=tool_context, + view_parameters=view_parameters, + ) + + assert result["status"] == "SUCCESS" + mock_client.execute_query.assert_called_once_with( + query=query, + instance_id=instance_id, + parameters=None, + parameter_types=None, + view_parameters=view_parameters, + ) diff --git a/tests/unittests/tools/bigtable/test_bigtable_toolset.py b/tests/unittests/tools/bigtable/test_bigtable_toolset.py index b5698cfc07..d6f8452358 100644 --- a/tests/unittests/tools/bigtable/test_bigtable_toolset.py +++ b/tests/unittests/tools/bigtable/test_bigtable_toolset.py @@ -19,6 +19,7 @@ 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 BigtableParameterizedViewTool 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 @@ -45,7 +46,7 @@ async def test_bigtable_toolset_tools_default(): tools = await toolset.get_tools() assert tools is not None - assert len(tools) == 7 + assert len(tools) == 8 assert all([isinstance(tool, GoogleTool) for tool in tools]) expected_tool_names = set([ @@ -54,6 +55,7 @@ async def test_bigtable_toolset_tools_default(): "list_tables", "get_table_info", "execute_sql", + "execute_sql_parameterized", "list_clusters", "get_cluster_info", ]) @@ -133,3 +135,115 @@ async def test_bigtable_toolset_unknown_tool(selected_tools, returned_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.asyncio +async def test_bigtable_toolset_query_tool_wrapped(): + """Test that execute_sql is wrapped in BigtableQueryTool.""" + credentials_config = BigtableCredentialsConfig( + client_id="abc", client_secret="def" + ) + toolset = BigtableToolset(credentials_config=credentials_config) + + tools = await toolset.get_tools() + query_tools = [tool for tool in tools if tool.name == "execute_sql"] + assert len(query_tools) == 1 + + parameterized_tools = [tool for tool in tools if tool.name == "execute_sql_parameterized"] + assert len(parameterized_tools) == 1 + tool = parameterized_tools[0] + assert isinstance(tool, BigtableParameterizedViewTool) + assert "view_parameters" in tool._ignore_params + assert tool._view_parameter_name is None + + +@pytest.mark.asyncio +async def test_bigtable_toolset_query_tool_wrapped_custom_mapping(): + """Test that BigtableParameterizedViewTool accepts custom mapping.""" + credentials_config = BigtableCredentialsConfig( + client_id="abc", client_secret="def" + ) + toolset = BigtableToolset( + credentials_config=credentials_config, + view_parameter_name="user_id", + ) + + tools = await toolset.get_tools() + parameterized_tools = [tool for tool in tools if tool.name == "execute_sql_parameterized"] + assert len(parameterized_tools) == 1 + tool = parameterized_tools[0] + assert isinstance(tool, BigtableParameterizedViewTool) + assert tool._view_parameter_name == "user_id" + + +@pytest.mark.asyncio +async def test_bigtable_parameterized_view_tool_execution(): + """Test that BigtableParameterizedViewTool maps attributes from tool_context to view_parameters.""" + # Define a dummy function to wrap that has 'view_parameters' parameter + def mock_execute_sql(view_parameters=None): + return {"status": "SUCCESS", "view_parameters": view_parameters} + + credentials = mock.create_autospec(Credentials, instance=True) + tool_settings = BigtableToolSettings() + tool_context = mock.create_autospec(ToolContext, instance=True) + tool_context.user_id = "test-user-123" + + # Create tool with custom mapping + tool = BigtableParameterizedViewTool( + func=mock_execute_sql, + view_parameter_name="user_id", + ) + + # Run the tool + res = await tool._run_async_with_credential( + credentials=credentials, + tool_settings=tool_settings, + args={}, + tool_context=tool_context, + ) + + assert res == { + "status": "SUCCESS", + "view_parameters": {"user_id": "test-user-123"}, + } + + +@pytest.mark.asyncio +async def test_bigtable_parameterized_view_tool_login_flow(): + """Test showing how user_id is updated in the session/context after login.""" + def mock_execute_sql(view_parameters=None): + return {"status": "SUCCESS", "view_parameters": view_parameters} + + from google.adk.sessions.session import Session + from google.adk.agents.invocation_context import InvocationContext + from google.adk.agents.context import Context + + session = Session(id="session-1", app_name="test-app", user_id="anonymous") + + invocation_context = mock.create_autospec(InvocationContext, instance=True) + invocation_context.session = session + type(invocation_context).user_id = property(lambda self: self.session.user_id) + + tool_context = Context(invocation_context=invocation_context) + assert tool_context.user_id == "anonymous" + + # Simulate login by updating user_id on session + session.user_id = "authenticated-user-999" + + credentials = mock.create_autospec(Credentials, instance=True) + tool = BigtableParameterizedViewTool( + func=mock_execute_sql, + view_parameter_name="user_id", + ) + + res = await tool._run_async_with_credential( + credentials=credentials, + tool_settings=BigtableToolSettings(), + args={}, + tool_context=tool_context, + ) + + assert res == { + "status": "SUCCESS", + "view_parameters": {"user_id": "authenticated-user-999"}, + } From c123232e65176767b723b647355cb21647b61997 Mon Sep 17 00:00:00 2001 From: Anwesha Das Date: Mon, 15 Jun 2026 14:53:06 -0400 Subject: [PATCH 2/2] feat(bigtable): Fallback to session state when injecting view parameters If the configured view_parameter_name is not a top-level property on the ToolContext (like user_id), check tool_context.state to support arbitrary application-specific session parameters (e.g. tenant_id). --- .../adk/tools/bigtable/bigtable_toolset.py | 18 +++++-- .../tools/bigtable/test_bigtable_toolset.py | 47 +++++++++++++++++++ 2 files changed, 62 insertions(+), 3 deletions(-) diff --git a/src/google/adk/tools/bigtable/bigtable_toolset.py b/src/google/adk/tools/bigtable/bigtable_toolset.py index c06edafb1c..deadf64e3c 100644 --- a/src/google/adk/tools/bigtable/bigtable_toolset.py +++ b/src/google/adk/tools/bigtable/bigtable_toolset.py @@ -99,10 +99,22 @@ async def _run_async_with_credential( ) -> Any: args_to_call = args.copy() signature = inspect.signature(self.func) - if "view_parameters" in signature.parameters: + if "view_parameters" in signature.parameters and self._view_parameter_name: view_params = {} - if self._view_parameter_name and hasattr(tool_context, self._view_parameter_name): - view_params[self._view_parameter_name] = getattr(tool_context, self._view_parameter_name) + # 1. Check if it's a strongly-typed top-level property (like 'user_id') + if hasattr(tool_context, self._view_parameter_name): + view_params[self._view_parameter_name] = getattr( + tool_context, self._view_parameter_name + ) + # 2. Fallback to checking application-level session state + elif ( + tool_context.state + and self._view_parameter_name in tool_context.state + ): + view_params[self._view_parameter_name] = tool_context.state[ + self._view_parameter_name + ] + args_to_call["view_parameters"] = view_params return await super()._run_async_with_credential( credentials, tool_settings, args_to_call, tool_context diff --git a/tests/unittests/tools/bigtable/test_bigtable_toolset.py b/tests/unittests/tools/bigtable/test_bigtable_toolset.py index d6f8452358..bdda9a9aa8 100644 --- a/tests/unittests/tools/bigtable/test_bigtable_toolset.py +++ b/tests/unittests/tools/bigtable/test_bigtable_toolset.py @@ -247,3 +247,50 @@ def mock_execute_sql(view_parameters=None): "status": "SUCCESS", "view_parameters": {"user_id": "authenticated-user-999"}, } + + +@pytest.mark.asyncio +async def test_bigtable_parameterized_view_tool_execution_session_state_fallback(): + """Test that BigtableParameterizedViewTool falls back to tool_context.state for custom parameters.""" + def mock_execute_sql(view_parameters=None): + return {"status": "SUCCESS", "view_parameters": view_parameters} + + from google.adk.sessions.session import Session + from google.adk.agents.invocation_context import InvocationContext + from google.adk.agents.context import Context + + # Create session with application-level state + session = Session( + id="session-1", + app_name="test-app", + user_id="user-123", + state={"tenant_id": "tenant-xyz"}, + ) + + invocation_context = mock.create_autospec(InvocationContext, instance=True) + invocation_context.session = session + type(invocation_context).user_id = property(lambda self: self.session.user_id) + + tool_context = Context(invocation_context=invocation_context) + + # Ensure 'tenant_id' is NOT a top-level property or attribute on tool_context + assert not hasattr(tool_context, "tenant_id") + assert "tenant_id" in tool_context.state + + credentials = mock.create_autospec(Credentials, instance=True) + tool = BigtableParameterizedViewTool( + func=mock_execute_sql, + view_parameter_name="tenant_id", + ) + + res = await tool._run_async_with_credential( + credentials=credentials, + tool_settings=BigtableToolSettings(), + args={}, + tool_context=tool_context, + ) + + assert res == { + "status": "SUCCESS", + "view_parameters": {"tenant_id": "tenant-xyz"}, + }