fix: use pessimistic json encoder in SQL Lab (#28266)

This commit is contained in:
Maxime Beauchemin 2024-05-06 12:23:50 -07:00 committed by GitHub
parent d1e13ab3e4
commit c10cee3a39
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 50 additions and 9 deletions

View File

@ -340,15 +340,15 @@ class SqlLabRestApi(BaseSupersetApi):
key = params.get("key")
rows = params.get("rows")
result = SqlExecutionResultsCommand(key=key, rows=rows).run()
# return the result without special encoding
return json_success(
json.dumps(
# Using pessimistic json serialization since some database drivers can return
# unserializeable types at times
payload = json.dumps(
result,
default=utils.json_iso_dttm_ser,
default=utils.pessimistic_json_iso_dttm_ser,
ignore_nan=True,
),
200,
)
return json_success(payload, 200)
@expose("/execute/", methods=("POST",))
@protect()

View File

@ -507,8 +507,8 @@ def json_iso_dttm_ser(obj: Any, pessimistic: bool = False) -> Any:
return base_json_conv(obj)
except TypeError as ex:
if pessimistic:
logger.error("Failed to serialize %s", obj)
return f"Unserializable [{type(obj)}]"
raise ex

View File

@ -14,6 +14,8 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
import datetime
import json
import os
from dataclasses import dataclass
from typing import Any, Optional
@ -31,8 +33,10 @@ from superset.utils.core import (
generic_find_fk_constraint_name,
get_datasource_full_name,
is_test,
json_iso_dttm_ser,
normalize_dttm_col,
parse_boolean_string,
pessimistic_json_iso_dttm_ser,
QueryObjectFilterClause,
remove_extra_adhoc_filters,
)
@ -396,3 +400,40 @@ def test_get_datasource_full_name():
get_datasource_full_name("db", "table", "catalog", None)
== "[db].[catalog].[table]"
)
def test_json_iso_dttm_ser():
data = {
"datetime": datetime.datetime(2021, 1, 1, 0, 0, 0),
"date": datetime.date(2021, 1, 1),
}
json_str = json.dumps(data, default=json_iso_dttm_ser)
reloaded_data = json.loads(json_str)
assert reloaded_data["datetime"] == "2021-01-01T00:00:00"
assert reloaded_data["date"] == "2021-01-01"
def test_pessimistic_json_iso_dttm_ser():
data = {
"datetime": datetime.datetime(2021, 1, 1, 0, 0, 0),
"date": datetime.date(2021, 1, 1),
"UNSERIALIZABLE": MagicMock(),
}
json_str = json.dumps(data, default=pessimistic_json_iso_dttm_ser)
reloaded_data = json.loads(json_str)
assert reloaded_data["datetime"] == "2021-01-01T00:00:00"
assert reloaded_data["date"] == "2021-01-01"
assert (
reloaded_data["UNSERIALIZABLE"]
== "Unserializable [<class 'unittest.mock.MagicMock'>]"
)
def test_pessimistic_json_iso_dttm_ser_nonutf8():
data = {
"INVALID_UTF8_BYTES": b"\xff",
}
assert isinstance(data["INVALID_UTF8_BYTES"], bytes)
json_str = json.dumps(data, default=pessimistic_json_iso_dttm_ser)
reloaded_data = json.loads(json_str)
assert reloaded_data["INVALID_UTF8_BYTES"] == "[bytes]"