fix(marshmallow): add compatibility layer for Flask-AppBuilder with marshmallow 4.x

superset/marshmallow_compatibility.py

@@ -52,7 +52,9 @@ def patch_marshmallow_for_flask_appbuilder() -> None:
     """
     import marshmallow
 
-    if getattr(marshmallow.Schema._init_fields, "_fab_patched", False):
+    if getattr(marshmallow.Schema._init_fields, "_fab_patched", False) or getattr(
+        marshmallow.Schema._init_fields, "__superset_compat_patched__", False
+    ):
         return  # already patched
 
     original_init_fields = marshmallow.Schema._init_fields
@@ -65,17 +67,23 @@ def patched_init_fields(self: "marshmallow.Schema") -> Any:
         # otherwise choke on; unknown non-FAB names are left to raise so genuine
         # schema bugs still surface.
         candidates = set(opts.fields or ()) | set(getattr(opts, "additional", ()) or ())
+        declared_fields = dict(getattr(self, "declared_fields", {}))
         for name in candidates:
-            if name not in self.declared_fields and _looks_like_fab_field(name):
+            if name not in declared_fields and _looks_like_fab_field(name):
                 logger.debug(
                     "marshmallow FAB compat: stubbing Raw field for %r on %s",
                     name,
                     type(self).__name__,
                 )
-                self.declared_fields[name] = marshmallow.fields.Raw()
+                declared_fields[name] = marshmallow.fields.Raw(
+                    allow_none=True,
+                    load_default=None,
+                )
+        self.declared_fields = declared_fields
         return original_init_fields(self)
 
     patched_init_fields._fab_patched = True  # type: ignore[attr-defined]
+    patched_init_fields.__superset_compat_patched__ = True  # type: ignore[attr-defined]
     marshmallow.Schema._init_fields = patched_init_fields
 
 

superset/models/helpers.py

@@ -976,6 +976,7 @@ def ensure_extra_json_is_not_none(
         self,
         _: str,
         value: Optional[dict[str, Any]],
+        **kwargs: Any,
     ) -> Any:
         if value is None:
             return "{}"

superset/themes/schemas.py

@@ -14,9 +14,10 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+from contextvars import ContextVar
 from typing import Any
 
-from marshmallow import fields, Schema, validates, ValidationError
+from marshmallow import fields, post_load, Schema, validates, ValidationError
 
 from superset.themes.utils import (
     is_valid_theme,
@@ -25,6 +26,17 @@
 )
 from superset.utils import json
 
+# Context variable for storing sanitized JSON data during validation
+sanitized_json_context: ContextVar[str | None] = ContextVar(
+    "sanitized_json_data", default=None
+)
+
+
+def _consume_sanitized_json_data(default: str) -> str:
+    sanitized_json_data = sanitized_json_context.get()
+    sanitized_json_context.set(None)
+    return sanitized_json_data or default
+
 
 def _sanitize_and_validate_theme_config(theme_config: dict[str, Any]) -> dict[str, Any]:
     """Sanitize and validate theme configuration.
@@ -77,7 +89,7 @@ def validate_json_data(self, value: dict[str, Any], **kwargs: Any) -> None:
                 value.clear()
                 value.update(sanitized_config)
             else:
-                self.context["sanitized_json_data"] = json.dumps(sanitized_config)
+                sanitized_json_context.set(json.dumps(sanitized_config))
 
 
 class ThemePostSchema(Schema):
@@ -91,6 +103,7 @@ def validate_theme_name(self, value: str, **kwargs: Any) -> None:
 
     @validates("json_data")
     def validate_and_sanitize_json_data(self, value: str, **kwargs: Any) -> None:
+        sanitized_json_context.set(None)
         # Parse JSON
         try:
             theme_config = json.loads(value) if isinstance(value, str) else value
@@ -104,7 +117,18 @@ def validate_and_sanitize_json_data(self, value: str, **kwargs: Any) -> None:
         # Note: This modifies the input data to ensure sanitized content is stored
         if sanitized_config != theme_config:
             # Re-serialize the sanitized config
-            self.context["sanitized_json_data"] = json.dumps(sanitized_config)
+            sanitized_json_context.set(json.dumps(sanitized_config))
+
+    def _apply_sanitized_json_data(self, data: dict[str, Any]) -> dict[str, Any]:
+        data["json_data"] = _consume_sanitized_json_data(data["json_data"])
+        return data
+
+    @post_load
+    def apply_sanitized_json_data(
+        self, data: dict[str, Any], **kwargs: Any
+    ) -> dict[str, Any]:
+        del kwargs
+        return self._apply_sanitized_json_data(data)
 
 
 class ThemePutSchema(Schema):
