From 13ff009d34836a80f107cb43a632df15f7c215e4 Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Mon, 21 Jul 2025 23:16:28 -0700 Subject: [PATCH] fix: Handle non-json-serializable values in the `execute_sql` tool This change takes cares of SQL results containing values that are not json serializable (e.g. datetime, bignumeric) by converting them to their string representation. PiperOrigin-RevId: 785719997 --- .../adk/tools/bigquery/metadata_tool.py | 2 + src/google/adk/tools/bigquery/query_tool.py | 16 +- .../bigquery/test_bigquery_query_tool.py | 144 ++++++++++++++++++ 3 files changed, 161 insertions(+), 1 deletion(-) diff --git a/src/google/adk/tools/bigquery/metadata_tool.py b/src/google/adk/tools/bigquery/metadata_tool.py index 64f23d07..d8f6613b 100644 --- a/src/google/adk/tools/bigquery/metadata_tool.py +++ b/src/google/adk/tools/bigquery/metadata_tool.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + from google.auth.credentials import Credentials from google.cloud import bigquery diff --git a/src/google/adk/tools/bigquery/query_tool.py b/src/google/adk/tools/bigquery/query_tool.py index 7406d9a4..cd929b29 100644 --- a/src/google/adk/tools/bigquery/query_tool.py +++ b/src/google/adk/tools/bigquery/query_tool.py @@ -12,7 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import functools +import json import types from typing import Callable @@ -159,7 +162,18 @@ def execute_sql( project=project_id, max_results=MAX_DOWNLOADED_QUERY_RESULT_ROWS, ) - rows = [{key: val for key, val in row.items()} for row in row_iterator] + rows = [] + for row in row_iterator: + row_values = {} + for key, val in row.items(): + try: + # if the json serialization of the value succeeds, use it as is + json.dumps(val) + except: + val = str(val) + row_values[key] = val + rows.append(row_values) + result = {"status": "SUCCESS", "rows": rows} if ( MAX_DOWNLOADED_QUERY_RESULT_ROWS is not None diff --git a/tests/unittests/tools/bigquery/test_bigquery_query_tool.py b/tests/unittests/tools/bigquery/test_bigquery_query_tool.py index c42e3881..18173399 100644 --- a/tests/unittests/tools/bigquery/test_bigquery_query_tool.py +++ b/tests/unittests/tools/bigquery/test_bigquery_query_tool.py @@ -14,11 +14,15 @@ from __future__ import annotations +import datetime +import decimal import os import textwrap from typing import Optional from unittest import mock +import dateutil +import dateutil.relativedelta from google.adk.tools import BaseTool from google.adk.tools.bigquery import BigQueryCredentialsConfig from google.adk.tools.bigquery import BigQueryToolset @@ -829,3 +833,143 @@ def test_execute_sql_no_default_auth( result = execute_sql(project, query, credentials, tool_config, tool_context) assert result == {"status": "SUCCESS", "rows": query_result} mock_default_auth.assert_not_called() + + +@pytest.mark.parametrize( + ("query", "query_result", "tool_result_rows"), + [ + pytest.param( + "SELECT [1,2,3] AS x", + [{"x": [1, 2, 3]}], + [{"x": [1, 2, 3]}], + id="ARRAY", + ), + pytest.param( + "SELECT TRUE AS x", [{"x": True}], [{"x": True}], id="BOOL" + ), + pytest.param( + "SELECT b'Hello World!' AS x", + [{"x": b"Hello World!"}], + [{"x": "b'Hello World!'"}], + id="BYTES", + ), + pytest.param( + "SELECT DATE '2025-07-21' AS x", + [{"x": datetime.date(2025, 7, 21)}], + [{"x": "2025-07-21"}], + id="DATE", + ), + pytest.param( + "SELECT DATETIME '2025-07-21 14:30:45' AS x", + [{"x": datetime.datetime(2025, 7, 21, 14, 30, 45)}], + [{"x": "2025-07-21 14:30:45"}], + id="DATETIME", + ), + pytest.param( + "SELECT ST_GEOGFROMTEXT('POINT(-122.21 47.48)') as x", + [{"x": "POINT(-122.21 47.48)"}], + [{"x": "POINT(-122.21 47.48)"}], + id="GEOGRAPHY", + ), + pytest.param( + "SELECT INTERVAL 10 DAY as x", + [{"x": dateutil.relativedelta.relativedelta(days=10)}], + [{"x": "relativedelta(days=+10)"}], + id="INTERVAL", + ), + pytest.param( + "SELECT JSON_OBJECT('name', 'Alice', 'age', 30) AS x", + [{"x": {"age": 30, "name": "Alice"}}], + [{"x": {"age": 30, "name": "Alice"}}], + id="JSON", + ), + pytest.param("SELECT 1 AS x", [{"x": 1}], [{"x": 1}], id="INT64"), + pytest.param( + "SELECT CAST(1.2 AS NUMERIC) AS x", + [{"x": decimal.Decimal("1.2")}], + [{"x": "1.2"}], + id="NUMERIC", + ), + pytest.param( + "SELECT CAST(1.2 AS BIGNUMERIC) AS x", + [{"x": decimal.Decimal("1.2")}], + [{"x": "1.2"}], + id="BIGNUMERIC", + ), + pytest.param( + "SELECT 1.23 AS x", [{"x": 1.23}], [{"x": 1.23}], id="FLOAT64" + ), + pytest.param( + "SELECT RANGE(DATE '2023-01-01', DATE '2023-01-31') as x", + [{ + "x": { + "start": datetime.date(2023, 1, 1), + "end": datetime.date(2023, 1, 31), + } + }], + [{ + "x": ( + "{'start': datetime.date(2023, 1, 1), 'end':" + " datetime.date(2023, 1, 31)}" + ) + }], + id="RANGE", + ), + pytest.param( + "SELECT 'abc' AS x", [{"x": "abc"}], [{"x": "abc"}], id="STRING" + ), + pytest.param( + "SELECT STRUCT('Alice' AS name, 30 AS age) as x", + [{"x": {"name": "Alice", "age": 30}}], + [{"x": {"name": "Alice", "age": 30}}], + id="STRUCT", + ), + pytest.param( + "SELECT TIME '10:30:45' as x", + [{"x": datetime.time(10, 30, 45)}], + [{"x": "10:30:45"}], + id="TIME", + ), + pytest.param( + "SELECT TIMESTAMP '2025-07-21 10:30:45-07:00' as x", + [{ + "x": datetime.datetime( + 2025, 7, 21, 17, 30, 45, tzinfo=datetime.timezone.utc + ) + }], + [{"x": "2025-07-21 17:30:45+00:00"}], + id="TIMESTAMP", + ), + pytest.param( + "SELECT NULL AS x", [{"x": None}], [{"x": None}], id="NULL" + ), + ], +) +@mock.patch.dict(os.environ, {}, clear=True) +@mock.patch("google.cloud.bigquery.Client.query_and_wait", autospec=True) +@mock.patch("google.cloud.bigquery.Client.query", autospec=True) +def test_execute_sql_result_dtype( + mock_query, mock_query_and_wait, query, query_result, tool_result_rows +): + """Test execute_sql tool invocation for various BigQuery data types. + + See all the supported BigQuery data types at + https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#data_type_list. + """ + project = "my_project" + statement_type = "SELECT" + credentials = mock.create_autospec(Credentials, instance=True) + tool_config = BigQueryToolConfig() + tool_context = mock.create_autospec(ToolContext, instance=True) + + # Simulate the result of query API + query_job = mock.create_autospec(bigquery.QueryJob) + query_job.statement_type = statement_type + mock_query.return_value = query_job + + # Simulate the result of query_and_wait API + mock_query_and_wait.return_value = query_result + + # 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}