encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-07-09 18:19:28 -07:00
Code Issues Packages Projects Releases Wiki Activity

feat: Introduce write protected mode to BigQuery tools

Browse Source 
This allows to protect against any write operations (e.g. update or delete a table), useful for some agents that must only be used in a read-only mode, while the user may have write permissions.

PiperOrigin-RevId: 769803741
This commit is contained in:
Google Team Member
2025-06-10 14:37:24 -07:00
committed by Copybara-Service
parent 77f44a4e45
commit 6c999caa41
10 changed files with 490 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unittests/tools/bigquery/test_bigquery_query_tool.py
+220
View File
@@ -0,0 +1,220 @@
# 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 textwrap
from typing import Optional
from google.adk.tools import BaseTool
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 pytest
async def get_tool(
name: str, tool_config: Optional[BigQueryToolConfig] = None
) -> BaseTool:
"""Get a tool from BigQuery toolset.
This method gets the tool view that an Agent using the BigQuery toolset would
see.
Returns:
The tool.
"""
credentials_config = BigQueryCredentialsConfig(
client_id="abc", client_secret="def"
)
toolset = BigQueryToolset(
credentials_config=credentials_config,
tool_filter=[name],
bigquery_tool_config=tool_config,
)
tools = await toolset.get_tools()
assert tools is not None
assert len(tools) == 1
return tools[0]
@pytest.mark.parametrize(
("tool_config",),
[
pytest.param(None, id="no-config"),
pytest.param(BigQueryToolConfig(), id="default-config"),
pytest.param(
BigQueryToolConfig(write_mode=WriteMode.BLOCKED),
id="explicit-no-write",
),
],
)
@pytest.mark.asyncio
async def test_execute_sql_declaration_read_only(tool_config):
"""Test BigQuery execute_sql tool declaration in read-only mode.
This test verifies that the execute_sql tool declaration reflects the
read-only capability.
"""
tool_name = "execute_sql"
tool = await get_tool(tool_name, tool_config)
assert tool.name == tool_name
assert tool.description == textwrap.dedent("""\
Run a BigQuery SQL query in the project and return the result.
Args:
project_id (str): The GCP project id in which the query should be
executed.
query (str): The BigQuery SQL query to be executed.
credentials (Credentials): The credentials to use for the request.
Returns:
dict: Dictionary representing the result of the query.
If the result contains the key "result_is_likely_truncated" with
value True, it means that there may be additional rows matching the
query not returned in the result.
Examples:
Fetch data or insights from a table:
>>> execute_sql("bigframes-dev",
... "SELECT island, COUNT(*) AS population "
... "FROM bigquery-public-data.ml_datasets.penguins GROUP BY island")
{
"status": "ERROR",
"rows": [
{
"island": "Dream",
"population": 124
},
{
"island": "Biscoe",
"population": 168
},
{
"island": "Torgersen",
"population": 52
}
]
}""")
@pytest.mark.parametrize(
("tool_config",),
[
pytest.param(
BigQueryToolConfig(write_mode=WriteMode.ALLOWED),
id="explicit-all-write",
),
],
)
@pytest.mark.asyncio
async def test_execute_sql_declaration_write(tool_config):
"""Test BigQuery execute_sql tool declaration with all writes enabled.
This test verifies that the execute_sql tool declaration reflects the write
capability.
"""
tool_name = "execute_sql"
tool = await get_tool(tool_name, tool_config)
assert tool.name == tool_name
assert tool.description == textwrap.dedent("""\
Run a BigQuery SQL query in the project and return the result.
Args:
project_id (str): The GCP project id in which the query should be
executed.
query (str): The BigQuery SQL query to be executed.
credentials (Credentials): The credentials to use for the request.
Returns:
dict: Dictionary representing the result of the query.
If the result contains the key "result_is_likely_truncated" with
value True, it means that there may be additional rows matching the
query not returned in the result.
Examples:
Fetch data or insights from a table:
>>> execute_sql("bigframes-dev",
... "SELECT island, COUNT(*) AS population "
... "FROM bigquery-public-data.ml_datasets.penguins GROUP BY island")
{
"status": "ERROR",
"rows": [
{
"island": "Dream",
"population": 124
},
{
"island": "Biscoe",
"population": 168
},
{
"island": "Torgersen",
"population": 52
}
]
}
Create a table from the result of a query:
>>> execute_sql("bigframes-dev",
... "CREATE TABLE my_project.my_dataset.my_table AS "
... "SELECT island, COUNT(*) AS population "
... "FROM bigquery-public-data.ml_datasets.penguins GROUP BY island")
{
"status": "SUCCESS",
"rows": []
}
Delete a table:
>>> execute_sql("bigframes-dev",
... "DROP TABLE my_project.my_dataset.my_table")
{
"status": "SUCCESS",
"rows": []
}
Copy a table to another table:
>>> execute_sql("bigframes-dev",
... "CREATE TABLE my_project.my_dataset.my_table_clone "
... "CLONE my_project.my_dataset.my_table")
{
"status": "SUCCESS",
"rows": []
}
Create a snapshot (a lightweight, read-optimized copy) of en existing
table:
>>> execute_sql("bigframes-dev",
... "CREATE SNAPSHOT TABLE my_project.my_dataset.my_table_snapshot "
... "CLONE my_project.my_dataset.my_table")
{
"status": "SUCCESS",
"rows": []
}
Notes:
- If a destination table already exists, there are a few ways to overwrite
it:
- Use "CREATE OR REPLACE TABLE" instead of "CREATE TABLE".
- First run "DROP TABLE", followed by "CREATE TABLE".
- To insert data into a table, use "INSERT INTO" statement.""")
tests/unittests/tools/bigquery/test_bigquery_tool.py
+25 -15
View File
@@ -92,11 +92,13 @@ class TestBigQueryTool:
The tool should properly inherit from FunctionTool while adding
Google API specific credential management capabilities.
"""
tool = BigQueryTool(func=sample_function, credentials=credentials_config)
tool = BigQueryTool(
func=sample_function, credentials_config=credentials_config
)
assert tool.func == sample_function
assert tool.credentials_manager is not None
assert isinstance(tool.credentials_manager, BigQueryCredentialsManager)
assert tool._credentials_manager is not None
assert isinstance(tool._credentials_manager, BigQueryCredentialsManager)
# Verify that 'credentials' parameter is ignored in function signature analysis
assert "credentials" in tool._ignore_params
@@ -106,10 +108,10 @@ class TestBigQueryTool:
Some tools might handle authentication externally or use service
accounts, so credential management should be optional.
"""
tool = BigQueryTool(func=sample_function, credentials=None)
tool = BigQueryTool(func=sample_function, credentials_config=None)
assert tool.func == sample_function
assert tool.credentials_manager is None
assert tool._credentials_manager is None
@pytest.mark.asyncio
async def test_run_async_with_valid_credentials(
@@ -120,12 +122,14 @@ class TestBigQueryTool:
This tests the main happy path where credentials are available
and the underlying function executes successfully.
"""
tool = BigQueryTool(func=sample_function, credentials=credentials_config)
tool = BigQueryTool(
func=sample_function, credentials_config=credentials_config
)
# Mock the credentials manager to return valid credentials
mock_creds = Mock(spec=Credentials)
with patch.object(
tool.credentials_manager,
tool._credentials_manager,
"get_valid_credentials",
return_value=mock_creds,
) as mock_get_creds:
@@ -147,11 +151,13 @@ class TestBigQueryTool:
When credentials aren't available and OAuth flow is needed,
the tool should return a user-friendly message rather than failing.
"""
tool = BigQueryTool(func=sample_function, credentials=credentials_config)
tool = BigQueryTool(
func=sample_function, credentials_config=credentials_config
)
# Mock credentials manager to return None (OAuth flow in progress)
with patch.object(
tool.credentials_manager, "get_valid_credentials", return_value=None
tool._credentials_manager, "get_valid_credentials", return_value=None
) as mock_get_creds:
result = await tool.run_async(
@@ -171,7 +177,7 @@ class TestBigQueryTool:
Tools without credential managers should execute normally,
passing None for credentials if the function accepts them.
"""
tool = BigQueryTool(func=sample_function, credentials=None)
tool = BigQueryTool(func=sample_function, credentials_config=None)
result = await tool.run_async(
args={"param1": "test_value"}, tool_context=mock_tool_context
@@ -190,12 +196,12 @@ class TestBigQueryTool:
which is important for tools that make async API calls.
"""
tool = BigQueryTool(
func=async_sample_function, credentials=credentials_config
func=async_sample_function, credentials_config=credentials_config
)
mock_creds = Mock(spec=Credentials)
with patch.object(
tool.credentials_manager,
tool._credentials_manager,
"get_valid_credentials",
return_value=mock_creds,
):
@@ -220,11 +226,13 @@ class TestBigQueryTool:
def failing_function(param1: str, credentials: Credentials = None) -> dict:
raise ValueError("Something went wrong")
tool = BigQueryTool(func=failing_function, credentials=credentials_config)
tool = BigQueryTool(
func=failing_function, credentials_config=credentials_config
)
mock_creds = Mock(spec=Credentials)
with patch.object(
tool.credentials_manager,
tool._credentials_manager,
"get_valid_credentials",
return_value=mock_creds,
):
@@ -250,7 +258,9 @@ class TestBigQueryTool:
) -> dict:
return {"success": True}
tool = BigQueryTool(func=complex_function, credentials=credentials_config)
tool = BigQueryTool(
func=complex_function, credentials_config=credentials_config
)
# The 'credentials' parameter should be ignored in mandatory args analysis
mandatory_args = tool._get_mandatory_args()
tests/unittests/tools/bigquery/test_bigquery_tool_config.py
+27
View File
@@ -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.bigquery.config import BigQueryToolConfig
import pytest
def test_bigquery_tool_config_experimental_warning():
"""Test BigQueryToolConfig experimental warning."""
with pytest.warns(
UserWarning,
match="Config defaults may have breaking change in the future.",
):
BigQueryToolConfig()
tests/unittests/tools/bigquery_tool/test_bigquery_toolset.py → tests/unittests/tools/bigquery/test_bigquery_toolset.py
View File