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
This commit is contained in:
Google Team Member
2025-07-21 23:16:28 -07:00
committed by Copybara-Service
parent 6cc3d9ddd1
commit 13ff009d34
3 changed files with 161 additions and 1 deletions
@@ -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
+15 -1
View File
@@ -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
@@ -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}