feat: set per-tool user agent in BQ calls and tool label in BQ jobs

This will help per tool usage for BigQuery tools.

PiperOrigin-RevId: 829142106

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

@@ -25,12 +25,16 @@ from ... import version
 USER_AGENT = f"adk-bigquery-tool google-adk/{version.__version__}"
 
 
+from typing import List
+from typing import Union
+
+
 def get_bigquery_client(
     *,
     project: Optional[str],
     credentials: Credentials,
     location: Optional[str] = None,
-    user_agent: Optional[str] = None,
+    user_agent: Optional[Union[str, List[str]]] = None,
 ) -> bigquery.Client:
   """Get a BigQuery client.
 
@@ -44,9 +48,16 @@ def get_bigquery_client(
     A BigQuery client.
   """
 
-  user_agent = f"{USER_AGENT} {user_agent}" if user_agent else USER_AGENT
+  user_agents = [USER_AGENT]
+  if user_agent:
+    if isinstance(user_agent, str):
+      user_agents.append(user_agent)
+    else:
+      user_agents.extend([ua for ua in user_agent if ua])
 
-  client_info = google.api_core.client_info.ClientInfo(user_agent=user_agent)
+  client_info = google.api_core.client_info.ClientInfo(
+      user_agent=" ".join(user_agents)
+  )
 
   bigquery_client = bigquery.Client(
       project=project,

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

@@ -51,7 +51,7 @@ def list_dataset_ids(
         project=project_id,
         credentials=credentials,
         location=settings.location,
-        user_agent=settings.application_name,
+        user_agent=[settings.application_name, "list_dataset_ids"],
     )
 
     datasets = []
@@ -123,7 +123,7 @@ def get_dataset_info(
         project=project_id,
         credentials=credentials,
         location=settings.location,
-        user_agent=settings.application_name,
+        user_agent=[settings.application_name, "get_dataset_info"],
     )
     dataset = bq_client.get_dataset(
         bigquery.DatasetReference(project_id, dataset_id)
@@ -162,7 +162,7 @@ def list_table_ids(
         project=project_id,
         credentials=credentials,
         location=settings.location,
-        user_agent=settings.application_name,
+        user_agent=[settings.application_name, "list_table_ids"],
     )
 
     tables = []
@@ -285,7 +285,7 @@ def get_table_info(
         project=project_id,
         credentials=credentials,
         location=settings.location,
-        user_agent=settings.application_name,
+        user_agent=[settings.application_name, "get_table_info"],
     )
     return bq_client.get_table(
         bigquery.TableReference(
@@ -579,8 +579,9 @@ def get_job_info(
         project=project_id,
         credentials=credentials,
         location=settings.location,
-        user_agent=settings.application_name,
+        user_agent=[settings.application_name, "get_job_info"],
     )
+
     job = bq_client.get_job(job_id)
     # We need to use _properties to get the job info because it contains all
     # the job info.

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

@@ -32,6 +32,160 @@ from .config import WriteMode
 BIGQUERY_SESSION_INFO_KEY = "bigquery_session_info"
 
 
+def _execute_sql(
+    project_id: str,
+    query: str,
+    credentials: Credentials,
+    settings: BigQueryToolConfig,
+    tool_context: ToolContext,
+    dry_run: bool = False,
+    caller_id: Optional[str] = None,
+) -> dict:
+  try:
+    # Validate compute project if applicable
+    if (
+        settings.compute_project_id
+        and project_id != settings.compute_project_id
+    ):
+      return {
+          "status": "ERROR",
+          "error_details": (
+              f"Cannot execute query in the project {project_id}, as the tool"
+              " is restricted to execute queries only in the project"
+              f" {settings.compute_project_id}."
+          ),
+      }
+
+    # Get BigQuery client
+    bq_client = client.get_bigquery_client(
+        project=project_id,
+        credentials=credentials,
+        location=settings.location,
+        user_agent=[settings.application_name, caller_id],
+    )
+
+    # BigQuery connection properties where applicable
+    bq_connection_properties = []
+
+    # BigQuery job labels if applicable
+    bq_job_labels = {}
+    if caller_id:
+      bq_job_labels["adk-bigquery-tool"] = caller_id
+
+    if not settings or settings.write_mode == WriteMode.BLOCKED:
+      dry_run_query_job = bq_client.query(
+          query,
+          project=project_id,
+          job_config=bigquery.QueryJobConfig(
+              dry_run=True, labels=bq_job_labels
+          ),
+      )
+      if dry_run_query_job.statement_type != "SELECT":
+        return {
+            "status": "ERROR",
+            "error_details": "Read-only mode only supports SELECT statements.",
+        }
+    elif settings.write_mode == WriteMode.PROTECTED:
+      # In protected write mode, write operation only to a temporary artifact is
+      # allowed. This artifact must have been created in a BigQuery session. In
+      # such a scenario, the session info (session id and the anonymous dataset
+      # containing the artifact) is persisted in the tool context.
+      bq_session_info = tool_context.state.get(BIGQUERY_SESSION_INFO_KEY, None)
+      if bq_session_info:
+        bq_session_id, bq_session_dataset_id = bq_session_info
+      else:
+        session_creator_job = bq_client.query(
+            "SELECT 1",
+            project=project_id,
+            job_config=bigquery.QueryJobConfig(
+                dry_run=True, create_session=True, labels=bq_job_labels
+            ),
+        )
+        bq_session_id = session_creator_job.session_info.session_id
+        bq_session_dataset_id = session_creator_job.destination.dataset_id
+
+        # Remember the BigQuery session info for subsequent queries
+        tool_context.state[BIGQUERY_SESSION_INFO_KEY] = (
+            bq_session_id,
+            bq_session_dataset_id,
+        )
+
+      # Session connection property will be set in the query execution
+      bq_connection_properties.append(
+          bigquery.ConnectionProperty("session_id", bq_session_id)
+      )
+
+      # Check the query type w.r.t. the BigQuery session
+      dry_run_query_job = bq_client.query(
+          query,
+          project=project_id,
+          job_config=bigquery.QueryJobConfig(
+              dry_run=True,
+              connection_properties=bq_connection_properties,
+              labels=bq_job_labels,
+          ),
+      )
+      if (
+          dry_run_query_job.statement_type != "SELECT"
+          and dry_run_query_job.destination.dataset_id != bq_session_dataset_id
+      ):
+        return {
+            "status": "ERROR",
+            "error_details": (
+                "Protected write mode only supports SELECT statements, or write"
+                " operations in the anonymous dataset of a BigQuery session."
+            ),
+        }
+
+    # Return the dry run characteristics of the query if requested
+    if dry_run:
+      dry_run_job = bq_client.query(
+          query,
+          project=project_id,
+          job_config=bigquery.QueryJobConfig(
+              dry_run=True,
+              connection_properties=bq_connection_properties,
+              labels=bq_job_labels,
+          ),
+      )
+      return {"status": "SUCCESS", "dry_run_info": dry_run_job.to_api_repr()}
+
+    # Finally execute the query, fetch the result, and return it
+    row_iterator = bq_client.query_and_wait(
+        query,
+        job_config=bigquery.QueryJobConfig(
+            connection_properties=bq_connection_properties,
+            labels=bq_job_labels,
+        ),
+        project=project_id,
+        max_results=settings.max_query_result_rows,
+    )
+    rows = []
+    for row in row_iterator:
+      row_values = {}
+      for key, val in row.items():
+        try:
+          # if the json serialization of the value succeeds, use it as is
+          json.dumps(val)
+        except:
+          val = str(val)
+        row_values[key] = val
+      rows.append(row_values)
+
+    result = {"status": "SUCCESS", "rows": rows}
+    if (
+        settings.max_query_result_rows is not None
+        and len(rows) == settings.max_query_result_rows
+    ):
+      result["result_is_likely_truncated"] = True
+    return result
+  except Exception as ex:  # pylint: disable=broad-except
+    return {
+        "status": "ERROR",
+        "error_details": str(ex),
+    }
+
+
 def execute_sql(
     project_id: str,
     query: str,
@@ -118,142 +272,15 @@ def execute_sql(
             }
           }
   """
-  try:
-    # Validate compute project if applicable
-    if (
-        settings.compute_project_id
-        and project_id != settings.compute_project_id
-    ):
-      return {
-          "status": "ERROR",
-          "error_details": (
-              f"Cannot execute query in the project {project_id}, as the tool"
-              " is restricted to execute queries only in the project"
-              f" {settings.compute_project_id}."
-          ),
-      }
-
-    # Get BigQuery client
-    bq_client = client.get_bigquery_client(
-        project=project_id,
-        credentials=credentials,
-        location=settings.location,
-        user_agent=settings.application_name,
-    )
-
-    # BigQuery connection properties where applicable
-    bq_connection_properties = None
-
-    if not settings or settings.write_mode == WriteMode.BLOCKED:
-      dry_run_query_job = bq_client.query(
-          query,
-          project=project_id,
-          job_config=bigquery.QueryJobConfig(dry_run=True),
-      )
-      if dry_run_query_job.statement_type != "SELECT":
-        return {
-            "status": "ERROR",
-            "error_details": "Read-only mode only supports SELECT statements.",
-        }
-    elif settings.write_mode == WriteMode.PROTECTED:
-      # In protected write mode, write operation only to a temporary artifact is
-      # allowed. This artifact must have been created in a BigQuery session. In
-      # such a scenario, the session info (session id and the anonymous dataset
-      # containing the artifact) is persisted in the tool context.
-      bq_session_info = tool_context.state.get(BIGQUERY_SESSION_INFO_KEY, None)
-      if bq_session_info:
-        bq_session_id, bq_session_dataset_id = bq_session_info
-      else:
-        session_creator_job = bq_client.query(
-            "SELECT 1",
-            project=project_id,
-            job_config=bigquery.QueryJobConfig(
-                dry_run=True, create_session=True
-            ),
-        )
-        bq_session_id = session_creator_job.session_info.session_id
-        bq_session_dataset_id = session_creator_job.destination.dataset_id
-
-        # Remember the BigQuery session info for subsequent queries
-        tool_context.state[BIGQUERY_SESSION_INFO_KEY] = (
-            bq_session_id,
-            bq_session_dataset_id,
-        )
-
-      # Session connection property will be set in the query execution
-      bq_connection_properties = [
-          bigquery.ConnectionProperty("session_id", bq_session_id)
-      ]
-
-      # Check the query type w.r.t. the BigQuery session
-      dry_run_query_job = bq_client.query(
-          query,
-          project=project_id,
-          job_config=bigquery.QueryJobConfig(
-              dry_run=True,
-              connection_properties=bq_connection_properties,
-          ),
-      )
-      if (
-          dry_run_query_job.statement_type != "SELECT"
-          and dry_run_query_job.destination.dataset_id != bq_session_dataset_id
-      ):
-        return {
-            "status": "ERROR",
-            "error_details": (
-                "Protected write mode only supports SELECT statements, or write"
-                " operations in the anonymous dataset of a BigQuery session."
-            ),
-        }
-
-    # Finally execute the query and fetch the result
-    if dry_run:
-      job_config_kwargs = {"dry_run": True}
-      if bq_connection_properties:
-        job_config_kwargs["connection_properties"] = bq_connection_properties
-      job_config = bigquery.QueryJobConfig(**job_config_kwargs)
-      dry_run_job = bq_client.query(
-          query,
-          project=project_id,
-          job_config=job_config,
-      )
-      return {"status": "SUCCESS", "dry_run_info": dry_run_job.to_api_repr()}
-
-    job_config = (
-        bigquery.QueryJobConfig(connection_properties=bq_connection_properties)
-        if bq_connection_properties
-        else None
-    )
-    row_iterator = bq_client.query_and_wait(
-        query,
-        job_config=job_config,
-        project=project_id,
-        max_results=settings.max_query_result_rows,
-    )
-    rows = []
-    for row in row_iterator:
-      row_values = {}
-      for key, val in row.items():
-        try:
-          # if the json serialization of the value succeeds, use it as is
-          json.dumps(val)
-        except:
-          val = str(val)
-        row_values[key] = val
-      rows.append(row_values)
-
-    result = {"status": "SUCCESS", "rows": rows}
-    if (
-        settings.max_query_result_rows is not None
-        and len(rows) == settings.max_query_result_rows
-    ):
-      result["result_is_likely_truncated"] = True
-    return result
-  except Exception as ex:  # pylint: disable=broad-except
-    return {
-        "status": "ERROR",
-        "error_details": str(ex),
-    }
+  return _execute_sql(
+      project_id=project_id,
+      query=query,
+      credentials=credentials,
+      settings=settings,
+      tool_context=tool_context,
+      dry_run=dry_run,
+      caller_id="execute_sql",
+  )
 
 
 def _execute_sql_write_mode(*args, **kwargs) -> dict:
@@ -892,7 +919,14 @@ def forecast(
     confidence_level => {confidence_level}
   )
   """
-  return execute_sql(project_id, query, credentials, settings, tool_context)
+  return _execute_sql(
+      project_id=project_id,
+      query=query,
+      credentials=credentials,
+      settings=settings,
+      tool_context=tool_context,
+      caller_id="forecast",
+  )
 
 
 def analyze_contribution(
@@ -1065,22 +1099,24 @@ def analyze_contribution(
       # session.
       settings.write_mode = WriteMode.PROTECTED
 
-    result = execute_sql(
-        project_id,
-        create_model_query,
-        credentials,
-        settings,
-        tool_context,
+    result = _execute_sql(
+        project_id=project_id,
+        query=create_model_query,
+        credentials=credentials,
+        settings=settings,
+        tool_context=tool_context,
+        caller_id="analyze_contribution",
     )
     if result["status"] != "SUCCESS":
       return result
 
-    result = execute_sql(
-        project_id,
-        get_insights_query,
-        credentials,
-        settings,
-        tool_context,
+    result = _execute_sql(
+        project_id=project_id,
+        query=get_insights_query,
+        credentials=credentials,
+        settings=settings,
+        tool_context=tool_context,
+        caller_id="analyze_contribution",
     )
   except Exception as ex:  # pylint: disable=broad-except
     return {
@@ -1292,22 +1328,24 @@ def detect_anomalies(
       # session.
       settings.write_mode = WriteMode.PROTECTED
 
-    result = execute_sql(
-        project_id,
-        create_model_query,
-        credentials,
-        settings,
-        tool_context,
+    result = _execute_sql(
+        project_id=project_id,
+        query=create_model_query,
+        credentials=credentials,
+        settings=settings,
+        tool_context=tool_context,
+        caller_id="detect_anomalies",
     )
     if result["status"] != "SUCCESS":
       return result
 
-    result = execute_sql(
-        project_id,
-        anomaly_detection_query,
-        credentials,
-        settings,
-        tool_context,
+    result = _execute_sql(
+        project_id=project_id,
+        query=anomaly_detection_query,
+        credentials=credentials,
+        settings=settings,
+        tool_context=tool_context,
+        caller_id="detect_anomalies",
     )
   except Exception as ex:  # pylint: disable=broad-except
     return {

tests/unittests/tools/bigquery/test_bigquery_client.py

@@ -156,6 +156,31 @@ def test_bigquery_client_user_agent_custom():
     assert expected_user_agents.issubset(actual_user_agents)
 
 
+def test_bigquery_client_user_agent_custom_list():
+  """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_agent1", "custom_user_agent2"],
+    )
+
+    # Verify that the tracking user agents were 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_agent1",
+        "custom_user_agent2",
+    }
+    actual_user_agents = set(client_info_arg.user_agent.split())
+    assert expected_user_agents.issubset(actual_user_agents)
+
+
 def test_bigquery_client_location_custom():
   """Test BigQuery client custom location."""
   # Trigger the BigQuery client creation

tests/unittests/tools/bigquery/test_bigquery_metadata_tool.py

@@ -183,10 +183,10 @@ def test_list_dataset_ids_bq_client_creation(mock_get_bigquery_client):
   assert (
       mock_get_bigquery_client.call_args.kwargs["credentials"] == bq_credentials
   )
-  assert (
-      mock_get_bigquery_client.call_args.kwargs["user_agent"]
-      == application_name
-  )
+  assert mock_get_bigquery_client.call_args.kwargs["user_agent"] == [
+      application_name,
+      "list_dataset_ids",
+  ]
 
 
 @mock.patch(
@@ -209,10 +209,10 @@ def test_get_dataset_info_bq_client_creation(mock_get_bigquery_client):
   assert (
       mock_get_bigquery_client.call_args.kwargs["credentials"] == bq_credentials
   )
-  assert (
-      mock_get_bigquery_client.call_args.kwargs["user_agent"]
-      == application_name
-  )
+  assert mock_get_bigquery_client.call_args.kwargs["user_agent"] == [
+      application_name,
+      "get_dataset_info",
+  ]
 
 
 @mock.patch(
@@ -235,10 +235,10 @@ def test_list_table_ids_bq_client_creation(mock_get_bigquery_client):
   assert (
       mock_get_bigquery_client.call_args.kwargs["credentials"] == bq_credentials
   )
-  assert (
-      mock_get_bigquery_client.call_args.kwargs["user_agent"]
-      == application_name
-  )
+  assert mock_get_bigquery_client.call_args.kwargs["user_agent"] == [
+      application_name,
+      "list_table_ids",
+  ]
 
 
 @mock.patch(
@@ -262,7 +262,33 @@ def test_get_table_info_bq_client_creation(mock_get_bigquery_client):
   assert (
       mock_get_bigquery_client.call_args.kwargs["credentials"] == bq_credentials
   )
-  assert (
-      mock_get_bigquery_client.call_args.kwargs["user_agent"]
-      == application_name
+  assert mock_get_bigquery_client.call_args.kwargs["user_agent"] == [
+      application_name,
+      "get_table_info",
+  ]
+
+
+@mock.patch(
+    "google.adk.tools.bigquery.client.get_bigquery_client", autospec=True
+)
+def test_get_job_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_job_id = "my_job_id"
+  bq_credentials = mock.create_autospec(Credentials, instance=True)
+  application_name = "my-agent"
+  tool_settings = BigQueryToolConfig(application_name=application_name)
+
+  metadata_tool.get_job_info(
+      bq_project, bq_job_id, bq_credentials, tool_settings
   )
+  mock_get_bigquery_client.assert_called_once()
+  assert len(mock_get_bigquery_client.call_args.kwargs) == 4
+  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,
+      "get_job_info",
+  ]

tests/unittests/tools/bigquery/test_bigquery_query_tool.py

@@ -1149,10 +1149,10 @@ def test_execute_sql_bq_client_creation(mock_get_bigquery_client):
   assert len(mock_get_bigquery_client.call_args.kwargs) == 4
   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
-  )
+  assert mock_get_bigquery_client.call_args.kwargs["user_agent"] == [
+      application_name,
+      "execute_sql",
+  ]
 
 
 def test_execute_sql_unexpected_project_id():
@@ -1177,10 +1177,10 @@ def test_execute_sql_unexpected_project_id():
   }
 
 
-# AI.Forecast calls execute_sql with a specific query statement. We need to
-# test that the query is properly constructed and call execute_sql with the
+# AI.Forecast calls _execute_sql with a specific query statement. We need to
+# test that the query is properly constructed and call _execute_sql with the
 # correct parameters exactly once.
-@mock.patch("google.adk.tools.bigquery.query_tool.execute_sql", autospec=True)
+@mock.patch("google.adk.tools.bigquery.query_tool._execute_sql", autospec=True)
 def test_forecast_with_table_id(mock_execute_sql):
   mock_credentials = mock.MagicMock(spec=Credentials)
   mock_settings = BigQueryToolConfig()
@@ -1210,18 +1210,19 @@ def test_forecast_with_table_id(mock_execute_sql):
   )
   """
   mock_execute_sql.assert_called_once_with(
-      "test-project",
-      expected_query,
-      mock_credentials,
-      mock_settings,
-      mock_tool_context,
+      project_id="test-project",
+      query=expected_query,
+      credentials=mock_credentials,
+      settings=mock_settings,
+      tool_context=mock_tool_context,
+      caller_id="forecast",
   )
 
 
-# AI.Forecast calls execute_sql with a specific query statement. We need to
-# test that the query is properly constructed and call execute_sql with the
+# AI.Forecast calls _execute_sql with a specific query statement. We need to
+# test that the query is properly constructed and call _execute_sql with the
 # correct parameters exactly once.
-@mock.patch("google.adk.tools.bigquery.query_tool.execute_sql", autospec=True)
+@mock.patch("google.adk.tools.bigquery.query_tool._execute_sql", autospec=True)
 def test_forecast_with_query_statement(mock_execute_sql):
   mock_credentials = mock.MagicMock(spec=Credentials)
   mock_settings = BigQueryToolConfig()
@@ -1249,11 +1250,12 @@ def test_forecast_with_query_statement(mock_execute_sql):
   )
   """
   mock_execute_sql.assert_called_once_with(
-      "test-project",
-      expected_query,
-      mock_credentials,
-      mock_settings,
-      mock_tool_context,
+      project_id="test-project",
+      query=expected_query,
+      credentials=mock_credentials,
+      settings=mock_settings,
+      tool_context=mock_tool_context,
+      caller_id="forecast",
   )
 
 
@@ -1277,10 +1279,10 @@ def test_forecast_with_invalid_id_cols():
   assert "All elements in id_cols must be strings." in result["error_details"]
 
 
-# analyze_contribution calls execute_sql twice. We need to test that the
-# queries are properly constructed and call execute_sql with the correct
+# analyze_contribution calls _execute_sql twice. We need to test that the
+# queries are properly constructed and call _execute_sql with the correct
 # parameters exactly twice.
-@mock.patch("google.adk.tools.bigquery.query_tool.execute_sql", autospec=True)
+@mock.patch("google.adk.tools.bigquery.query_tool._execute_sql", autospec=True)
 @mock.patch("uuid.uuid4", autospec=True)
 def test_analyze_contribution_with_table_id(mock_uuid, mock_execute_sql):
   """Test analyze_contribution tool invocation with a table id."""
@@ -1313,25 +1315,27 @@ def test_analyze_contribution_with_table_id(mock_uuid, mock_execute_sql):
 
   assert mock_execute_sql.call_count == 2
   mock_execute_sql.assert_any_call(
-      "test-project",
-      expected_create_model_query,
-      mock_credentials,
-      mock_settings,
-      mock_tool_context,
+      project_id="test-project",
+      query=expected_create_model_query,
+      credentials=mock_credentials,
+      settings=mock_settings,
+      tool_context=mock_tool_context,
+      caller_id="analyze_contribution",
   )
   mock_execute_sql.assert_any_call(
-      "test-project",
-      expected_get_insights_query,
-      mock_credentials,
-      mock_settings,
-      mock_tool_context,
+      project_id="test-project",
+      query=expected_get_insights_query,
+      credentials=mock_credentials,
+      settings=mock_settings,
+      tool_context=mock_tool_context,
+      caller_id="analyze_contribution",
   )
 
 
-# analyze_contribution calls execute_sql twice. We need to test that the
-# queries are properly constructed and call execute_sql with the correct
+# analyze_contribution calls _execute_sql twice. We need to test that the
+# queries are properly constructed and call _execute_sql with the correct
 # parameters exactly twice.
-@mock.patch("google.adk.tools.bigquery.query_tool.execute_sql", autospec=True)
+@mock.patch("google.adk.tools.bigquery.query_tool._execute_sql", autospec=True)
 @mock.patch("uuid.uuid4", autospec=True)
 def test_analyze_contribution_with_query_statement(mock_uuid, mock_execute_sql):
   """Test analyze_contribution tool invocation with a query statement."""
@@ -1365,18 +1369,20 @@ def test_analyze_contribution_with_query_statement(mock_uuid, mock_execute_sql):
 
   assert mock_execute_sql.call_count == 2
   mock_execute_sql.assert_any_call(
-      "test-project",
-      expected_create_model_query,
-      mock_credentials,
-      mock_settings,
-      mock_tool_context,
+      project_id="test-project",
+      query=expected_create_model_query,
+      credentials=mock_credentials,
+      settings=mock_settings,
+      tool_context=mock_tool_context,
+      caller_id="analyze_contribution",
   )
   mock_execute_sql.assert_any_call(
-      "test-project",
-      expected_get_insights_query,
-      mock_credentials,
-      mock_settings,
-      mock_tool_context,
+      project_id="test-project",
+      query=expected_get_insights_query,
+      credentials=mock_credentials,
+      settings=mock_settings,
+      tool_context=mock_tool_context,
+      caller_id="analyze_contribution",
   )
 
 
@@ -1404,10 +1410,10 @@ def test_analyze_contribution_with_invalid_dimension_id_cols():
   )
 
 
-# detect_anomalies calls execute_sql twice. We need to test that
-# the queries are properly constructed and call execute_sql with the correct
+# detect_anomalies calls _execute_sql twice. We need to test that
+# the queries are properly constructed and call _execute_sql with the correct
 # parameters exactly twice.
-@mock.patch("google.adk.tools.bigquery.query_tool.execute_sql", autospec=True)
+@mock.patch("google.adk.tools.bigquery.query_tool._execute_sql", autospec=True)
 @mock.patch("uuid.uuid4", autospec=True)
 def test_detect_anomalies_with_table_id(mock_uuid, mock_execute_sql):
   """Test time series anomaly detection tool invocation with a table id."""
@@ -1440,25 +1446,27 @@ def test_detect_anomalies_with_table_id(mock_uuid, mock_execute_sql):
 
   assert mock_execute_sql.call_count == 2
   mock_execute_sql.assert_any_call(
-      "test-project",
-      expected_create_model_query,
-      mock_credentials,
-      mock_settings,
-      mock_tool_context,
+      project_id="test-project",
+      query=expected_create_model_query,
+      credentials=mock_credentials,
+      settings=mock_settings,
+      tool_context=mock_tool_context,
+      caller_id="detect_anomalies",
   )
   mock_execute_sql.assert_any_call(
-      "test-project",
-      expected_anomaly_detection_query,
-      mock_credentials,
-      mock_settings,
-      mock_tool_context,
+      project_id="test-project",
+      query=expected_anomaly_detection_query,
+      credentials=mock_credentials,
+      settings=mock_settings,
+      tool_context=mock_tool_context,
+      caller_id="detect_anomalies",
   )
 
 
-# detect_anomalies calls execute_sql twice. We need to test that
-# the queries are properly constructed and call execute_sql with the correct
+# detect_anomalies calls _execute_sql twice. We need to test that
+# the queries are properly constructed and call _execute_sql with the correct
 # parameters exactly twice.
-@mock.patch("google.adk.tools.bigquery.query_tool.execute_sql", autospec=True)
+@mock.patch("google.adk.tools.bigquery.query_tool._execute_sql", autospec=True)
 @mock.patch("uuid.uuid4", autospec=True)
 def test_detect_anomalies_with_custom_params(mock_uuid, mock_execute_sql):
   """Test time series anomaly detection tool invocation with a table id."""
@@ -1494,25 +1502,27 @@ def test_detect_anomalies_with_custom_params(mock_uuid, mock_execute_sql):
 
   assert mock_execute_sql.call_count == 2
   mock_execute_sql.assert_any_call(
-      "test-project",
-      expected_create_model_query,
-      mock_credentials,
-      mock_settings,
-      mock_tool_context,
+      project_id="test-project",
+      query=expected_create_model_query,
+      credentials=mock_credentials,
+      settings=mock_settings,
+      tool_context=mock_tool_context,
+      caller_id="detect_anomalies",
   )
   mock_execute_sql.assert_any_call(
-      "test-project",
-      expected_anomaly_detection_query,
-      mock_credentials,
-      mock_settings,
-      mock_tool_context,
+      project_id="test-project",
+      query=expected_anomaly_detection_query,
+      credentials=mock_credentials,
+      settings=mock_settings,
+      tool_context=mock_tool_context,
+      caller_id="detect_anomalies",
   )
 
 
-# detect_anomalies calls execute_sql twice. We need to test that
-# the queries are properly constructed and call execute_sql with the correct
+# detect_anomalies calls _execute_sql twice. We need to test that
+# the queries are properly constructed and call _execute_sql with the correct
 # parameters exactly twice.
-@mock.patch("google.adk.tools.bigquery.query_tool.execute_sql", autospec=True)
+@mock.patch("google.adk.tools.bigquery.query_tool._execute_sql", autospec=True)
 @mock.patch("uuid.uuid4", autospec=True)
 def test_detect_anomalies_on_target_table(mock_uuid, mock_execute_sql):
   """Test time series anomaly detection tool with target data is provided."""
@@ -1550,22 +1560,24 @@ def test_detect_anomalies_on_target_table(mock_uuid, mock_execute_sql):
 
   assert mock_execute_sql.call_count == 2
   mock_execute_sql.assert_any_call(
-      "test-project",
-      expected_create_model_query,
-      mock_credentials,
-      mock_settings,
-      mock_tool_context,
+      project_id="test-project",
+      query=expected_create_model_query,
+      credentials=mock_credentials,
+      settings=mock_settings,
+      tool_context=mock_tool_context,
+      caller_id="detect_anomalies",
   )
   mock_execute_sql.assert_any_call(
-      "test-project",
-      expected_anomaly_detection_query,
-      mock_credentials,
-      mock_settings,
-      mock_tool_context,
+      project_id="test-project",
+      query=expected_anomaly_detection_query,
+      credentials=mock_credentials,
+      settings=mock_settings,
+      tool_context=mock_tool_context,
+      caller_id="detect_anomalies",
   )
 
 
-def test_detect_anomalies__with_invalid_id_cols():
+def test_detect_anomalies_with_invalid_id_cols():
   """Test time series anomaly detection tool invocation with invalid times_series_id_cols."""
   mock_credentials = mock.MagicMock(spec=Credentials)
   mock_settings = BigQueryToolConfig()
@@ -1587,3 +1599,123 @@ def test_detect_anomalies__with_invalid_id_cols():
       "All elements in times_series_id_cols must be strings."
       in result["error_details"]
   )
+
+
+@pytest.mark.parametrize(
+    ("write_mode", "dry_run", "query_call_count", "query_and_wait_call_count"),
+    [
+        pytest.param(WriteMode.ALLOWED, False, 0, 1, id="write-allowed"),
+        pytest.param(WriteMode.ALLOWED, True, 1, 0, id="write-allowed-dry-run"),
+        pytest.param(WriteMode.BLOCKED, False, 1, 1, id="write-blocked"),
+        pytest.param(WriteMode.BLOCKED, True, 2, 0, id="write-blocked-dry-run"),
+        pytest.param(WriteMode.PROTECTED, False, 2, 1, id="write-protected"),
+        pytest.param(
+            WriteMode.PROTECTED, True, 3, 0, id="write-protected-dry-run"
+        ),
+    ],
+)
+def test_execute_sql_job_labels(
+    write_mode, dry_run, query_call_count, query_and_wait_call_count
+):
+  """Test execute_sql tool for job label."""
+  project = "my_project"
+  query = "SELECT 123 AS num"
+  statement_type = "SELECT"
+  credentials = mock.create_autospec(Credentials, instance=True)
+  tool_settings = BigQueryToolConfig(write_mode=write_mode)
+  tool_context = mock.create_autospec(ToolContext, instance=True)
+  tool_context.state.get.return_value = None
+
+  with mock.patch("google.cloud.bigquery.Client", autospec=False) as Client:
+    bq_client = Client.return_value
+
+    query_job = mock.create_autospec(bigquery.QueryJob)
+    query_job.statement_type = statement_type
+    bq_client.query.return_value = query_job
+
+    execute_sql(
+        project,
+        query,
+        credentials,
+        tool_settings,
+        tool_context,
+        dry_run=dry_run,
+    )
+
+    assert bq_client.query.call_count == query_call_count
+    assert bq_client.query_and_wait.call_count == query_and_wait_call_count
+    for call_args_list in [
+        bq_client.query.call_args_list,
+        bq_client.query_and_wait.call_args_list,
+    ]:
+      for call_args in call_args_list:
+        _, mock_kwargs = call_args
+        assert mock_kwargs["job_config"].labels == {
+            "adk-bigquery-tool": "execute_sql"
+        }
+
+
+@pytest.mark.parametrize(
+    ("tool_call", "expected_label"),
+    [
+        pytest.param(
+            lambda tool_context: forecast(
+                project_id="test-project",
+                history_data="SELECT * FROM `test-dataset.test-table`",
+                timestamp_col="ts_col",
+                data_col="data_col",
+                credentials=mock.create_autospec(Credentials, instance=True),
+                settings=BigQueryToolConfig(write_mode=WriteMode.ALLOWED),
+                tool_context=tool_context,
+            ),
+            "forecast",
+            id="forecast",
+        ),
+        pytest.param(
+            lambda tool_context: analyze_contribution(
+                project_id="test-project",
+                input_data="test-dataset.test-table",
+                dimension_id_cols=["dim1", "dim2"],
+                contribution_metric="SUM(metric)",
+                is_test_col="is_test",
+                credentials=mock.create_autospec(Credentials, instance=True),
+                settings=BigQueryToolConfig(write_mode=WriteMode.ALLOWED),
+                tool_context=tool_context,
+            ),
+            "analyze_contribution",
+            id="analyze-contribution",
+        ),
+        pytest.param(
+            lambda tool_context: detect_anomalies(
+                project_id="test-project",
+                history_data="SELECT * FROM `test-dataset.test-table`",
+                times_series_timestamp_col="ts_timestamp",
+                times_series_data_col="ts_data",
+                credentials=mock.create_autospec(Credentials, instance=True),
+                settings=BigQueryToolConfig(write_mode=WriteMode.ALLOWED),
+                tool_context=tool_context,
+            ),
+            "detect_anomalies",
+            id="detect-anomalies",
+        ),
+    ],
+)
+def test_ml_tool_job_labels(tool_call, expected_label):
+  """Test ML tools for job label."""
+
+  with mock.patch("google.cloud.bigquery.Client", autospec=False) as Client:
+    bq_client = Client.return_value
+
+    tool_context = mock.create_autospec(ToolContext, instance=True)
+    tool_context.state.get.return_value = None
+    tool_call(tool_context)
+
+    for call_args_list in [
+        bq_client.query.call_args_list,
+        bq_client.query_and_wait.call_args_list,
+    ]:
+      for call_args in call_args_list:
+        _, mock_kwargs = call_args
+        assert mock_kwargs["job_config"].labels == {
+            "adk-bigquery-tool": expected_label
+        }