fix(permalink): migrate to marshmallow codec (#24166)

superset/dashboards/permalink/api.py

@@ -30,7 +30,7 @@
 )
 from superset.dashboards.permalink.commands.get import GetDashboardPermalinkCommand
 from superset.dashboards.permalink.exceptions import DashboardPermalinkInvalidStateError
-from superset.dashboards.permalink.schemas import DashboardPermalinkPostSchema
+from superset.dashboards.permalink.schemas import DashboardPermalinkStateSchema
 from superset.extensions import event_logger
 from superset.key_value.exceptions import KeyValueAccessDeniedError
 from superset.views.base_api import BaseSupersetApi, requires_json
@@ -39,13 +39,13 @@
 
 
 class DashboardPermalinkRestApi(BaseSupersetApi):
-    add_model_schema = DashboardPermalinkPostSchema()
+    add_model_schema = DashboardPermalinkStateSchema()
     method_permission_name = MODEL_API_RW_METHOD_PERMISSION_MAP
     allow_browser_login = True
     class_permission_name = "DashboardPermalinkRestApi"
     resource_name = "dashboard"
     openapi_spec_tag = "Dashboard Permanent Link"
-    openapi_spec_component_schemas = (DashboardPermalinkPostSchema,)
+    openapi_spec_component_schemas = (DashboardPermalinkStateSchema,)
 
     @expose("/<pk>/permalink", methods=["POST"])
     @protect()

superset/dashboards/permalink/commands/base.py

@@ -17,13 +17,18 @@
 from abc import ABC
 
 from superset.commands.base import BaseCommand
+from superset.dashboards.permalink.schemas import DashboardPermalinkSchema
 from superset.key_value.shared_entries import get_permalink_salt
-from superset.key_value.types import JsonKeyValueCodec, KeyValueResource, SharedKey
+from superset.key_value.types import (
+    KeyValueResource,
+    MarshmallowKeyValueCodec,
+    SharedKey,
+)
 
 
 class BaseDashboardPermalinkCommand(BaseCommand, ABC):
     resource = KeyValueResource.DASHBOARD_PERMALINK
-    codec = JsonKeyValueCodec()
+    codec = MarshmallowKeyValueCodec(DashboardPermalinkSchema())
 
     @property
     def salt(self) -> str:

superset/dashboards/permalink/commands/create.py

@@ -23,6 +23,7 @@
 from superset.dashboards.permalink.exceptions import DashboardPermalinkCreateFailedError
 from superset.dashboards.permalink.types import DashboardPermalinkState
 from superset.key_value.commands.upsert import UpsertKeyValueCommand
+from superset.key_value.exceptions import KeyValueCodecEncodeException
 from superset.key_value.utils import encode_permalink_key, get_deterministic_uuid
 from superset.utils.core import get_user_id
 
@@ -62,6 +63,8 @@ def run(self) -> str:
             ).run()
             assert key.id  # for type checks
             return encode_permalink_key(key=key.id, salt=self.salt)
+        except KeyValueCodecEncodeException as ex:
+            raise DashboardPermalinkCreateFailedError(str(ex)) from ex
         except SQLAlchemyError as ex:
             logger.exception("Error running create command")
             raise DashboardPermalinkCreateFailedError() from ex

superset/dashboards/permalink/commands/get.py

@@ -25,7 +25,11 @@
 from superset.dashboards.permalink.exceptions import DashboardPermalinkGetFailedError
 from superset.dashboards.permalink.types import DashboardPermalinkValue
 from superset.key_value.commands.get import GetKeyValueCommand
-from superset.key_value.exceptions import KeyValueGetFailedError, KeyValueParseKeyError
+from superset.key_value.exceptions import (
+    KeyValueCodecDecodeException,
+    KeyValueGetFailedError,
+    KeyValueParseKeyError,
+)
 from superset.key_value.utils import decode_permalink_id
 
 logger = logging.getLogger(__name__)
@@ -51,6 +55,7 @@ def run(self) -> Optional[DashboardPermalinkValue]:
             return None
         except (
             DashboardNotFoundError,
+            KeyValueCodecDecodeException,
             KeyValueGetFailedError,
             KeyValueParseKeyError,
         ) as ex:

superset/dashboards/permalink/schemas.py

@@ -17,7 +17,7 @@
 from marshmallow import fields, Schema
 
 
-class DashboardPermalinkPostSchema(Schema):
+class DashboardPermalinkStateSchema(Schema):
     dataMask = fields.Dict(
         required=False,
         allow_none=True,
@@ -48,3 +48,12 @@ class DashboardPermalinkPostSchema(Schema):
         allow_none=True,
         description="Optional anchor link added to url hash",
     )
