Skip to content

Commit e5c0d8b

Browse files
authored
fix(bigquery-analytics): route Storage Write API appends to the dataset's region (#2)
The Storage Write API `AppendRows` streaming RPC does not auto-populate the request-routing header, so the plugin's writes were always routed to the US multiregion. Writes to a dataset in any other region (e.g. northamerica-northeast1) failed with a "session not found" / stream-not-found error and no rows were ever written — which surfaced to users as session_id (and every other column) failing to propagate. Set the `x-goog-request-params: write_stream=<stream>` routing header on the append_rows call, matching what google.cloud.bigquery_storage_v1.writer does internally, so requests reach the region that owns the write stream. US-multiregion behavior is unchanged. Adds a regression test asserting the routing header is passed. Fixes GoogleCloudPlatform/BigQuery-Agent-Analytics-SDK#262
1 parent d0ab39f commit e5c0d8b

2 files changed

Lines changed: 64 additions & 1 deletion

File tree

src/google/adk/plugins/bigquery_agent_analytics_plugin.py

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1171,7 +1171,21 @@ async def requests_iter():
11711171
yield req
11721172

11731173
async def perform_write():
1174-
responses = await self.write_client.append_rows(requests_iter())
1174+
# The AppendRows streaming RPC does not auto-populate the
1175+
# request-routing header, so writes to any region other than
1176+
# the US multiregion fail with a "session not found" /
1177+
# stream-not-found error. Set the routing header explicitly
1178+
# (same as google.cloud.bigquery_storage_v1.writer) so the
1179+
# request reaches the region that owns the write stream.
1180+
responses = await self.write_client.append_rows(
1181+
requests_iter(),
1182+
metadata=(
1183+
(
1184+
"x-goog-request-params",
1185+
f"write_stream={self.write_stream}",
1186+
),
1187+
),
1188+
)
11751189
async for response in responses:
11761190
error = getattr(response, "error", None)
11771191
error_code = getattr(error, "code", None)

tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -635,6 +635,55 @@ async def test_event_denylist(
635635
await asyncio.sleep(0.01)
636636
mock_write_client.append_rows.assert_called_once()
637637

638+
@pytest.mark.asyncio
639+
async def test_append_rows_sets_regional_routing_header(
640+
self,
641+
mock_write_client,
642+
callback_context,
643+
mock_auth_default,
644+
mock_bq_client,
645+
mock_to_arrow_schema,
646+
dummy_arrow_schema,
647+
mock_asyncio_to_thread,
648+
):
649+
"""Regression test for cross-region writes (issue #262).
650+
651+
The Storage Write API streaming AppendRows RPC does not
652+
auto-populate the request-routing header, so writes to a dataset
653+
outside the US multiregion (e.g. northamerica-northeast1) fail with
654+
a "session not found" / stream-not-found error unless the header is
655+
set explicitly. Assert the header is passed to append_rows so the
656+
request reaches the region that owns the write stream.
657+
"""
658+
_ = mock_auth_default
659+
_ = mock_bq_client
660+
config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig()
661+
async with managed_plugin(
662+
PROJECT_ID,
663+
DATASET_ID,
664+
table_id=TABLE_ID,
665+
config=config,
666+
location="northamerica-northeast1",
667+
) as plugin:
668+
await plugin._ensure_started()
669+
mock_write_client.append_rows.reset_mock()
670+
llm_request = llm_request_lib.LlmRequest(
671+
model="gemini-pro",
672+
contents=[types.Content(parts=[types.Part(text="Prompt")])],
673+
)
674+
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
675+
await plugin.before_model_callback(
676+
callback_context=callback_context, llm_request=llm_request
677+
)
678+
await asyncio.sleep(0.01) # Allow background task to run
679+
mock_write_client.append_rows.assert_called_once()
680+
metadata = mock_write_client.append_rows.call_args.kwargs.get("metadata")
681+
assert metadata is not None, "append_rows must receive routing metadata"
682+
assert (
683+
"x-goog-request-params",
684+
f"write_stream={DEFAULT_STREAM_NAME}",
685+
) in tuple(metadata)
686+
638687
@pytest.mark.asyncio
639688
async def test_content_formatter(
640689
self,

0 commit comments

Comments
 (0)