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
This commit is contained in:
Google Team Member
2025-09-12 16:52:16 -07:00
committed by Copybara-Service
parent b3b31a9ffb
commit 4c00b86e33
7 changed files with 52 additions and 9 deletions
@@ -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