From 4c00b86e33d51b6a728180d8accb81f7a8aa4fc2 Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Fri, 12 Sep 2025 16:51:39 -0700 Subject: [PATCH] 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 | 18 ++++++++++++++++-- src/google/adk/tools/bigquery/config.py | 9 +++++++++ .../adk/tools/bigquery/metadata_tool.py | 4 ++++ src/google/adk/tools/bigquery/query_tool.py | 1 + .../tools/bigquery/test_bigquery_client.py | 19 +++++++++++++++++-- .../bigquery/test_bigquery_metadata_tool.py | 8 ++++---- .../bigquery/test_bigquery_query_tool.py | 2 +- 7 files changed, 52 insertions(+), 9 deletions(-) diff --git a/src/google/adk/tools/bigquery/client.py b/src/google/adk/tools/bigquery/client.py index 328afbb6..b23626f8 100644 --- a/src/google/adk/tools/bigquery/client.py +++ b/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 diff --git a/src/google/adk/tools/bigquery/config.py b/src/google/adk/tools/bigquery/config.py index 823c88e7..adaa8234 100644 --- a/src/google/adk/tools/bigquery/config.py +++ b/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): diff --git a/src/google/adk/tools/bigquery/metadata_tool.py b/src/google/adk/tools/bigquery/metadata_tool.py index 44fcdfbb..8cfc95c1 100644 --- a/src/google/adk/tools/bigquery/metadata_tool.py +++ b/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( diff --git a/src/google/adk/tools/bigquery/query_tool.py b/src/google/adk/tools/bigquery/query_tool.py index e66ea3ac..03be2866 100644 --- a/src/google/adk/tools/bigquery/query_tool.py +++ b/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, ) diff --git a/tests/unittests/tools/bigquery/test_bigquery_client.py b/tests/unittests/tools/bigquery/test_bigquery_client.py index 6decb6eb..15e2decd 100644 --- a/tests/unittests/tools/bigquery/test_bigquery_client.py +++ b/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(): - """Test BigQuery client project.""" +def test_bigquery_client_default(): + """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" diff --git a/tests/unittests/tools/bigquery/test_bigquery_metadata_tool.py b/tests/unittests/tools/bigquery/test_bigquery_metadata_tool.py index 9de1654b..2bcd587d 100644 --- a/tests/unittests/tools/bigquery/test_bigquery_metadata_tool.py +++ b/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 diff --git a/tests/unittests/tools/bigquery/test_bigquery_query_tool.py b/tests/unittests/tools/bigquery/test_bigquery_query_tool.py index 48b6e392..95abbe49 100644 --- a/tests/unittests/tools/bigquery/test_bigquery_query_tool.py +++ b/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 (