@@ -118,6 +142,7 @@ def validate_theme_name(self, value: str, **kwargs: Any) -> None:
 
     @validates("json_data")
     def validate_and_sanitize_json_data(self, value: str, **kwargs: Any) -> None:
+        sanitized_json_context.set(None)
         # Parse JSON
         try:
             theme_config = json.loads(value) if isinstance(value, str) else value
@@ -131,7 +156,15 @@ def validate_and_sanitize_json_data(self, value: str, **kwargs: Any) -> None:
         # Note: This modifies the input data to ensure sanitized content is stored
         if sanitized_config != theme_config:
             # Re-serialize the sanitized config
-            self.context["sanitized_json_data"] = json.dumps(sanitized_config)
+            sanitized_json_context.set(json.dumps(sanitized_config))
+
+    @post_load
+    def apply_sanitized_json_data(
+        self, data: dict[str, Any], **kwargs: Any
+    ) -> dict[str, Any]:
+        del kwargs
+        data["json_data"] = _consume_sanitized_json_data(data["json_data"])
+        return data
 
 
 openapi_spec_methods_override = {

tests/unit_tests/test_marshmallow_compatibility.py

@@ -0,0 +1,182 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Unit tests for the marshmallow compatibility shim."""
+
+from unittest.mock import patch
+
+import pytest
+from marshmallow import fields, Schema
+
+from superset.marshmallow_compatibility import patch_marshmallow_for_flask_appbuilder
+
+
+@pytest.fixture(autouse=True)
+def restore_schema_init_fields():
+    """Restore Schema._init_fields after each test."""
+    original_method = Schema._init_fields
+    try:
+        yield
+    finally:
+        Schema._init_fields = original_method
+
+
+class TestMarshmallowCompatibility:
+    """Test cases for the marshmallow 4.x compatibility module."""
+
+    def test_patch_marshmallow_for_flask_appbuilder_applies_patch(self):
+        """Test that the patch function correctly replaces Schema._init_fields."""
+        original_method = Schema._init_fields
+
+        patch_marshmallow_for_flask_appbuilder()
+
+        assert callable(Schema._init_fields)
+        assert getattr(Schema._init_fields, "__superset_compat_patched__", False)
+        if not getattr(original_method, "__superset_compat_patched__", False):
+            assert Schema._init_fields != original_method
+
+    def test_patch_functionality_with_real_schema_creation(self):
+        """Test that the patch works with actual schema creation scenarios."""
+        patch_marshmallow_for_flask_appbuilder()
+
+        class TestSchema(Schema):
+            name = fields.Str()
+            age = fields.Int()
+
+        schema = TestSchema()
+        assert "name" in schema.declared_fields
+        assert "age" in schema.declared_fields
+        assert isinstance(schema.declared_fields["name"], fields.Str)
+        assert isinstance(schema.declared_fields["age"], fields.Int)
+
+    def test_patch_handles_schema_with_no_fields(self):
+        """Test that the patch works with schemas that have no declared fields."""
+        patch_marshmallow_for_flask_appbuilder()
+
+        class EmptySchema(Schema):
+            pass
+
+        schema = EmptySchema()
+        assert hasattr(schema, "declared_fields")
+
+    def test_patch_pre_stubs_missing_fab_fields_before_initialization(self):
+        """Test that FAB-looking missing fields are stubbed before init runs."""
+
+        def fake_init_fields(self: Schema):
+            assert "missing_field" in self.declared_fields
+            return "initialized"
+
+        with patch.object(Schema, "_init_fields", fake_init_fields):
+            patch_marshmallow_for_flask_appbuilder()
+
+            schema = object.__new__(Schema)
+            schema.declared_fields = {}
+            schema.opts = type(
+                "Opts",
+                (),
+                {"fields": ("missing_field",), "additional": ()},
+            )()
+
+            result = Schema._init_fields(schema)
+
+        assert result == "initialized"
+        assert isinstance(schema.declared_fields["missing_field"], fields.Raw)
+        assert schema.declared_fields["missing_field"].allow_none is True
+        assert schema.declared_fields["missing_field"].load_default is None
+
+    def test_raw_field_creation_and_configuration(self):
+        """Test that Raw fields can be created with the expected configuration."""
+        # Test creating a Raw field with our configuration
+        raw_field = fields.Raw(allow_none=True, load_default=None)
+
+        assert isinstance(raw_field, fields.Raw)
+        assert raw_field.allow_none is True
+        assert raw_field.load_default is None
+
+    def test_schema_declared_fields_manipulation(self):
+        """Test that we can manipulate schema declared_fields."""
+
+        class TestSchema(Schema):
+            existing_field = fields.Str()
+
+        schema = TestSchema()
+
+        # Verify initial state
+        assert "existing_field" in schema.declared_fields
+        assert isinstance(schema.declared_fields["existing_field"], fields.Str)
+
+        # Test adding a new field
+        schema.declared_fields["new_field"] = fields.Raw(
+            allow_none=True, load_default=None
+        )
+
+        # Verify the new field was added
+        assert "new_field" in schema.declared_fields
+        assert isinstance(schema.declared_fields["new_field"], fields.Raw)
+        assert schema.declared_fields["new_field"].allow_none is True
+        assert schema.declared_fields["new_field"].load_default is None
+
+    def test_flask_appbuilder_field_names_list(self):
+        """Test that we have the correct list of Flask-AppBuilder field names."""
+        # Common Flask-AppBuilder auto-generated field names that our fix handles
+        expected_fab_fields = [
+            "permission_id",
+            "view_menu_id",
+            "db_id",
+            "chart_id",
+            "dashboard_id",
+            "user_id",
+        ]
+
+        # Verify these are strings (field names)
+        for field_name in expected_fab_fields:
+            assert isinstance(field_name, str)
+            assert len(field_name) > 0
+            assert "_id" in field_name
+
+    def test_patch_function_is_callable(self):
+        """Test that the patch function can be called without errors."""
+        patch_marshmallow_for_flask_appbuilder()
+        patched_method = Schema._init_fields
+
+        patch_marshmallow_for_flask_appbuilder()
+        patch_marshmallow_for_flask_appbuilder()
+
+        assert Schema._init_fields is patched_method
+
+    def test_marshmallow_schema_basic_functionality(self):
+        """Test basic marshmallow schema functionality still works."""
+
+        class UserSchema(Schema):
+            name = fields.Str(required=True)
+            email = fields.Email()
+            age = fields.Int(validate=lambda x: x > 0)
+
+        schema = UserSchema()
+
+        # Test serialization
+        data = {"name": "John Doe", "email": "john@example.com", "age": 30}
+        result = schema.load(data)
+        assert result["name"] == "John Doe"
+        assert result["email"] == "john@example.com"
+        assert result["age"] == 30
+
+        from marshmallow import ValidationError
+
+        # Test validation - missing required field should raise error
+        with pytest.raises(ValidationError):
+            schema.load({"email": "john@example.com", "age": 30})  # Missing name

tests/unit_tests/themes/api_test.py

@@ -15,8 +15,9 @@
 # specific language governing permissions and limitations
 # under the License.
 
-
 from superset.themes.api import ThemeRestApi
+from superset.themes.schemas import ThemePostSchema, ThemePutSchema
+from superset.utils import json
 
 
 class TestThemeRestApi:
@@ -104,8 +105,6 @@ def test_bulk_delete_enabled(self):
 
     def test_custom_schemas_configured(self):
         """Test that custom schemas are properly configured"""
-        from superset.themes.schemas import ThemePostSchema, ThemePutSchema
-
         api = ThemeRestApi()
         assert isinstance(api.add_model_schema, ThemePostSchema)
         assert isinstance(api.edit_model_schema, ThemePutSchema)
@@ -121,3 +120,43 @@ def test_list_columns_include_new_fields(self):
         expected_new_fields = ["is_system", "uuid"]
         for field in expected_new_fields:
             assert field in ThemeRestApi.list_columns
+
+    def test_theme_post_schema_uses_sanitized_json_data(self):
+        """Test that sanitized theme JSON is written back during load."""
+        schema = ThemePostSchema()
+
+        result = schema.load(
+            {
+                "theme_name": "test_theme",
+                "json_data": json.dumps(
+                    {
+                        "token": {
+                            "brandSpinnerUrl": "javascript:alert('xss')",
+                            "colorPrimary": "#ff0000",
+                        }
+                    }
+                ),
+            }
+        )
+
+        assert json.loads(result["json_data"])["token"]["brandSpinnerUrl"] == ""
+
+    def test_theme_put_schema_uses_sanitized_json_data(self):
+        """Test that update schema also writes back sanitized theme JSON."""
+        schema = ThemePutSchema()
+
+        result = schema.load(
+            {
+                "theme_name": "test_theme",
+                "json_data": json.dumps(
+                    {
+                        "token": {
+                            "brandSpinnerUrl": "javascript:alert('xss')",
+                            "colorPrimary": "#ff0000",
+                        }
+                    }
+                ),
+            }
+        )
+
+        assert json.loads(result["json_data"])["token"]["brandSpinnerUrl"] == ""