(superset) 01/04: fix: use pessimistic json encoder in SQL Lab (#28266)

[mi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

michaelsmolina pushed a commit to branch 4.0
in repository https://gitbox.apache.org/repos/asf/superset.git

commit d379184b759af791263cdc197c4a2c584b476c30
Author: Maxime Beauchemin <ma...@gmail.com>
AuthorDate: Mon May 6 12:23:50 2024 -0700

    fix: use pessimistic json encoder in SQL Lab (#28266)
---
 superset/sqllab/api.py              | 16 +++++++--------
 superset/utils/core.py              |  2 +-
 tests/unit_tests/utils/test_core.py | 41 +++++++++++++++++++++++++++++++++++++
 3 files changed, 50 insertions(+), 9 deletions(-)

diff --git a/superset/sqllab/api.py b/superset/sqllab/api.py
index eeb95e6aad..f5e80ab24b 100644
--- a/superset/sqllab/api.py
+++ b/superset/sqllab/api.py
@@ -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(
-                result,
-                default=utils.json_iso_dttm_ser,
-                ignore_nan=True,
-            ),
-            200,
+
+        # Using pessimistic json serialization since some database drivers can return
+        # unserializeable types at times
+        payload = json.dumps(
+            result,
+            default=utils.pessimistic_json_iso_dttm_ser,
+            ignore_nan=True,
         )
+        return json_success(payload, 200)
 
     @expose("/execute/", methods=("POST",))
     @protect()
diff --git a/superset/utils/core.py b/superset/utils/core.py
index df12aafe07..f605ef99c1 100644
--- a/superset/utils/core.py
+++ b/superset/utils/core.py
@@ -505,8 +505,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
 
 
diff --git a/tests/unit_tests/utils/test_core.py b/tests/unit_tests/utils/test_core.py
index 3e73637dad..4bb8f95102 100644
--- a/tests/unit_tests/utils/test_core.py
+++ b/tests/unit_tests/utils/test_core.py
@@ -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_constraint_name,
     generic_find_fk_constraint_name,
     is_test,
+    json_iso_dttm_ser,
     normalize_dttm_col,
     parse_boolean_string,
+    pessimistic_json_iso_dttm_ser,
     QueryObjectFilterClause,
     remove_extra_adhoc_filters,
 )
@@ -371,3 +375,40 @@ def test_generic_find_fk_constraint_none_exist():
     )
 
     assert result is None
+
+
+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]"