refactor: Migration of json utilities from core (apache#28522)

Co-authored-by: Eyal Ezer <eyal.ezer@ge.com>

superset/annotation_layers/annotations/schemas.py

@@ -20,7 +20,7 @@
 from marshmallow.validate import Length
 
 from superset.exceptions import SupersetException
-from superset.utils import core as utils
+from superset.utils import json as json_utils
 
 openapi_spec_methods_override = {
     "get": {"get": {"summary": "Get an annotation layer"}},
@@ -51,7 +51,7 @@
 
 def validate_json(value: Union[bytes, bytearray, str]) -> None:
     try:
-        utils.validate_json(value)
+        json_utils.validate_json(value)
     except SupersetException as ex:
         raise ValidationError("JSON not valid") from ex
 

superset/charts/data/api.py

@@ -21,7 +21,6 @@
 import logging
 from typing import Any, TYPE_CHECKING
 
-import simplejson
 from flask import current_app, g, make_response, request, Response
 from flask_appbuilder.api import expose, protect
 from flask_babel import gettext as _
@@ -47,11 +46,11 @@
 from superset.exceptions import QueryObjectValidationError
 from superset.extensions import event_logger
 from superset.models.sql_lab import Query
+from superset.utils import json as json_utils
 from superset.utils.core import (
     create_zip,
     DatasourceType,
     get_user_id,
-    json_int_dttm_ser,
 )
 from superset.utils.decorators import logs_context
 from superset.views.base import CsvResponse, generate_download_headers, XlsxResponse
@@ -396,9 +395,9 @@ def _process_data(query_data: Any) -> Any:
             )
 
         if result_format == ChartDataResultFormat.JSON:
-            response_data = simplejson.dumps(
+            response_data = json_utils.dumps(
                 {"result": result["queries"]},
-                default=json_int_dttm_ser,
+                default=json_utils.json_int_dttm_ser,
                 ignore_nan=True,
             )
             resp = make_response(response_data, 200)

superset/common/query_object.py

@@ -43,10 +43,10 @@
     get_column_names,
     get_metric_names,
     is_adhoc_metric,
-    json_int_dttm_ser,
     QueryObjectFilterClause,
 )
 from superset.utils.hashing import md5_sha_from_dict
+from superset.utils.json import json_int_dttm_ser
 
 if TYPE_CHECKING:
     from superset.connectors.sqla.models import BaseDatasource

superset/dashboards/schemas.py

@@ -24,7 +24,7 @@
 from superset import security_manager
 from superset.exceptions import SupersetException
 from superset.tags.models import TagType
-from superset.utils import core as utils
+from superset.utils import json as json_utils
 
 get_delete_ids_schema = {"type": "array", "items": {"type": "integer"}}
 get_export_ids_schema = {"type": "array", "items": {"type": "integer"}}
@@ -88,7 +88,7 @@
 
 def validate_json(value: Union[bytes, bytearray, str]) -> None:
     try:
-        utils.validate_json(value)
+        json_utils.validate_json(value)
     except SupersetException as ex:
         raise ValidationError("JSON not valid") from ex
 

superset/embedded/view.py

@@ -14,7 +14,6 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-import json
 from typing import Callable
 
 from flask import abort, request
@@ -25,7 +24,7 @@
 from superset import event_logger, is_feature_enabled
 from superset.daos.dashboard import EmbeddedDashboardDAO
 from superset.superset_typing import FlaskResponse
-from superset.utils import core as utils
+from superset.utils import json as json_utils
 from superset.views.base import BaseSupersetView, common_bootstrap_payload
 
 
@@ -87,7 +86,7 @@ def embedded(
         return self.render_template(
             "superset/spa.html",
             entry="embedded",
-            bootstrap_data=json.dumps(
-                bootstrap_data, default=utils.pessimistic_json_iso_dttm_ser
+            bootstrap_data=json_utils.dumps(
+                bootstrap_data, default=json_utils.pessimistic_json_iso_dttm_ser
             ),
         )

superset/key_value/utils.py

@@ -26,7 +26,7 @@
 
 from superset.key_value.exceptions import KeyValueParseKeyError
 from superset.key_value.types import KeyValueFilter, KeyValueResource
-from superset.utils.core import json_dumps_w_dates
+from superset.utils.json import json_dumps_w_dates
 
 HASHIDS_MIN_LENGTH = 11
 

superset/models/core.py

@@ -76,7 +76,7 @@
 from superset.result_set import SupersetResultSet
 from superset.sql_parse import Table
 from superset.superset_typing import OAuth2ClientConfig, ResultSetColumnType
-from superset.utils import cache as cache_util, core as utils
+from superset.utils import cache as cache_util, core as utils, json as json_utils
 from superset.utils.backports import StrEnum
 from superset.utils.core import DatasourceName, get_username
 from superset.utils.oauth2 import get_oauth2_access_token
@@ -601,7 +601,7 @@ def column_needs_conversion(df_series: pd.Series) -> bool:
 
         for col, coltype in df.dtypes.to_dict().items():
             if coltype == numpy.object_ and column_needs_conversion(df[col]):
-                df[col] = df[col].apply(utils.json_dumps_w_dates)
+                df[col] = df[col].apply(json_utils.json_dumps_w_dates)
         return df
 
     @property
@@ -957,7 +957,7 @@ def get_pk_constraint(self, table: Table) -> dict[str, Any]:
 
             def _convert(value: Any) -> Any:
                 try:
-                    return utils.base_json_conv(value)
+                    return json_utils.base_json_conv(value)
                 except TypeError:
                     return None
 

superset/models/dashboard.py

@@ -51,7 +51,7 @@
 from superset.tasks.thumbnails import cache_dashboard_thumbnail
 from superset.tasks.utils import get_current_user
 from superset.thumbnails.digest import get_dashboard_digest
-from superset.utils import core as utils
+from superset.utils import core as utils, json as json_utils
 
 metadata = Model.metadata  # pylint: disable=no-member
 config = app.config
@@ -372,7 +372,7 @@ def export_dashboards(  # pylint: disable=too-many-locals
 
         return json.dumps(
             {"dashboards": copied_dashboards, "datasources": eager_datasources},
-            cls=utils.DashboardEncoder,
+            cls=json_utils.DashboardEncoder,
             indent=4,
         )
 

superset/result_set.py

@@ -28,7 +28,7 @@
 
 from superset.db_engine_specs import BaseEngineSpec
 from superset.superset_typing import DbapiDescription, DbapiResult, ResultSetColumnType
-from superset.utils import core as utils
+from superset.utils import core as utils, json as json_utils
 from superset.utils.core import GenericDataType
 
 logger = logging.getLogger(__name__)
@@ -61,7 +61,7 @@ def dedup(l: list[str], suffix: str = "__", case_sensitive: bool = True) -> list
 
 
 def stringify(obj: Any) -> str:
-    return json.dumps(obj, default=utils.json_iso_dttm_ser)
+    return json_utils.dumps(obj, default=json_utils.json_iso_dttm_ser)
 
 
 def stringify_values(array: NDArray[Any]) -> NDArray[Any]:

superset/sql_lab.py

@@ -24,7 +24,6 @@
 
 import backoff
 import msgpack
-import simplejson as json
 from celery.exceptions import SoftTimeLimitExceeded
 from flask_babel import gettext as __
 
@@ -59,8 +58,8 @@
 )
 from superset.sqllab.limiting_factor import LimitingFactor
 from superset.sqllab.utils import write_ipc_buffer
+from superset.utils import json as json_utils
 from superset.utils.core import (
-    json_iso_dttm_ser,
     override_user,
     QuerySource,
     zlib_compress,
@@ -349,9 +348,13 @@ def _serialize_payload(
 ) -> Union[bytes, str]:
     logger.debug("Serializing to msgpack: %r", use_msgpack)
     if use_msgpack:
-        return msgpack.dumps(payload, default=json_iso_dttm_ser, use_bin_type=True)
+        return msgpack.dumps(
+            payload, default=json_utils.json_iso_dttm_ser, use_bin_type=True
+        )
 
-    return json.dumps(payload, default=json_iso_dttm_ser, ignore_nan=True)
+    return json_utils.dumps(
+        payload, default=json_utils.json_iso_dttm_ser, ignore_nan=True
+    )
 
 
 def _serialize_and_expand_data(

superset/sqllab/api.py

@@ -18,7 +18,6 @@
 from typing import Any, cast, Optional
 from urllib import parse
 
-import simplejson as json
 from flask import request, Response
 from flask_appbuilder import permission_name
 from flask_appbuilder.api import expose, protect, rison, safe
@@ -62,7 +61,7 @@
 from superset.sqllab.utils import bootstrap_sqllab_data
 from superset.sqllab.validators import CanAccessQueryValidatorImpl
 from superset.superset_typing import FlaskResponse
-from superset.utils import core as utils
+from superset.utils import core as utils, json as json_utils
 from superset.views.base import CsvResponse, generate_download_headers, json_success
 from superset.views.base_api import BaseSupersetApi, requires_json, statsd_metrics
 
@@ -132,9 +131,9 @@ def get(self) -> Response:
         result = bootstrap_sqllab_data(user_id)
 
         return json_success(
-            json.dumps(
+            json_utils.dumps(
                 {"result": result},
-                default=utils.json_iso_dttm_ser,
+                default=json_utils.json_iso_dttm_ser,
                 ignore_nan=True,
             ),
             200,
@@ -343,9 +342,9 @@ def get_results(self, **kwargs: Any) -> FlaskResponse:
 
         # Using pessimistic json serialization since some database drivers can return
         # unserializeable types at times
-        payload = json.dumps(
+        payload = json_utils.dumps(
             result,
-            default=utils.pessimistic_json_iso_dttm_ser,
+            default=json_utils.pessimistic_json_iso_dttm_ser,
             ignore_nan=True,
         )
         return json_success(payload, 200)

superset/sqllab/execution_context_convertor.py

@@ -19,11 +19,9 @@
 import logging
 from typing import Any, TYPE_CHECKING
 
-import simplejson as json
-
-import superset.utils.core as utils
 from superset.sqllab.command_status import SqlJsonExecutionStatus
 from superset.sqllab.utils import apply_display_max_row_configuration_if_require
+from superset.utils import json as json_utils
 
 logger = logging.getLogger(__name__)
 
@@ -52,23 +50,16 @@ def set_payload(
 
     def serialize_payload(self) -> str:
         if self._exc_status == SqlJsonExecutionStatus.HAS_RESULTS:
-            sql_results = apply_display_max_row_configuration_if_require(
-                self.payload, self._max_row_in_display_configuration
+            return json_utils.dumps(
+                apply_display_max_row_configuration_if_require(
+                    self.payload, self._max_row_in_display_configuration
+                ),
+                default=json_utils.pessimistic_json_iso_dttm_ser,
+                ignore_nan=True,
             )
-            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
+        return json_utils.dumps(
+            {"query": self.payload},
+            default=json_utils.json_int_dttm_ser,
+            ignore_nan=True,
         )

superset/utils/cache.py

@@ -30,8 +30,8 @@
 from superset import db
 from superset.extensions import cache_manager
 from superset.models.cache import CacheKey
-from superset.utils.core import json_int_dttm_ser
 from superset.utils.hashing import md5_sha_from_dict
+from superset.utils.json import json_int_dttm_ser
 
 if TYPE_CHECKING:
     from superset.stats_logger import BaseStatsLogger