google-adk 1.2.0__py3-none-any.whl → 1.4.0__py3-none-any.whl

google/adk/tools/bigquery/bigquery_credentials.py

@@ -27,16 +27,18 @@ from google.oauth2.credentials import Credentials
 from pydantic import BaseModel
 from pydantic import model_validator
 
-from ...auth import AuthConfig
-from ...auth import AuthCredential
-from ...auth import AuthCredentialTypes
-from ...auth import OAuth2Auth
+from ...auth.auth_credential import AuthCredential
+from ...auth.auth_credential import AuthCredentialTypes
+from ...auth.auth_credential import OAuth2Auth
+from ...auth.auth_tool import AuthConfig
+from ...utils.feature_decorator import experimental
 from ..tool_context import ToolContext
 
 BIGQUERY_TOKEN_CACHE_KEY = "bigquery_token_cache"
 BIGQUERY_DEFAULT_SCOPE = ["https://www.googleapis.com/auth/bigquery"]
 
 
+@experimental
 class BigQueryCredentialsConfig(BaseModel):
   """Configuration for Google API tools. (Experimental)"""
 

google/adk/tools/bigquery/bigquery_tool.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
 
 import inspect
 from typing import Any
@@ -21,12 +22,15 @@ from typing import Optional
 from google.oauth2.credentials import Credentials
 from typing_extensions import override
 
+from ...utils.feature_decorator import experimental
 from ..function_tool import FunctionTool
 from ..tool_context import ToolContext
 from .bigquery_credentials import BigQueryCredentialsConfig
 from .bigquery_credentials import BigQueryCredentialsManager
+from .config import BigQueryToolConfig
 
 
+@experimental
 class BigQueryTool(FunctionTool):
   """GoogleApiTool class for tools that call Google APIs.
 
@@ -41,21 +45,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 +79,12 @@ class BigQueryTool(FunctionTool):
     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 +94,7 @@ class BigQueryTool(FunctionTool):
       # 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 +106,7 @@ class BigQueryTool(FunctionTool):
   async def _run_async_with_credential(
       self,
       credentials: Credentials,
+      tool_config: BigQueryToolConfig,
       args: dict[str, Any],
       tool_context: ToolContext,
   ) -> Any:
@@ -113,4 +124,6 @@ class BigQueryTool(FunctionTool):
     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)

google/adk/tools/bigquery/bigquery_toolset.py

@@ -26,10 +26,13 @@ from . import query_tool
 from ...tools.base_tool import BaseTool
 from ...tools.base_toolset import BaseToolset
 from ...tools.base_toolset import ToolPredicate
+from ...utils.feature_decorator import experimental
 from .bigquery_credentials import BigQueryCredentialsConfig
 from .bigquery_tool import BigQueryTool
+from .config import BigQueryToolConfig
 
 
+@experimental
 class BigQueryToolset(BaseToolset):
   """BigQuery Toolset contains tools for interacting with BigQuery data and metadata."""
 
@@ -38,9 +41,11 @@ class BigQueryToolset(BaseToolset):
       *,
       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 +69,15 @@ class BigQueryToolset(BaseToolset):
     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),
         ]
     ]
 

google/adk/tools/bigquery/client.py

@@ -18,16 +18,20 @@ import google.api_core.client_info
 from google.cloud import bigquery
 from google.oauth2.credentials import Credentials
 
-USER_AGENT = "adk-bigquery-tool"
+from ... import version
 
+USER_AGENT = f"adk-bigquery-tool google-adk/{version.__version__}"
 
-def get_bigquery_client(*, credentials: Credentials) -> bigquery.Client:
+
+def get_bigquery_client(
+    *, project: str, credentials: Credentials
+) -> bigquery.Client:
   """Get a BigQuery client."""
 
   client_info = google.api_core.client_info.ClientInfo(user_agent=USER_AGENT)
 
   bigquery_client = bigquery.Client(
-      credentials=credentials, client_info=client_info
+      project=project, credentials=credentials, client_info=client_info
   )
 
   return bigquery_client

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.
+  """

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]:
@@ -42,7 +42,9 @@ def list_dataset_ids(project_id: str, credentials: Credentials) -> list[str]:
        'bbc_news']
   """
   try:
-    bq_client = client.get_bigquery_client(credentials=credentials)
+    bq_client = client.get_bigquery_client(
+        project=project_id, credentials=credentials
+    )
 
     datasets = []
     for dataset in bq_client.list_datasets(project_id):
@@ -69,40 +71,46 @@ def get_dataset_info(
       dict: Dictionary representing the properties of the dataset.
 
   Examples:
-      >>> get_dataset_info("bigquery-public-data", "penguins")
+      >>> get_dataset_info("bigquery-public-data", "cdc_places")
       {
         "kind": "bigquery#dataset",
-        "etag": "PNC5907iQbzeVcAru/2L3A==",
-        "id": "bigquery-public-data:ml_datasets",
-        "selfLink":
-          "https://bigquery.googleapis.com/bigquery/v2/projects/bigquery-public-data/datasets/ml_datasets",
+        "etag": "fz9BaiXKgbGi53EpI2rJug==",
+        "id": "bigquery-public-data:cdc_places",
+        "selfLink": "https://content-bigquery.googleapis.com/bigquery/v2/projects/bigquery-public-data/datasets/cdc_places",
         "datasetReference": {
-            "datasetId": "ml_datasets",
-            "projectId": "bigquery-public-data"
+          "datasetId": "cdc_places",
+          "projectId": "bigquery-public-data"
         },
+        "description": "Local Data for Better Health, County Data",
         "access": [
-            {
-                "role": "OWNER",
-                "groupByEmail": "cloud-datasets-eng@google.com"
-            },
-            {
-                "role": "READER",
-                "iamMember": "allUsers"
-            },
-            {
-                "role": "READER",
-                "groupByEmail": "bqml-eng@google.com"
-            }
+          {
+            "role": "WRITER",
+            "specialGroup": "projectWriters"
+          },
+          {
+            "role": "OWNER",
+            "specialGroup": "projectOwners"
+          },
+          {
+            "role": "OWNER",
+            "userByEmail": "some-redacted-email@bigquery-public-data.iam.gserviceaccount.com"
+          },
+          {
+            "role": "READER",
+            "specialGroup": "projectReaders"
+          }
         ],
-        "creationTime": "1553208775542",
-        "lastModifiedTime": "1686338918114",
+        "creationTime": "1640891845643",
+        "lastModifiedTime": "1640891845643",
         "location": "US",
         "type": "DEFAULT",
         "maxTimeTravelHours": "168"
       }
   """
   try:
