You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
6cc3d9ddd1
commit
13ff009d34
@@ -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
|
||||
|
||||
|
||||
@@ -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}
|
||||
|
||||
Reference in New Issue
Block a user