Skip to content

Commit c263426

Browse files
haiyuan-eng-googlecopybara-github
authored andcommitted
fix: fix dataset location handling in BigQueryAgentAnalyticsPlugin
Remove `location=self.location` from the `bigquery.Client()` constructor. When the client has no default location and `client.query(sql)` is called without an explicit `location` parameter, the BQ API infers the job location from the dataset referenced in the DDL statement. Traced through the BQ Python client: 1. `client.query(sql)` — `location` param defaults to `None`, falls back to `self.location` which is also `None` 2. `_to_query_request()` — when `location is None`, the `"location"` key is **not included** in the API request 3. BQ API — infers location from the dataset referenced in `CREATE OR REPLACE VIEW` **One line of production code changed.** Everything else is test updates. | Operation | Before | After | |---|---|---| | Table CRUD | Works | Same | | Storage Write API | Works | Same | | View creation (non-US dataset) | **Silent failure** | **Works** — BQ infers location | Co-authored-by: Haiyuan Cao <haiyuan@google.com> PiperOrigin-RevId: 905242480
1 parent 7a7083c commit c263426

2 files changed

Lines changed: 98 additions & 1 deletion

File tree

src/google/adk/plugins/bigquery_agent_analytics_plugin.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2184,7 +2184,6 @@ async def _lazy_setup(self, **kwargs) -> None:
21842184
self._executor,
21852185
lambda: bigquery.Client(
21862186
project=self.project_id,
2187-
location=self.location,
21882187
credentials=self._credentials,
21892188
),
21902189
)

tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7246,3 +7246,101 @@ async def test_no_a2a_interaction_for_no_metadata(
72467246

72477247
await asyncio.sleep(0.05)
72487248
assert mock_write_client.append_rows.call_count == 0
7249+
7250+
7251+
# ================================================================
7252+
# TEST CLASS: Dataset location handling (Issue #5476)
7253+
# ================================================================
7254+
class TestDatasetLocationHandling:
7255+
"""Tests that BQ client is created without a default location.
7256+
7257+
When location is omitted from bigquery.Client(), client.query()
7258+
sends no location field in the API request, letting BigQuery
7259+
infer location from the referenced dataset. This prevents
7260+
silent view-creation failures for non-US datasets.
7261+
"""
7262+
7263+
@pytest.mark.asyncio
7264+
async def test_client_created_without_location(
7265+
self,
7266+
mock_auth_default,
7267+
mock_to_arrow_schema,
7268+
mock_asyncio_to_thread,
7269+
):
7270+
"""bigquery.Client is created without a location parameter."""
7271+
with mock.patch.object(bigquery, "Client", autospec=True) as mock_bq_cls:
7272+
mock_bq_cls.return_value.get_table.side_effect = (
7273+
cloud_exceptions.NotFound("table")
7274+
)
7275+
mock_bq_cls.return_value.create_table.return_value = None
7276+
7277+
async with managed_plugin(
7278+
project_id=PROJECT_ID,
7279+
dataset_id=DATASET_ID,
7280+
table_id=TABLE_ID,
7281+
location="europe-west1",
7282+
config=bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
7283+
create_views=False,
7284+
),
7285+
) as plugin:
7286+
await plugin._ensure_started()
7287+
7288+
mock_bq_cls.assert_called_once()
7289+
_, kwargs = mock_bq_cls.call_args
7290+
assert "location" not in kwargs
7291+
7292+
@pytest.mark.asyncio
7293+
async def test_view_query_omits_location(
7294+
self,
7295+
mock_auth_default,
7296+
mock_to_arrow_schema,
7297+
mock_asyncio_to_thread,
7298+
):
7299+
"""View creation DDL queries do not pass an explicit location."""
7300+
with mock.patch.object(bigquery, "Client", autospec=True) as mock_bq_cls:
7301+
mock_client = mock_bq_cls.return_value
7302+
mock_client.get_table.return_value = mock.MagicMock()
7303+
mock_client.query.return_value.result.return_value = None
7304+
7305+
async with managed_plugin(
7306+
project_id=PROJECT_ID,
7307+
dataset_id=DATASET_ID,
7308+
table_id=TABLE_ID,
7309+
config=bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
7310+
create_views=True,
7311+
),
7312+
) as plugin:
7313+
await plugin._ensure_started()
7314+
7315+
assert mock_client.query.call_count > 0
7316+
for call in mock_client.query.call_args_list:
7317+
_, kwargs = call
7318+
# No explicit location — BQ infers from dataset
7319+
assert "location" not in kwargs
7320+
7321+
@pytest.mark.asyncio
7322+
async def test_view_error_still_logged(
7323+
self,
7324+
mock_auth_default,
7325+
mock_to_arrow_schema,
7326+
mock_asyncio_to_thread,
7327+
):
7328+
"""View creation errors are logged but not raised."""
7329+
with mock.patch.object(bigquery, "Client", autospec=True) as mock_bq_cls:
7330+
mock_client = mock_bq_cls.return_value
7331+
mock_client.get_table.return_value = mock.MagicMock()
7332+
mock_client.query.return_value.result.side_effect = Exception(
7333+
"view error"
7334+
)
7335+
7336+
# Should not raise
7337+
async with managed_plugin(
7338+
project_id=PROJECT_ID,
7339+
dataset_id=DATASET_ID,
7340+
table_id=TABLE_ID,
7341+
config=bigquery_agent_analytics_plugin.BigQueryLoggerConfig(
7342+
create_views=True,
7343+
),
7344+
) as plugin:
7345+
await plugin._ensure_started()
7346+
assert plugin._started

0 commit comments

Comments
 (0)