-    bq_client = client.get_bigquery_client(credentials=credentials)
+    bq_client = client.get_bigquery_client(
+        project=project_id, credentials=credentials
+    )
     dataset = bq_client.get_dataset(
         bigquery.DatasetReference(project_id, dataset_id)
     )
@@ -128,16 +136,14 @@ def list_table_ids(
       list[str]: List of the tables ids present in the dataset.
 
   Examples:
-      >>> list_table_ids("bigquery-public-data", "ml_datasets")
-      ['census_adult_income',
-       'credit_card_default',
-       'holidays_and_events_for_forecasting',
-       'iris',
-       'penguins',
-       'ulb_fraud_detection']
+      >>> list_table_ids("bigquery-public-data", "cdc_places")
+      ['chronic_disease_indicators',
+       'local_data_for_better_health_county_data']
   """
   try:
-    bq_client = client.get_bigquery_client(credentials=credentials)
+    bq_client = client.get_bigquery_client(
+        project=project_id, credentials=credentials
+    )
 
     tables = []
     for table in bq_client.list_tables(
@@ -167,76 +173,93 @@ def get_table_info(
       dict: Dictionary representing the properties of the table.
 
   Examples:
-      >>> get_table_info("bigquery-public-data", "ml_datasets", "penguins")
+      >>> get_table_info("bigquery-public-data", "cdc_places", "local_data_for_better_health_county_data")
       {
         "kind": "bigquery#table",
-        "etag": "X0ZkRohSGoYvWemRYEgOHA==",
-        "id": "bigquery-public-data:ml_datasets.penguins",
-        "selfLink":
-        "https://bigquery.googleapis.com/bigquery/v2/projects/bigquery-public-data/datasets/ml_datasets/tables/penguins",
+        "etag": "wx23aDqmgc39oUSiNuYTAA==",
+        "id": "bigquery-public-data:cdc_places.local_data_for_better_health_county_data",
+        "selfLink": "https://content-bigquery.googleapis.com/bigquery/v2/projects/bigquery-public-data/datasets/cdc_places/tables/local_data_for_better_health_county_data",
         "tableReference": {
-            "projectId": "bigquery-public-data",
-            "datasetId": "ml_datasets",
-            "tableId": "penguins"
+          "projectId": "bigquery-public-data",
+          "datasetId": "cdc_places",
+          "tableId": "local_data_for_better_health_county_data"
         },
+        "description": "Local Data for Better Health, County Data",
         "schema": {
-            "fields": [
-                {
-                    "name": "species",
-                    "type": "STRING",
-                    "mode": "REQUIRED"
-                },
-                {
-                    "name": "island",
-                    "type": "STRING",
-                    "mode": "NULLABLE"
-                },
-                {
-                    "name": "culmen_length_mm",
-                    "type": "FLOAT",
-                    "mode": "NULLABLE"
-                },
-                {
-                    "name": "culmen_depth_mm",
-                    "type": "FLOAT",
-                    "mode": "NULLABLE"
-                },
-                {
-                    "name": "flipper_length_mm",
-                    "type": "FLOAT",
-                    "mode": "NULLABLE"
-                },
-                {
-                    "name": "body_mass_g",
-                    "type": "FLOAT",
-                    "mode": "NULLABLE"
-                },
-                {
-                    "name": "sex",
-                    "type": "STRING",
-                    "mode": "NULLABLE"
-                }
-            ]
+          "fields": [
+            {
+              "name": "year",
+              "type": "INTEGER",
+              "mode": "NULLABLE"
+            },
+            {
+              "name": "stateabbr",
+              "type": "STRING",
+              "mode": "NULLABLE"
+            },
+            {
+              "name": "statedesc",
+              "type": "STRING",
+              "mode": "NULLABLE"
+            },
+            {
+              "name": "locationname",
+              "type": "STRING",
+              "mode": "NULLABLE"
+            },
+            {
+              "name": "datasource",
+              "type": "STRING",
+              "mode": "NULLABLE"
+            },
+            {
+              "name": "category",
+              "type": "STRING",
+              "mode": "NULLABLE"
+            },
+            {
+              "name": "measure",
+              "type": "STRING",
+              "mode": "NULLABLE"
+            },
+            {
+              "name": "data_value_unit",
+              "type": "STRING",
+              "mode": "NULLABLE"
+            },
+            {
+              "name": "data_value_type",
+              "type": "STRING",
+              "mode": "NULLABLE"
+            },
+            {
+              "name": "data_value",
+              "type": "FLOAT",
+              "mode": "NULLABLE"
+            }
+          ]
         },
-        "numBytes": "28947",
-        "numLongTermBytes": "28947",
-        "numRows": "344",
-        "creationTime": "1619804743188",
-        "lastModifiedTime": "1634584675234",
+        "numBytes": "234849",
+        "numLongTermBytes": "0",
+        "numRows": "1000",
+        "creationTime": "1640891846119",
+        "lastModifiedTime": "1749427268137",
         "type": "TABLE",
         "location": "US",
-        "numTimeTravelPhysicalBytes": "0",
-        "numTotalLogicalBytes": "28947",
-        "numActiveLogicalBytes": "0",
-        "numLongTermLogicalBytes": "28947",
-        "numTotalPhysicalBytes": "5350",
-        "numActivePhysicalBytes": "0",
-        "numLongTermPhysicalBytes": "5350",
-        "numCurrentPhysicalBytes": "5350"
+        "numTimeTravelPhysicalBytes": "285737",
+        "numTotalLogicalBytes": "234849",
+        "numActiveLogicalBytes": "234849",
+        "numLongTermLogicalBytes": "0",
+        "numTotalPhysicalBytes": "326557",
+        "numActivePhysicalBytes": "326557",
+        "numLongTermPhysicalBytes": "0",
+        "numCurrentPhysicalBytes": "40820"
       }
   """
   try:
-    bq_client = client.get_bigquery_client(credentials=credentials)
+    bq_client = client.get_bigquery_client(
+        project=project_id, credentials=credentials
+    )
     return bq_client.get_table(
         bigquery.TableReference(
             bigquery.DatasetReference(project_id, dataset_id), table_id