+
+
+class DashboardPermalinkSchema(Schema):
+    dashboardId = fields.String(
+        required=True,
+        allow_none=False,
+        metadata={"description": "The id or slug of the dasbhoard"},
+    )
+    state = fields.Nested(DashboardPermalinkStateSchema())

superset/explore/permalink/api.py

@@ -32,7 +32,7 @@
 from superset.explore.permalink.commands.create import CreateExplorePermalinkCommand
 from superset.explore.permalink.commands.get import GetExplorePermalinkCommand
 from superset.explore.permalink.exceptions import ExplorePermalinkInvalidStateError
-from superset.explore.permalink.schemas import ExplorePermalinkPostSchema
+from superset.explore.permalink.schemas import ExplorePermalinkStateSchema
 from superset.extensions import event_logger
 from superset.key_value.exceptions import KeyValueAccessDeniedError
 from superset.views.base_api import BaseSupersetApi, requires_json, statsd_metrics
@@ -41,13 +41,13 @@
 
 
 class ExplorePermalinkRestApi(BaseSupersetApi):
-    add_model_schema = ExplorePermalinkPostSchema()
+    add_model_schema = ExplorePermalinkStateSchema()
     method_permission_name = MODEL_API_RW_METHOD_PERMISSION_MAP
     allow_browser_login = True
     class_permission_name = "ExplorePermalinkRestApi"
     resource_name = "explore"
     openapi_spec_tag = "Explore Permanent Link"
-    openapi_spec_component_schemas = (ExplorePermalinkPostSchema,)
+    openapi_spec_component_schemas = (ExplorePermalinkStateSchema,)
 
     @expose("/permalink", methods=["POST"])
     @protect()

superset/explore/permalink/commands/base.py

@@ -17,13 +17,18 @@
 from abc import ABC
 
 from superset.commands.base import BaseCommand
+from superset.explore.permalink.schemas import ExplorePermalinkSchema
 from superset.key_value.shared_entries import get_permalink_salt
-from superset.key_value.types import JsonKeyValueCodec, KeyValueResource, SharedKey
+from superset.key_value.types import (
+    KeyValueResource,
+    MarshmallowKeyValueCodec,
+    SharedKey,
+)
 
 
 class BaseExplorePermalinkCommand(BaseCommand, ABC):
     resource: KeyValueResource = KeyValueResource.EXPLORE_PERMALINK
-    codec = JsonKeyValueCodec()
+    codec = MarshmallowKeyValueCodec(ExplorePermalinkSchema())
 
     @property
     def salt(self) -> str:

superset/explore/permalink/commands/create.py

@@ -23,6 +23,7 @@
 from superset.explore.permalink.exceptions import ExplorePermalinkCreateFailedError
 from superset.explore.utils import check_access as check_chart_access
 from superset.key_value.commands.create import CreateKeyValueCommand
+from superset.key_value.exceptions import KeyValueCodecEncodeException
 from superset.key_value.utils import encode_permalink_key
 from superset.utils.core import DatasourceType
 
@@ -58,6 +59,8 @@ def run(self) -> str:
             if key.id is None:
                 raise ExplorePermalinkCreateFailedError("Unexpected missing key id")
             return encode_permalink_key(key=key.id, salt=self.salt)
+        except KeyValueCodecEncodeException as ex:
+            raise ExplorePermalinkCreateFailedError(str(ex)) from ex
         except SQLAlchemyError as ex:
             logger.exception("Error running create command")
             raise ExplorePermalinkCreateFailedError() from ex

superset/explore/permalink/commands/get.py

@@ -25,7 +25,11 @@
 from superset.explore.permalink.types import ExplorePermalinkValue
 from superset.explore.utils import check_access as check_chart_access
 from superset.key_value.commands.get import GetKeyValueCommand
-from superset.key_value.exceptions import KeyValueGetFailedError, KeyValueParseKeyError
+from superset.key_value.exceptions import (
+    KeyValueCodecDecodeException,
+    KeyValueGetFailedError,
+    KeyValueParseKeyError,
+)
 from superset.key_value.utils import decode_permalink_id
 from superset.utils.core import DatasourceType
 
@@ -59,6 +63,7 @@ def run(self) -> Optional[ExplorePermalinkValue]:
             return None
         except (
             DatasetNotFoundError,
+            KeyValueCodecDecodeException,
             KeyValueGetFailedError,
             KeyValueParseKeyError,
         ) as ex:

superset/explore/permalink/schemas.py

@@ -17,7 +17,7 @@
 from marshmallow import fields, Schema
 
 
-class ExplorePermalinkPostSchema(Schema):
+class ExplorePermalinkStateSchema(Schema):
     formData = fields.Dict(
         required=True,
         allow_none=False,
@@ -37,3 +37,27 @@ class ExplorePermalinkPostSchema(Schema):
         allow_none=True,
         description="URL Parameters",
     )
