feat: allow setting agent/application name for BigQuery tools

This will allow tracking of tool usage per agent/application.

PiperOrigin-RevId: 800607186

contributing/samples/bigquery/agent.py

@@ -25,13 +25,18 @@ import google.auth
 # Define an appropriate credential type
 CREDENTIALS_TYPE = AuthCredentialTypes.OAUTH2
 
+# Define an appropriate application name
+BIGQUERY_AGENT_NAME = "adk_sample_bigquery_agent"
+
 
 # Define BigQuery tool config with write mode set to allowed. Note that this is
 # only to demonstrate the full capability of the BigQuery tools. In production
 # you may want to change to BLOCKED (default write mode, effectively makes the
 # tool read-only) or PROTECTED (only allows writes in the anonymous dataset of a
 # BigQuery session) write mode.
-tool_config = BigQueryToolConfig(write_mode=WriteMode.ALLOWED)
+tool_config = BigQueryToolConfig(
+    write_mode=WriteMode.ALLOWED, application_name=BIGQUERY_AGENT_NAME
+)
 
 if CREDENTIALS_TYPE == AuthCredentialTypes.OAUTH2:
   # Initiaze the tools to do interactive OAuth
@@ -64,7 +69,7 @@ bigquery_toolset = BigQueryToolset(
 # debug CLI
 root_agent = LlmAgent(
     model="gemini-2.0-flash",
-    name="bigquery_agent",
+    name=BIGQUERY_AGENT_NAME,
     description=(
         "Agent to answer questions about BigQuery data and models and execute"
         " SQL queries."

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

@@ -26,11 +26,16 @@ USER_AGENT = f"adk-bigquery-tool google-adk/{version.__version__}"
 
 
 def get_bigquery_client(
-    *, project: Optional[str], credentials: Credentials
+    *,
+    project: Optional[str],
+    credentials: Credentials,
+    user_agent: Optional[str] = None,
 ) -> bigquery.Client:
   """Get a BigQuery client."""
 
-  client_info = google.api_core.client_info.ClientInfo(user_agent=USER_AGENT)
+  user_agent = f"{USER_AGENT} {user_agent}" if user_agent else USER_AGENT
+
+  client_info = google.api_core.client_info.ClientInfo(user_agent=user_agent)
 
   bigquery_client = bigquery.Client(
       project=project, credentials=credentials, client_info=client_info

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

@@ -15,8 +15,10 @@
 from __future__ import annotations
 
 from enum import Enum
+from typing import Optional
 
 from pydantic import BaseModel
+from pydantic import field_validator
 
 from ...utils.feature_decorator import experimental
 
@@ -58,4 +60,21 @@ class BigQueryToolConfig(BaseModel):
   max_query_result_rows: int = 50
   """Maximum number of rows to return from a query.
 
-  By default, the query result will be limited to 50 rows."""
+  By default, the query result will be limited to 50 rows.
+  """
+
+  application_name: Optional[str] = None
+  """Name of the application using the BigQuery tools.
+
+  By default, no particular application name will be set in the BigQuery
+  interaction. But if the the tool user (agent builder) wants to differentiate
+  their application/agent for tracking or support purpose, they can set this field.
+  """
+
+  @field_validator('application_name')
+  @classmethod
+  def validate_application_name(cls, v):
+    """Validate the application name."""
+    if v and ' ' in v:
+      raise ValueError('Application name should not contain spaces.')
+    return v

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

@@ -18,9 +18,12 @@ from google.auth.credentials import Credentials
 from google.cloud import bigquery
 
 from . import client
+from .config import BigQueryToolConfig
 
 
-def list_dataset_ids(project_id: str, credentials: Credentials) -> list[str]:
+def list_dataset_ids(
+    project_id: str, credentials: Credentials, settings: BigQueryToolConfig
+) -> list[str]:
   """List BigQuery dataset ids in a Google Cloud project.
 
   Args:
@@ -45,7 +48,9 @@ def list_dataset_ids(project_id: str, credentials: Credentials) -> list[str]:
   """
   try:
     bq_client = client.get_bigquery_client(
-        project=project_id, credentials=credentials
+        project=project_id,
+        credentials=credentials,
+        user_agent=settings.application_name,
     )
 
     datasets = []
@@ -60,7 +65,10 @@ def list_dataset_ids(project_id: str, credentials: Credentials) -> list[str]:
 
 
 def get_dataset_info(
-    project_id: str, dataset_id: str, credentials: Credentials
+    project_id: str,
+    dataset_id: str,
+    credentials: Credentials,
+    settings: BigQueryToolConfig,
 ) -> dict:
   """Get metadata information about a BigQuery dataset.
 
@@ -111,7 +119,9 @@ def get_dataset_info(
   """
   try:
     bq_client = client.get_bigquery_client(
-        project=project_id, credentials=credentials
+        project=project_id,
+        credentials=credentials,
+        user_agent=settings.application_name,
     )
     dataset = bq_client.get_dataset(
         bigquery.DatasetReference(project_id, dataset_id)
@@ -125,7 +135,10 @@ def get_dataset_info(
 
 
 def list_table_ids(
-    project_id: str, dataset_id: str, credentials: Credentials
+    project_id: str,
+    dataset_id: str,
+    credentials: Credentials,
+    settings: BigQueryToolConfig,
 ) -> list[str]:
   """List table ids in a BigQuery dataset.
 
@@ -144,7 +157,9 @@ def list_table_ids(
   """
   try:
     bq_client = client.get_bigquery_client(
-        project=project_id, credentials=credentials
+        project=project_id,
+        credentials=credentials,
+        user_agent=settings.application_name,
     )
 
     tables = []
@@ -161,7 +176,11 @@ def list_table_ids(
 
 
 def get_table_info(
-    project_id: str, dataset_id: str, table_id: str, credentials: Credentials
+    project_id: str,
+    dataset_id: str,
+    table_id: str,
+    credentials: Credentials,
+    settings: BigQueryToolConfig,
 ) -> dict:
   """Get metadata information about a BigQuery table.
 
@@ -260,7 +279,9 @@ def get_table_info(
   """
   try:
     bq_client = client.get_bigquery_client(
-        project=project_id, credentials=credentials
+        project=project_id,
+        credentials=credentials,
+        user_agent=settings.application_name,
     )
     return bq_client.get_table(
         bigquery.TableReference(

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

@@ -80,7 +80,9 @@ def execute_sql(
   try:
     # Get BigQuery client
     bq_client = client.get_bigquery_client(
-        project=project_id, credentials=credentials
+        project=project_id,
+        credentials=credentials,
+        user_agent=settings.application_name,
     )
 
     # BigQuery connection properties where applicable

tests/unittests/tools/bigquery/test_bigquery_client.py

@@ -15,9 +15,9 @@
 from __future__ import annotations
 
 import os
-import re
 from unittest import mock
 
+import google.adk
 from google.adk.tools.bigquery.client import get_bigquery_client
 from google.auth.exceptions import DefaultCredentialsError
 from google.oauth2.credentials import Credentials
@@ -109,8 +109,8 @@ def test_bigquery_client_project_set_with_env():
       assert client.project == "test-gcp-project"
 
 
-def test_bigquery_client_user_agent():
-  """Test BigQuery client user agent."""
+def test_bigquery_client_user_agent_default():
+  """Test BigQuery client default user agent."""
   with mock.patch(
       "google.cloud.bigquery.client.Connection", autospec=True
   ) as mock_connection:
@@ -123,7 +123,33 @@ def test_bigquery_client_user_agent():
     # Verify that the tracking user agent was set
     client_info_arg = mock_connection.call_args[1].get("client_info")
     assert client_info_arg is not None
-    assert re.search(
-        r"adk-bigquery-tool google-adk/([0-9A-Za-z._\-+/]+)",
-        client_info_arg.user_agent,
+    expected_user_agents = {
+        "adk-bigquery-tool",
+        f"google-adk/{google.adk.__version__}",
+    }
+    actual_user_agents = set(client_info_arg.user_agent.split())
+    assert expected_user_agents.issubset(actual_user_agents)
+
+
+def test_bigquery_client_user_agent_custom():
+  """Test BigQuery client custom user agent."""
+  with mock.patch(
+      "google.cloud.bigquery.client.Connection", autospec=True
+  ) as mock_connection:
+    # Trigger the BigQuery client creation
+    get_bigquery_client(
+        project="test-gcp-project",
+        credentials=mock.create_autospec(Credentials, instance=True),
+        user_agent="custom_user_agent",
     )
+
+    # Verify that the tracking user agent was set
+    client_info_arg = mock_connection.call_args[1].get("client_info")
+    assert client_info_arg is not None
+    expected_user_agents = {
+        "adk-bigquery-tool",
+        f"google-adk/{google.adk.__version__}",
+        "custom_user_agent",
+    }
+    actual_user_agents = set(client_info_arg.user_agent.split())
+    assert expected_user_agents.issubset(actual_user_agents)

tests/unittests/tools/bigquery/test_bigquery_metadata_tool.py

@@ -18,19 +18,22 @@ import os
 from unittest import mock
 
 from google.adk.tools.bigquery import metadata_tool
+from google.adk.tools.bigquery.config import BigQueryToolConfig
 from google.auth.exceptions import DefaultCredentialsError
 from google.cloud import bigquery
 from google.oauth2.credentials import Credentials
-import pytest
 
 
 @mock.patch.dict(os.environ, {}, clear=True)
 @mock.patch("google.cloud.bigquery.Client.list_datasets", autospec=True)
 @mock.patch("google.auth.default", autospec=True)
-def test_list_dataset_ids(mock_default_auth, mock_list_datasets):
-  """Test list_dataset_ids tool invocation."""
+def test_list_dataset_ids_no_default_auth(
+    mock_default_auth, mock_list_datasets
+):
+  """Test list_dataset_ids tool invocation involves no default auth."""
   project = "my_project_id"
   mock_credentials = mock.create_autospec(Credentials, instance=True)
+  tool_settings = BigQueryToolConfig()
 
   # Simulate the behavior of default auth - on purpose throw exception when
   # the default auth is called
@@ -42,7 +45,9 @@ def test_list_dataset_ids(mock_default_auth, mock_list_datasets):
       bigquery.DatasetReference(project, "dataset1"),
       bigquery.DatasetReference(project, "dataset2"),
   ]
-  result = metadata_tool.list_dataset_ids(project, mock_credentials)
+  result = metadata_tool.list_dataset_ids(
+      project, mock_credentials, tool_settings
+  )
   assert result == ["dataset1", "dataset2"]
   mock_default_auth.assert_not_called()
 
@@ -50,9 +55,10 @@ def test_list_dataset_ids(mock_default_auth, mock_list_datasets):
 @mock.patch.dict(os.environ, {}, clear=True)
 @mock.patch("google.cloud.bigquery.Client.get_dataset", autospec=True)
 @mock.patch("google.auth.default", autospec=True)
-def test_get_dataset_info(mock_default_auth, mock_get_dataset):
-  """Test get_dataset_info tool invocation."""
+def test_get_dataset_info_no_default_auth(mock_default_auth, mock_get_dataset):
+  """Test get_dataset_info tool invocation involves no default auth."""
   mock_credentials = mock.create_autospec(Credentials, instance=True)
+  tool_settings = BigQueryToolConfig()
 
   # Simulate the behavior of default auth - on purpose throw exception when
   # the default auth is called
@@ -64,7 +70,7 @@ def test_get_dataset_info(mock_default_auth, mock_get_dataset):
       Credentials, instance=True
   )
   result = metadata_tool.get_dataset_info(
-      "my_project_id", "my_dataset_id", mock_credentials
+      "my_project_id", "my_dataset_id", mock_credentials, tool_settings
   )
   assert result != {
       "status": "ERROR",
@@ -76,12 +82,13 @@ def test_get_dataset_info(mock_default_auth, mock_get_dataset):
 @mock.patch.dict(os.environ, {}, clear=True)
 @mock.patch("google.cloud.bigquery.Client.list_tables", autospec=True)
 @mock.patch("google.auth.default", autospec=True)
-def test_list_table_ids(mock_default_auth, mock_list_tables):
-  """Test list_table_ids tool invocation."""
+def test_list_table_ids_no_default_auth(mock_default_auth, mock_list_tables):
+  """Test list_table_ids tool invocation involves no default auth."""
   project = "my_project_id"
   dataset = "my_dataset_id"
   dataset_ref = bigquery.DatasetReference(project, dataset)
   mock_credentials = mock.create_autospec(Credentials, instance=True)
+  tool_settings = BigQueryToolConfig()
 
   # Simulate the behavior of default auth - on purpose throw exception when
   # the default auth is called
@@ -93,7 +100,9 @@ def test_list_table_ids(mock_default_auth, mock_list_tables):
       bigquery.TableReference(dataset_ref, "table1"),
       bigquery.TableReference(dataset_ref, "table2"),
   ]
-  result = metadata_tool.list_table_ids(project, dataset, mock_credentials)
+  result = metadata_tool.list_table_ids(
+      project, dataset, mock_credentials, tool_settings
+  )
   assert result == ["table1", "table2"]
   mock_default_auth.assert_not_called()
 
@@ -101,9 +110,10 @@ def test_list_table_ids(mock_default_auth, mock_list_tables):
 @mock.patch.dict(os.environ, {}, clear=True)
 @mock.patch("google.cloud.bigquery.Client.get_table", autospec=True)
 @mock.patch("google.auth.default", autospec=True)
-def test_get_table_info(mock_default_auth, mock_get_table):
-  """Test get_table_info tool invocation."""
+def test_get_table_info_no_default_auth(mock_default_auth, mock_get_table):
+  """Test get_table_info tool invocation involves no default auth."""
   mock_credentials = mock.create_autospec(Credentials, instance=True)
+  tool_settings = BigQueryToolConfig()
 
   # Simulate the behavior of default auth - on purpose throw exception when
   # the default auth is called
@@ -113,10 +123,116 @@ def test_get_table_info(mock_default_auth, mock_get_table):
 
   mock_get_table.return_value = mock.create_autospec(Credentials, instance=True)
   result = metadata_tool.get_table_info(
-      "my_project_id", "my_dataset_id", "my_table_id", mock_credentials
+      "my_project_id",
+      "my_dataset_id",
+      "my_table_id",
+      mock_credentials,
+      tool_settings,
   )
   assert result != {
       "status": "ERROR",
       "error_details": "Your default credentials were not found",
   }
   mock_default_auth.assert_not_called()
+
+
+@mock.patch(
+    "google.adk.tools.bigquery.client.get_bigquery_client", autospec=True
+)
+def test_list_dataset_ids_bq_client_creation(mock_get_bigquery_client):
+  """Test BigQuery client creation params during list_dataset_ids tool invocation."""
+  bq_project = "my_project_id"
+  bq_credentials = mock.create_autospec(Credentials, instance=True)
+  application_name = "my-agent"
+  tool_settings = BigQueryToolConfig(application_name=application_name)
+
+  metadata_tool.list_dataset_ids(bq_project, bq_credentials, tool_settings)
+  mock_get_bigquery_client.assert_called_once()
+  assert len(mock_get_bigquery_client.call_args.kwargs) == 3
+  assert mock_get_bigquery_client.call_args.kwargs["project"] == bq_project
+  assert (
+      mock_get_bigquery_client.call_args.kwargs["credentials"] == bq_credentials
+  )
+  assert (
+      mock_get_bigquery_client.call_args.kwargs["user_agent"]
+      == application_name
+  )
+
+
+@mock.patch(
+    "google.adk.tools.bigquery.client.get_bigquery_client", autospec=True
+)
+def test_get_dataset_info_bq_client_creation(mock_get_bigquery_client):
+  """Test BigQuery client creation params during get_dataset_info tool invocation."""
+  bq_project = "my_project_id"
+  bq_dataset = "my_dataset_id"
+  bq_credentials = mock.create_autospec(Credentials, instance=True)
+  application_name = "my-agent"
+  tool_settings = BigQueryToolConfig(application_name=application_name)
+
+  metadata_tool.get_dataset_info(
+      bq_project, bq_dataset, bq_credentials, tool_settings
+  )
+  mock_get_bigquery_client.assert_called_once()
+  assert len(mock_get_bigquery_client.call_args.kwargs) == 3
+  assert mock_get_bigquery_client.call_args.kwargs["project"] == bq_project
+  assert (
+      mock_get_bigquery_client.call_args.kwargs["credentials"] == bq_credentials
+  )
+  assert (
+      mock_get_bigquery_client.call_args.kwargs["user_agent"]
+      == application_name
+  )
+
+
+@mock.patch(
+    "google.adk.tools.bigquery.client.get_bigquery_client", autospec=True
+)
+def test_list_table_ids_bq_client_creation(mock_get_bigquery_client):
+  """Test BigQuery client creation params during list_table_ids tool invocation."""
+  bq_project = "my_project_id"
+  bq_dataset = "my_dataset_id"
+  bq_credentials = mock.create_autospec(Credentials, instance=True)
+  application_name = "my-agent"
+  tool_settings = BigQueryToolConfig(application_name=application_name)
+
+  metadata_tool.list_table_ids(
+      bq_project, bq_dataset, bq_credentials, tool_settings
+  )
+  mock_get_bigquery_client.assert_called_once()
+  assert len(mock_get_bigquery_client.call_args.kwargs) == 3
+  assert mock_get_bigquery_client.call_args.kwargs["project"] == bq_project
+  assert (
+      mock_get_bigquery_client.call_args.kwargs["credentials"] == bq_credentials
+  )
+  assert (
+      mock_get_bigquery_client.call_args.kwargs["user_agent"]
+      == application_name
+  )
+
+
+@mock.patch(
+    "google.adk.tools.bigquery.client.get_bigquery_client", autospec=True
+)
+def test_get_table_info_bq_client_creation(mock_get_bigquery_client):
+  """Test BigQuery client creation params during get_table_info tool invocation."""
+  bq_project = "my_project_id"
+  bq_dataset = "my_dataset_id"
+  bq_table = "my_table_id"
+  bq_credentials = mock.create_autospec(Credentials, instance=True)
+  application_name = "my-agent"
+  tool_settings = BigQueryToolConfig(application_name=application_name)
+
+  metadata_tool.get_table_info(
+      bq_project, bq_dataset, bq_table, bq_credentials, tool_settings
+  )
+  mock_get_bigquery_client.assert_called_once()
+  assert len(mock_get_bigquery_client.call_args.kwargs) == 3
+  assert mock_get_bigquery_client.call_args.kwargs["project"] == bq_project
+  assert (
+      mock_get_bigquery_client.call_args.kwargs["credentials"] == bq_credentials
+  )
+  assert (
+      mock_get_bigquery_client.call_args.kwargs["user_agent"]
+      == application_name
+  )

tests/unittests/tools/bigquery/test_bigquery_query_tool.py

@@ -983,3 +983,26 @@ def test_execute_sql_result_dtype(
   # Test the tool worked without invoking default auth
   result = execute_sql(project, query, credentials, tool_settings, tool_context)
   assert result == {"status": "SUCCESS", "rows": tool_result_rows}
+
+
+@mock.patch(
+    "google.adk.tools.bigquery.client.get_bigquery_client", autospec=True
+)
+def test_execute_sql_bq_client_creation(mock_get_bigquery_client):
+  """Test BigQuery client creation params during execute_sql tool invocation."""
+  project = "my_project_id"
+  query = "SELECT 1"
+  credentials = mock.create_autospec(Credentials, instance=True)
+  application_name = "my-agent"
+  tool_settings = BigQueryToolConfig(application_name=application_name)
+  tool_context = mock.create_autospec(ToolContext, instance=True)
+
+  execute_sql(project, query, credentials, tool_settings, tool_context)
+  mock_get_bigquery_client.assert_called_once()
+  assert len(mock_get_bigquery_client.call_args.kwargs) == 3
+  assert mock_get_bigquery_client.call_args.kwargs["project"] == project
+  assert mock_get_bigquery_client.call_args.kwargs["credentials"] == credentials
+  assert (
+      mock_get_bigquery_client.call_args.kwargs["user_agent"]
+      == application_name
+  )

tests/unittests/tools/bigquery/test_bigquery_tool_config.py

@@ -25,3 +25,12 @@ def test_bigquery_tool_config_experimental_warning():
       match="Config defaults may have breaking change in the future.",
   ):
     BigQueryToolConfig()
+
+
+def test_bigquery_tool_config_invalid_application_name():
+  """Test BigQueryToolConfig with invalid application name."""
+  with pytest.raises(
+      ValueError,
+      match="Application name should not contain spaces.",
+  ):
+    BigQueryToolConfig(application_name="my agent")