feat: Introduce write protected mode to BigQuery tools

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

Author: google-genai-bot

contributing/samples/bigquery/agent.py

@@ -17,11 +17,15 @@
 from google.adk.agents import llm_agent
 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 google.auth
 
 RUN_WITH_ADC = False
 
 
+tool_config = BigQueryToolConfig(write_mode=WriteMode.ALLOWED)
+
 if RUN_WITH_ADC:
   # Initialize the tools to use the application default credentials.
   application_default_credentials, _ = google.auth.default()
@@ -37,7 +41,9 @@
       client_secret=os.getenv("OAUTH_CLIENT_SECRET"),
   )
 
-bigquery_toolset = BigQueryToolset(credentials_config=credentials_config)
+bigquery_toolset = BigQueryToolset(
+    credentials_config=credentials_config, bigquery_tool_config=tool_config
+)
 
 # The variable name `root_agent` determines what your root agent is for the
 # debug CLI

src/google/adk/tools/bigquery/bigquery_tool.py

@@ -25,6 +25,7 @@
 from ..tool_context import ToolContext
 from .bigquery_credentials import BigQueryCredentialsConfig
 from .bigquery_credentials import BigQueryCredentialsManager
+from .config import BigQueryToolConfig
 
 
 class BigQueryTool(FunctionTool):
@@ -41,21 +42,27 @@ class BigQueryTool(FunctionTool):
   def __init__(
       self,
       func: Callable[..., Any],
-      credentials: Optional[BigQueryCredentialsConfig] = None,
+      *,
+      credentials_config: Optional[BigQueryCredentialsConfig] = None,
+      bigquery_tool_config: Optional[BigQueryToolConfig] = None,
   ):
     """Initialize the Google API tool.
 
     Args:
         func: callable that impelments the tool's logic, can accept one
           'credential" parameter
-        credentials: credentials used to call Google API. If None, then we don't
-          hanlde the auth logic
+        credentials_config: credentials config used to call Google API. If None,
+          then we don't hanlde the auth logic
     """
     super().__init__(func=func)
     self._ignore_params.append("credentials")
