feat: Provide location config for BigQuery tools

Right now the tolls are always running against multi-region US by default. With this change the agent builder can scope the tools to data and compute in a particular BigQuery location.

PiperOrigin-RevId: 806473857

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

@@ -29,16 +29,30 @@ def get_bigquery_client(
     *,
     project: Optional[str],
     credentials: Credentials,
+    location: Optional[str] = None,
     user_agent: Optional[str] = None,
 ) -> bigquery.Client:
-  """Get a BigQuery client."""
+  """Get a BigQuery client.
+
+  Args:
+    project: The GCP project ID.
+    credentials: The credentials to use for the request.
+    location: The location of the BigQuery client.
+    user_agent: The user agent to use for the request.
+
+  Returns:
+    A BigQuery client.
+  """
 
   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
+      project=project,
+      credentials=credentials,
+      location=location,
+      client_info=client_info,
   )
 
   return bigquery_client

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

@@ -82,6 +82,15 @@ class BigQueryToolConfig(BaseModel):
   operations (such as query execution) in a specific project.
   """
 
+  location: Optional[str] = None
+  """BigQuery location to use for the data and compute.
+
+  This can be set if the BigQuery tools are expected to process data in a
+  particular BigQuery location. If not set, then location would be automatically
+  determined based on the data location in the query. For all supported
+  locations, see https://cloud.google.com/bigquery/docs/locations.
+  """
+
   @field_validator('application_name')
   @classmethod
   def validate_application_name(cls, v):

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

@@ -50,6 +50,7 @@ def list_dataset_ids(
     bq_client = client.get_bigquery_client(
         project=project_id,
         credentials=credentials,
+        location=settings.location,
         user_agent=settings.application_name,
     )
 
@@ -121,6 +122,7 @@ def get_dataset_info(
     bq_client = client.get_bigquery_client(
         project=project_id,
         credentials=credentials,
+        location=settings.location,
         user_agent=settings.application_name,
     )
     dataset = bq_client.get_dataset(
@@ -159,6 +161,7 @@ def list_table_ids(
     bq_client = client.get_bigquery_client(
         project=project_id,
         credentials=credentials,
+        location=settings.location,
         user_agent=settings.application_name,
     )
 
@@ -281,6 +284,7 @@ def get_table_info(
     bq_client = client.get_bigquery_client(
         project=project_id,
         credentials=credentials,
+        location=settings.location,
         user_agent=settings.application_name,
     )
     return bq_client.get_table(

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

@@ -97,6 +97,7 @@ def execute_sql(
     bq_client = client.get_bigquery_client(
         project=project_id,
         credentials=credentials,
+        location=settings.location,
         user_agent=settings.application_name,
     )
 

tests/unittests/tools/bigquery/test_bigquery_client.py

@@ -23,8 +23,8 @@ from google.auth.exceptions import DefaultCredentialsError
 from google.oauth2.credentials import Credentials
 
 
-def test_bigquery_client_project():
+def test_bigquery_client_default():
-  """Test BigQuery client project."""
+  """Test the default BigQuery client properties."""
   # Trigger the BigQuery client creation
   client = get_bigquery_client(
       project="test-gcp-project",
@@ -33,6 +33,7 @@ def test_bigquery_client_project():
 
   # Verify that the client has the desired project set
   assert client.project == "test-gcp-project"
+  assert client.location is None
 
 
 def test_bigquery_client_project_set_explicit():
@@ -153,3 +154,17 @@ def test_bigquery_client_user_agent_custom():
     }
     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
+  client = get_bigquery_client(
+      project="test-gcp-project",
+      credentials=mock.create_autospec(Credentials, instance=True),
+      location="us-central1",
+  )
+
+  # Verify that the client has the desired project set
+  assert client.project == "test-gcp-project"
+  assert client.location == "us-central1"

tests/unittests/tools/bigquery/test_bigquery_metadata_tool.py

@@ -148,7 +148,7 @@ def test_list_dataset_ids_bq_client_creation(mock_get_bigquery_client):
 
   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 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
@@ -174,7 +174,7 @@ def test_get_dataset_info_bq_client_creation(mock_get_bigquery_client):
       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 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
@@ -200,7 +200,7 @@ def test_list_table_ids_bq_client_creation(mock_get_bigquery_client):
       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 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
@@ -227,7 +227,7 @@ def test_get_table_info_bq_client_creation(mock_get_bigquery_client):
       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 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

tests/unittests/tools/bigquery/test_bigquery_query_tool.py

@@ -1000,7 +1000,7 @@ def test_execute_sql_bq_client_creation(mock_get_bigquery_client):
 
   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 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 (