mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
22eb7e5b06
commit
ffbb0b37e1
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user