+
+
+class ExplorePermalinkSchema(Schema):
+    chartId = fields.Integer(
+        required=False,
+        allow_none=True,
+        metadata={"description": "The id of the chart"},
+    )
+    datasourceType = fields.String(
+        required=True,
+        allow_none=False,
+        metadata={"description": "The type of the datasource"},
+    )
+    datasourceId = fields.Integer(
+        required=False,
+        allow_none=True,
+        metadata={"description": "The id of the datasource"},
+    )
+    datasource = fields.String(
+        required=False,
+        allow_none=True,
+        metadata={"description": "The fully qualified datasource reference"},
+    )
+    state = fields.Nested(ExplorePermalinkStateSchema())

superset/key_value/exceptions.py

@@ -52,3 +52,15 @@ class KeyValueUpsertFailedError(UpdateFailedError):
 
 class KeyValueAccessDeniedError(ForbiddenError):
     message = _("You don't have permission to modify the value.")
+
+
+class KeyValueCodecException(SupersetException):
+    pass
+
+
+class KeyValueCodecEncodeException(KeyValueCodecException):
+    message = _("Unable to encode value")
+
+
+class KeyValueCodecDecodeException(KeyValueCodecException):
+    message = _("Unable to decode value")

superset/key_value/types.py

@@ -24,6 +24,13 @@
 from typing import Any, Optional, TypedDict
 from uuid import UUID
 
+from marshmallow import Schema, ValidationError
+
+from superset.key_value.exceptions import (
+    KeyValueCodecDecodeException,
+    KeyValueCodecEncodeException,
+)
+
 
 @dataclass
 class Key:
@@ -61,10 +68,16 @@ def decode(self, value: bytes) -> Any:
 
 class JsonKeyValueCodec(KeyValueCodec):
     def encode(self, value: dict[Any, Any]) -> bytes:
-        return bytes(json.dumps(value), encoding="utf-8")
+        try:
+            return bytes(json.dumps(value), encoding="utf-8")
+        except TypeError as ex:
+            raise KeyValueCodecEncodeException(str(ex)) from ex
 
     def decode(self, value: bytes) -> dict[Any, Any]:
-        return json.loads(value)
+        try:
+            return json.loads(value)
+        except TypeError as ex:
+            raise KeyValueCodecDecodeException(str(ex)) from ex
 
 
 class PickleKeyValueCodec(KeyValueCodec):
@@ -73,3 +86,22 @@ def encode(self, value: dict[Any, Any]) -> bytes:
 
     def decode(self, value: bytes) -> dict[Any, Any]:
         return pickle.loads(value)
+
+
+class MarshmallowKeyValueCodec(JsonKeyValueCodec):
+    def __init__(self, schema: Schema):
+        self.schema = schema
+
+    def encode(self, value: dict[Any, Any]) -> bytes:
+        try:
+            obj = self.schema.dump(value)
+            return super().encode(obj)
+        except ValidationError as ex:
+            raise KeyValueCodecEncodeException(message=str(ex)) from ex
+
+    def decode(self, value: bytes) -> dict[Any, Any]:
+        try:
+            obj = super().decode(value)
+            return self.schema.load(obj)
+        except ValidationError as ex:
+            raise KeyValueCodecEncodeException(message=str(ex)) from ex

tests/integration_tests/explore/permalink/api_tests.py

@@ -22,8 +22,9 @@
 from sqlalchemy.orm import Session
 
 from superset import db
+from superset.explore.permalink.schemas import ExplorePermalinkSchema
 from superset.key_value.models import KeyValueEntry
-from superset.key_value.types import JsonKeyValueCodec, KeyValueResource
+from superset.key_value.types import KeyValueResource, MarshmallowKeyValueCodec
 from superset.key_value.utils import decode_permalink_id, encode_permalink_key
 from superset.models.slice import Slice
 from superset.utils.core import DatasourceType
@@ -94,14 +95,17 @@ def test_get_missing_chart(
     chart_id = 1234
     entry = KeyValueEntry(
         resource=KeyValueResource.EXPLORE_PERMALINK,
-        value=JsonKeyValueCodec().encode(
+        value=MarshmallowKeyValueCodec(ExplorePermalinkSchema()).encode(
             {
                 "chartId": chart_id,
                 "datasourceId": chart.datasource.id,
-                "datasourceType": DatasourceType.TABLE,
-                "formData": {
-                    "slice_id": chart_id,
-                    "datasource": f"{chart.datasource.id}__{chart.datasource.type}",
+                "datasourceType": DatasourceType.TABLE.value,
+                "state": {
+                    "urlParams": [["foo", "bar"]],
+                    "formData": {
+                        "slice_id": chart_id,
+                        "datasource": f"{chart.datasource.id}__{chart.datasource.type}",
+                    },
                 },
             }
         ),