From c10cee3a39c4b75273dc051670d20b21f6294151 Mon Sep 17 00:00:00 2001
From: Maxime Beauchemin <maximebeauchemin@gmail.com>
Date: Mon, 6 May 2024 12:23:50 -0700
Subject: [PATCH] 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 24c04dfe5..6d6579c94 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 514cbac75..6b44fda4e 100644
--- a/superset/utils/core.py
+++ b/superset/utils/core.py
@@ -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
 
 
diff --git a/tests/unit_tests/utils/test_core.py b/tests/unit_tests/utils/test_core.py
index 2ebec87c2..d6e8d9183 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_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]"