Skip to content

fix(integrations): run BigQuery execute_sql off the event loop - #7161

Open
yumizu-da wants to merge 1 commit into
google:mainfrom
yumizu-da:fix/bigquery-execute-sql-async
Open

yumizu-da wants to merge 1 commit into
google:mainfrom
yumizu-da:fix/bigquery-execute-sql-async

Conversation

@yumizu-da

@yumizu-da yumizu-da commented Sep 17, 2026

Copy link
Copy Markdown

Please ensure you have read the contribution guide before creating a pull request.

Link to Issue or Description of Change

1. Link to an existing issue (if applicable):

2. Or, if no issue exists, describe the change:

Problem:
execute_sql in integrations/bigquery/query_tool.py is a plain def calling a blocking BigQuery client, and ADK awaits a sync tool inline on the running loop. Nothing else in the process runs until the query returns, so on a server handling more than one session the others stall behind a call that is only waiting on the network. Spanner had the same problem and was fixed in 1dbceccf; Bigtable in 72f3e7e1. BigQuery has no asyncio.to_thread anywhere.

Solution:
Make execute_sql a coroutine and run the existing _execute_sql helper with await asyncio.to_thread(...), the same pattern Spanner and Bigtable use. _execute_sql stays sync because forecast, analyze_contribution and detect_anomalies call it directly; those three block the loop too, but that's a separate change. get_execute_sql and the docstring-swapping helper return Awaitable[dict] now, and the write-mode variants are clones of execute_sql.__code__, so they become coroutines with it.

Testing Plan

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.

One regression test added to tests/unittests/integrations/bigquery/test_bigquery_query_tool.py:

  • test_execute_sql_leaves_the_event_loop_free_while_querying - the stubbed client waits inside query_and_wait on a threading.Event that only a concurrently scheduled asyncio task can set. If the query ran on the loop thread the release could never happen (bounded by a 10s timeout, so a regression fails instead of hanging CI).

The rest of the diff converts the existing execute_sql tests to coroutines. test_tool_call_doesnt_change_global_settings and test_tool_call_doesnt_mutate_job_labels are parametrized over both execute_sql and the three ML tools that stay sync, so they await only when the result is awaitable.

$ pytest tests/unittests/integrations/bigquery -q
229 passed, 5 warnings, 4 subtests passed in 3.75s

tox across 3.10-3.14 is clean apart from the two test_import_loading.py allowlist cases on 3.11-3.13, which fail identically on unchanged main here.

Manual End-to-End (E2E) Tests:

Run the repro script from #7160: it stands in a BigQuery client that takes 2 seconds and counts how often a 100 ms heartbeat task gets to run while the query is in flight. Before this change the heartbeat does not tick at all until the query returns; with it the interval holds.

# before, on main at 7ae1c9b0
heartbeat ran 3 times
longest gap between heartbeats: 2.108s
result: {'status': 'SUCCESS', 'rows': [{'num': 123}]}

# with the fix
heartbeat ran 23 times
longest gap between heartbeats: 0.102s
result: {'status': 'SUCCESS', 'rows': [{'num': 123}]}

Checklist

  • I have read the CONTRIBUTING.md document.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end.
  • Any dependent changes have been merged and published in downstream modules.

Additional context

Only execute_sql here. The other BigQuery tools still run the blocking client on the loop and can follow the same pattern separately.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

BigQuery execute_sql blocks the event loop for the whole query

2 participants