feat: allow setting max_billed_bytes in BigQuery tools config

This will allow users to configure a limit to access in ADK tools on the charges for queries.

PiperOrigin-RevId: 831921163
This commit is contained in:
Google Team Member
2025-11-13 10:42:56 -08:00
committed by Copybara-Service
parent 22eb7e5b06
commit ffbb0b37e1
4 changed files with 72 additions and 4 deletions
@@ -1826,3 +1826,26 @@ def test_execute_sql_no_truncation():
# Check no truncation flag when fewer rows than limit
assert result["status"] == "SUCCESS"
assert "result_is_likely_truncated" not in result
def test_execute_sql_maximum_bytes_billed_config():
"""Test execute_sql tool respects maximum_bytes_billed from config."""
project = "my_project"
query = "SELECT 123 AS num"
statement_type = "SELECT"
credentials = mock.create_autospec(Credentials, instance=True)
tool_config = BigQueryToolConfig(maximum_bytes_billed=11_000_000)
tool_context = mock.create_autospec(ToolContext, instance=True)
with mock.patch("google.cloud.bigquery.Client", autospec=False) as Client:
bq_client = Client.return_value
query_job = mock.create_autospec(bigquery.QueryJob)
query_job.statement_type = statement_type
bq_client.query.return_value = query_job
execute_sql(project, query, credentials, tool_config, tool_context)
# Check that maximum_bytes_billed was called with config value
bq_client.query_and_wait.assert_called_once()
call_args = bq_client.query_and_wait.call_args
assert call_args.kwargs["job_config"].maximum_bytes_billed == 11_000_000