Skip to content

Commit fa6e461

Browse files
committed
fix: Address code review comments for BigQuery max rows configuration
- Create private constant _DEFAULT_MAX_DOWNLOADED_QUERY_RESULT_ROWS = 50 - Consolidate max_downloaded_rows logic to avoid duplication - Remove redundant None check in truncation logic
1 parent 1fb693a commit fa6e461

File tree

1 file changed

+4
-6
lines changed

1 file changed

+4
-6
lines changed

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

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
from .config import WriteMode
2929

3030
BIGQUERY_SESSION_INFO_KEY = "bigquery_session_info"
31+
_DEFAULT_MAX_DOWNLOADED_QUERY_RESULT_ROWS = 50
3132

3233

3334
def execute_sql(
@@ -156,11 +157,12 @@ def execute_sql(
156157
if bq_connection_properties
157158
else None
158159
)
160+
max_downloaded_rows = max_rows or (config.max_downloaded_rows if config else _DEFAULT_MAX_DOWNLOADED_QUERY_RESULT_ROWS)
159161
row_iterator = bq_client.query_and_wait(
160162
query,
161163
job_config=job_config,
162164
project=project_id,
163-
max_results=max_rows or (config.max_downloaded_rows if config else 50),
165+
max_results=max_downloaded_rows,
164166
)
165167
rows = []
166168
for row in row_iterator:
@@ -175,11 +177,7 @@ def execute_sql(
175177
rows.append(row_values)
176178

177179
result = {"status": "SUCCESS", "rows": rows}
178-
max_downloaded_rows = max_rows or (config.max_downloaded_rows if config else 50)
179-
if (
180-
max_downloaded_rows is not None
181-
and len(rows) == max_downloaded_rows
182-
):
180+
if len(rows) == max_downloaded_rows:
183181
result["result_is_likely_truncated"] = True
184182
return result
185183
except Exception as ex:

0 commit comments

Comments
 (0)