fix: utf-16 json encoder support (#28486)

Co-authored-by: Eyal Ezer <[email protected]>
apache · May 15, 2024 · 5f714b7 · 5f714b7
1 parent f420005
commit 5f714b7
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 8 deletions.
diff --git a/superset/sqllab/execution_context_convertor.py b/superset/sqllab/execution_context_convertor.py
@@ -52,13 +52,22 @@ def set_payload(
 
     def serialize_payload(self) -> str:
         if self._exc_status == SqlJsonExecutionStatus.HAS_RESULTS:
-            return json.dumps(
-                apply_display_max_row_configuration_if_require(
-                    self.payload, self._max_row_in_display_configuration
-                ),
-                default=utils.pessimistic_json_iso_dttm_ser,
-                ignore_nan=True,
+            sql_results = apply_display_max_row_configuration_if_require(
+                self.payload, self._max_row_in_display_configuration
             )
+            try:
+                return json.dumps(
+                    sql_results,
+                    default=utils.pessimistic_json_iso_dttm_ser,
+                    ignore_nan=True,
+                )
+            except UnicodeDecodeError:
+                return json.dumps(
+                    sql_results,
+                    default=utils.pessimistic_json_iso_dttm_ser,
+                    ensure_ascii=False,
+                    ignore_nan=True,
+                )
 
         return json.dumps(
             {"query": self.payload}, default=utils.json_int_dttm_ser, ignore_nan=True

diff --git a/superset/utils/core.py b/superset/utils/core.py
@@ -482,7 +482,10 @@ def base_json_conv(obj: Any) -> Any:
         try:
             return obj.decode("utf-8")
         except Exception:  # pylint: disable=broad-except
-            return "[bytes]"
+            try:
+                return obj.decode("utf-16")
+            except Exception:  # pylint: disable=broad-except
+                return "[bytes]"
 
     raise TypeError(f"Unserializable object {obj} of type {type(obj)}")
 

diff --git a/tests/integration_tests/utils_tests.py b/tests/integration_tests/utils_tests.py
@@ -127,7 +127,8 @@ def test_base_json_conv(self):
         assert isinstance(base_json_conv(time()), str)
         assert isinstance(base_json_conv(timedelta(0)), str)
         assert isinstance(base_json_conv(b""), str)
-        assert base_json_conv(bytes("", encoding="utf-16")) == "[bytes]"
+        assert isinstance(base_json_conv(b"\xff\xfe"), str)
+        assert base_json_conv(b"\xff") == "[bytes]"
 
         with pytest.raises(TypeError):
             base_json_conv(np.datetime64())

diff --git a/tests/unit_tests/utils/test_core.py b/tests/unit_tests/utils/test_core.py
@@ -437,3 +437,13 @@ def test_pessimistic_json_iso_dttm_ser_nonutf8():
     json_str = json.dumps(data, default=pessimistic_json_iso_dttm_ser)
     reloaded_data = json.loads(json_str)
     assert reloaded_data["INVALID_UTF8_BYTES"] == "[bytes]"
+
+
+def test_pessimistic_json_iso_dttm_ser_utf16():
+    data = {
+        "VALID_UTF16_BYTES": b"\xff\xfeS0\x930k0a0o0\x16NLu",
+    }
+    assert isinstance(data["VALID_UTF16_BYTES"], bytes)
+    json_str = json.dumps(data, default=pessimistic_json_iso_dttm_ser)
+    reloaded_data = json.loads(json_str)
+    assert reloaded_data["VALID_UTF16_BYTES"] == "こんにちは世界"