feat: Make MAX_DOWNLOADED_QUERY_RESULT_ROWS configurable

- Add max_downloaded_rows field to BigQueryToolConfig with default value 50
- Add optional max_rows parameter to execute_sql function
- Parameter takes precedence over config when both provided
- Maintain backward compatibility with existing code
- Add comprehensive tests for config and parameter functionality
- Update truncation logic to use resolved max_downloaded_rows value

Author: lupuletic

src/google/adk/tools/bigquery/config.py

@@ -54,3 +54,10 @@ class BigQueryToolConfig(BaseModel):
   By default, the tool will allow only read operations. This behaviour may
   change in future versions.
   """
+
+  max_downloaded_rows: int = 50
+  """Maximum number of rows to download from query results.
+
+  By default, limits query results to 50 rows to prevent excessive memory usage.
+  Set to a higher value if you need more rows returned.
+  """

src/google/adk/tools/bigquery/query_tool.py

@@ -27,7 +27,6 @@
 from .config import BigQueryToolConfig
 from .config import WriteMode
 
-MAX_DOWNLOADED_QUERY_RESULT_ROWS = 50
 BIGQUERY_SESSION_INFO_KEY = "bigquery_session_info"
 
 
@@ -37,6 +36,7 @@ def execute_sql(
     credentials: Credentials,
     config: BigQueryToolConfig,
     tool_context: ToolContext,
+    max_rows: int | None = None,
 ) -> dict:
   """Run a BigQuery or BigQuery ML SQL query in the project and return the result.
 
@@ -160,7 +160,7 @@ def execute_sql(
         query,
         job_config=job_config,
         project=project_id,
-        max_results=MAX_DOWNLOADED_QUERY_RESULT_ROWS,
+        max_results=max_rows or (config.max_downloaded_rows if config else 50),
     )
     rows = []
     for row in row_iterator:
@@ -175,9 +175,10 @@ def execute_sql(
       rows.append(row_values)
 
     result = {"status": "SUCCESS", "rows": rows}
+    max_downloaded_rows = max_rows or (config.max_downloaded_rows if config else 50)
     if (
-        MAX_DOWNLOADED_QUERY_RESULT_ROWS is not None
-        and len(rows) == MAX_DOWNLOADED_QUERY_RESULT_ROWS
+        max_downloaded_rows is not None
+        and len(rows) == max_downloaded_rows
     ):
       result["result_is_likely_truncated"] = True
     return result

tests/unittests/tools/bigquery/test_bigquery_query_tool.py

@@ -973,3 +973,86 @@ def test_execute_sql_result_dtype(
   # Test the tool worked without invoking default auth
   result = execute_sql(project, query, credentials, tool_config, tool_context)
   assert result == {"status": "SUCCESS", "rows": tool_result_rows}
+
+
+def test_execute_sql_max_rows_config():
+  """Test execute_sql tool respects max_downloaded_rows from config."""
+  project = "my_project"
+  query = "SELECT 123 AS num"
+  statement_type = "SELECT"
+  query_result = [{"num": i} for i in range(20)]  # 20 rows
+  credentials = mock.create_autospec(Credentials, instance=True)
+  tool_config = BigQueryToolConfig(max_downloaded_rows=10)
+  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
+    bq_client.query_and_wait.return_value = query_result[:10]
+
+    result = execute_sql(project, query, credentials, tool_config, tool_context)
+    
+    # Check that max_results 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["max_results"] == 10
+    
+    # Check truncation flag is set
+    assert result["status"] == "SUCCESS"
+    assert result["result_is_likely_truncated"] is True
+
+
+def test_execute_sql_max_rows_parameter():
+  """Test execute_sql tool respects max_rows parameter."""
+  project = "my_project"
+  query = "SELECT 123 AS num"
+  statement_type = "SELECT"
+  query_result = [{"num": i} for i in range(20)]  # 20 rows
+  credentials = mock.create_autospec(Credentials, instance=True)
+  tool_config = BigQueryToolConfig(max_downloaded_rows=10)
+  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
+    bq_client.query_and_wait.return_value = query_result[:5]
+
+    # Parameter should override config
+    result = execute_sql(project, query, credentials, tool_config, tool_context, max_rows=5)
+    
+    # Check that max_results was called with parameter value
+    bq_client.query_and_wait.assert_called_once()
+    call_args = bq_client.query_and_wait.call_args
+    assert call_args.kwargs["max_results"] == 5
+    
+    # Check truncation flag is set
+    assert result["status"] == "SUCCESS"
+    assert result["result_is_likely_truncated"] is True
+
+
+def test_execute_sql_no_truncation():
+  """Test execute_sql tool when results are not truncated."""
+  project = "my_project"
+  query = "SELECT 123 AS num"
+  statement_type = "SELECT"
+  query_result = [{"num": i} for i in range(3)]  # Only 3 rows
+  credentials = mock.create_autospec(Credentials, instance=True)
+  tool_config = BigQueryToolConfig(max_downloaded_rows=10)
+  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
+    bq_client.query_and_wait.return_value = query_result
+
+    result = execute_sql(project, query, credentials, tool_config, tool_context)
+    
+    # Check no truncation flag when fewer rows than limit
+    assert result["status"] == "SUCCESS"
+    assert "result_is_likely_truncated" not in result

tests/unittests/tools/bigquery/test_bigquery_tool_config.py

@@ -25,3 +25,17 @@ def test_bigquery_tool_config_experimental_warning():
       match="Config defaults may have breaking change in the future.",
   ):
     BigQueryToolConfig()
+
+
+def test_bigquery_tool_config_max_downloaded_rows_default():
+  """Test BigQueryToolConfig max_downloaded_rows default value."""
+  with pytest.warns(UserWarning):
+    config = BigQueryToolConfig()
+  assert config.max_downloaded_rows == 50
+
+
+def test_bigquery_tool_config_max_downloaded_rows_custom():
+  """Test BigQueryToolConfig max_downloaded_rows custom value."""
+  with pytest.warns(UserWarning):
+    config = BigQueryToolConfig(max_downloaded_rows=100)
+  assert config.max_downloaded_rows == 100