-    self.credentials_manager = (
-        BigQueryCredentialsManager(credentials) if credentials else None
+    self._ignore_params.append("config")
+    self._credentials_manager = (
+        BigQueryCredentialsManager(credentials_config)
+        if credentials_config
+        else None
     )
+    self._tool_config = bigquery_tool_config
 
   @override
   async def run_async(
@@ -69,12 +76,12 @@ async def run_async(
     try:
       # Get valid credentials
       credentials = (
-          await self.credentials_manager.get_valid_credentials(tool_context)
-          if self.credentials_manager
+          await self._credentials_manager.get_valid_credentials(tool_context)
+          if self._credentials_manager
           else None
       )
 
-      if credentials is None and self.credentials_manager:
+      if credentials is None and self._credentials_manager:
         # OAuth flow in progress
         return (
             "User authorization is required to access Google services for"
@@ -84,7 +91,7 @@ async def run_async(
       # Execute the tool's specific logic with valid credentials
 
       return await self._run_async_with_credential(
-          credentials, args, tool_context
+          credentials, self._tool_config, args, tool_context
       )
 
     except Exception as ex:
@@ -96,6 +103,7 @@ async def run_async(
   async def _run_async_with_credential(
       self,
       credentials: Credentials,
+      tool_config: BigQueryToolConfig,
       args: dict[str, Any],
       tool_context: ToolContext,
   ) -> Any:
@@ -113,4 +121,6 @@ async def _run_async_with_credential(
     signature = inspect.signature(self.func)
     if "credentials" in signature.parameters:
       args_to_call["credentials"] = credentials
+    if "config" in signature.parameters:
+      args_to_call["config"] = tool_config
     return await super().run_async(args=args_to_call, tool_context=tool_context)

src/google/adk/tools/bigquery/bigquery_toolset.py

@@ -28,6 +28,7 @@
 from ...tools.base_toolset import ToolPredicate
 from .bigquery_credentials import BigQueryCredentialsConfig
 from .bigquery_tool import BigQueryTool
+from .config import BigQueryToolConfig
 
 
 class BigQueryToolset(BaseToolset):
@@ -38,9 +39,11 @@ def __init__(
       *,
       tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
       credentials_config: Optional[BigQueryCredentialsConfig] = None,
+      bigquery_tool_config: Optional[BigQueryToolConfig] = None,
   ):
-    self._credentials_config = credentials_config
     self.tool_filter = tool_filter
+    self._credentials_config = credentials_config
+    self._tool_config = bigquery_tool_config
 
   def _is_tool_selected(
       self, tool: BaseTool, readonly_context: ReadonlyContext
@@ -64,14 +67,15 @@ async def get_tools(
     all_tools = [
         BigQueryTool(
             func=func,
-            credentials=self._credentials_config,
+            credentials_config=self._credentials_config,
+            bigquery_tool_config=self._tool_config,
         )
         for func in [
             metadata_tool.get_dataset_info,
             metadata_tool.get_table_info,
             metadata_tool.list_dataset_ids,
             metadata_tool.list_table_ids,
-            query_tool.execute_sql,
+            query_tool.get_execute_sql(self._tool_config),
         ]
     ]
 

src/google/adk/tools/bigquery/config.py

@@ -0,0 +1,46 @@
+# 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 enum import Enum
+
+from pydantic import BaseModel
+
+from ...utils.feature_decorator import experimental
+
+
+class WriteMode(Enum):
+  """Write mode indicating what levels of write operations are allowed in BigQuery."""
+
+  BLOCKED = 'blocked'
+  """No write operations are allowed.
+  
+  This mode implies that only read (i.e. SELECT query) operations are allowed.
+  """
+
+  ALLOWED = 'allowed'
+  """All write operations are allowed."""
+
+
+@experimental('Config defaults may have breaking change in the future.')
+class BigQueryToolConfig(BaseModel):
+  """Configuration for BigQuery tools."""
+
+  write_mode: WriteMode = WriteMode.BLOCKED
+  """Write mode for BigQuery tools.
+
+  By default, the tool will allow only read operations. This behaviour may
+  change in future versions.
+  """

src/google/adk/tools/bigquery/metadata_tool.py

@@ -15,7 +15,7 @@
 from google.cloud import bigquery
 from google.oauth2.credentials import Credentials
 
-from ...tools.bigquery import client
+from . import client
 
 
 def list_dataset_ids(project_id: str, credentials: Credentials) -> list[str]:

src/google/adk/tools/bigquery/query_tool.py

@@ -12,14 +12,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import functools
+import types
+from typing import Callable
+
+from google.cloud import bigquery
 from google.oauth2.credentials import Credentials
 
-from ...tools.bigquery import client
+from . import client
+from .config import BigQueryToolConfig
+from .config import WriteMode
 
 MAX_DOWNLOADED_QUERY_RESULT_ROWS = 50
 
 
-def execute_sql(project_id: str, query: str, credentials: Credentials) -> dict:
+def execute_sql(
+    project_id: str,
+    query: str,
+    credentials: Credentials,
+    config: BigQueryToolConfig,
+) -> dict:
   """Run a BigQuery SQL query in the project and return the result.
 
   Args:
@@ -35,34 +47,49 @@ def execute_sql(project_id: str, query: str, credentials: Credentials) -> dict:
             query not returned in the result.
 
   Examples:
-      >>> execute_sql("bigframes-dev",
-      ... "SELECT island, COUNT(*) AS population "
-      ... "FROM bigquery-public-data.ml_datasets.penguins GROUP BY island")
-      {
-        "rows": [
-            {
-                "island": "Dream",
-                "population": 124
-            },
-            {
-                "island": "Biscoe",
-                "population": 168
-            },
-            {
-                "island": "Torgersen",
-                "population": 52
-            }
-        ]
-      }
+      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
+                }
+            ]
+          }
   """
 
   try:
     bq_client = client.get_bigquery_client(credentials=credentials)
+    if not config or config.write_mode == WriteMode.BLOCKED:
+      query_job = bq_client.query(
+          query,
+          project=project_id,
+          job_config=bigquery.QueryJobConfig(dry_run=True),
+      )
+      if query_job.statement_type != "SELECT":
+        return {
+            "status": "ERROR",
+            "error_details": "Read-only mode only supports SELECT statements.",
+        }
+
     row_iterator = bq_client.query_and_wait(
         query, project=project_id, max_results=MAX_DOWNLOADED_QUERY_RESULT_ROWS
     )
     rows = [{key: val for key, val in row.items()} for row in row_iterator]
-    result = {"rows": rows}
+    result = {"status": "SUCCESS", "rows": rows}
     if (
         MAX_DOWNLOADED_QUERY_RESULT_ROWS is not None
         and len(rows) == MAX_DOWNLOADED_QUERY_RESULT_ROWS
@@ -74,3 +101,92 @@ def execute_sql(project_id: str, query: str, credentials: Credentials) -> dict:
         "status": "ERROR",
         "error_details": str(ex),
     }
+
+
+_execute_sql_write_examples = """
+      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.
+  """
+
+
+def get_execute_sql(config: BigQueryToolConfig) -> Callable[..., dict]:
+  """Get the execute_sql tool customized as per the given tool config.
+
+  Args:
+      config: BigQuery tool configuration indicating the behavior of the
+        execute_sql tool.
+
+  Returns:
+      callable[..., dict]: A version of the execute_sql tool respecting the tool
+      config.
+  """
+
+  if not config or config.write_mode == WriteMode.BLOCKED:
+    return execute_sql
+
+  # Create a new function object using the original function's code and globals.
+  # We pass the original code, globals, name, defaults, and closure.
+  # This creates a raw function object without copying other metadata yet.
+  execute_sql_wrapper = types.FunctionType(
+      execute_sql.__code__,
+      execute_sql.__globals__,
+      execute_sql.__name__,
+      execute_sql.__defaults__,
+      execute_sql.__closure__,
+  )
+
+  # Use functools.update_wrapper to copy over other essential attributes
+  # from the original function to the new one.
+  # This includes __name__, __qualname__, __module__, __annotations__, etc.
+  # It specifically allows us to then set __doc__ separately.
+  functools.update_wrapper(execute_sql_wrapper, execute_sql)
+
+  # Now, set the new docstring
+  execute_sql_wrapper.__doc__ += _execute_sql_write_examples
+
+  return execute_sql_wrapper