chore(key-value): use json serialization for main resources (#23888)

superset/dashboards/permalink/commands/base.py

@@ -18,11 +18,12 @@
 
 from superset.commands.base import BaseCommand
 from superset.key_value.shared_entries import get_permalink_salt
-from superset.key_value.types import KeyValueResource, SharedKey
+from superset.key_value.types import JsonKeyValueCodec, KeyValueResource, SharedKey
 
 
 class BaseDashboardPermalinkCommand(BaseCommand, ABC):
     resource = KeyValueResource.DASHBOARD_PERMALINK
+    codec = JsonKeyValueCodec()
 
     @property
     def salt(self) -> str:

superset/dashboards/permalink/commands/create.py

@@ -58,6 +58,7 @@ def run(self) -> str:
                 resource=self.resource,
                 key=get_deterministic_uuid(self.salt, (user_id, value)),
                 value=value,
+                codec=self.codec,
             ).run()
             assert key.id  # for type checks
             return encode_permalink_key(key=key.id, salt=self.salt)

superset/dashboards/permalink/commands/get.py

@@ -39,7 +39,11 @@ def run(self) -> Optional[DashboardPermalinkValue]:
         self.validate()
         try:
             key = decode_permalink_id(self.key, salt=self.salt)
-            command = GetKeyValueCommand(resource=self.resource, key=key)
+            command = GetKeyValueCommand(
+                resource=self.resource,
+                key=key,
+                codec=self.codec,
+            )
             value: Optional[DashboardPermalinkValue] = command.run()
             if value:
                 DashboardDAO.get_by_id_or_slug(value["dashboardId"])

superset/explore/permalink/commands/base.py

@@ -18,11 +18,12 @@
 
 from superset.commands.base import BaseCommand
 from superset.key_value.shared_entries import get_permalink_salt
-from superset.key_value.types import KeyValueResource, SharedKey
+from superset.key_value.types import JsonKeyValueCodec, KeyValueResource, SharedKey
 
 
 class BaseExplorePermalinkCommand(BaseCommand, ABC):
     resource: KeyValueResource = KeyValueResource.EXPLORE_PERMALINK
+    codec = JsonKeyValueCodec()
 
     @property
     def salt(self) -> str:

superset/explore/permalink/commands/create.py

@@ -52,6 +52,7 @@ def run(self) -> str:
             command = CreateKeyValueCommand(
                 resource=self.resource,
                 value=value,
+                codec=self.codec,
             )
             key = command.run()
             if key.id is None:

superset/explore/permalink/commands/get.py

@@ -43,6 +43,7 @@ def run(self) -> Optional[ExplorePermalinkValue]:
             value: Optional[ExplorePermalinkValue] = GetKeyValueCommand(
                 resource=self.resource,
                 key=key,
+                codec=self.codec,
             ).run()
             if value:
                 chart_id: Optional[int] = value.get("chartId")

superset/extensions/metastore_cache.py

@@ -23,10 +23,11 @@
 from flask_caching import BaseCache
 
 from superset.key_value.exceptions import KeyValueCreateFailedError
-from superset.key_value.types import KeyValueResource
+from superset.key_value.types import KeyValueResource, PickleKeyValueCodec
 from superset.key_value.utils import get_uuid_namespace
 
 RESOURCE = KeyValueResource.METASTORE_CACHE
+CODEC = PickleKeyValueCodec()
 
 
 class SupersetMetastoreCache(BaseCache):
@@ -68,6 +69,7 @@ def set(self, key: str, value: Any, timeout: Optional[int] = None) -> bool:
             resource=RESOURCE,
             key=self.get_key(key),
             value=value,
+            codec=CODEC,
             expires_on=self._get_expiry(timeout),
         ).run()
         return True
@@ -80,6 +82,7 @@ def add(self, key: str, value: Any, timeout: Optional[int] = None) -> bool:
             CreateKeyValueCommand(
                 resource=RESOURCE,
                 value=value,
+                codec=CODEC,
                 key=self.get_key(key),
                 expires_on=self._get_expiry(timeout),
             ).run()
@@ -92,7 +95,11 @@ def get(self, key: str) -> Any:
         # pylint: disable=import-outside-toplevel
         from superset.key_value.commands.get import GetKeyValueCommand
 
-        return GetKeyValueCommand(resource=RESOURCE, key=self.get_key(key)).run()
+        return GetKeyValueCommand(
+            resource=RESOURCE,
+            key=self.get_key(key),
+            codec=CODEC,
+        ).run()
 
     def has(self, key: str) -> bool:
         entry = self.get(key)

superset/key_value/commands/create.py

@@ -15,7 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 import logging
-import pickle
 from datetime import datetime
 from typing import Any, Optional, Union
 from uuid import UUID
@@ -26,7 +25,7 @@
 from superset.commands.base import BaseCommand
 from superset.key_value.exceptions import KeyValueCreateFailedError
 from superset.key_value.models import KeyValueEntry
-from superset.key_value.types import Key, KeyValueResource
+from superset.key_value.types import Key, KeyValueCodec, KeyValueResource
 from superset.utils.core import get_user_id
 
 logger = logging.getLogger(__name__)
@@ -35,13 +34,15 @@
 class CreateKeyValueCommand(BaseCommand):
     resource: KeyValueResource
     value: Any
+    codec: KeyValueCodec
     key: Optional[Union[int, UUID]]
     expires_on: Optional[datetime]
 
-    def __init__(
+    def __init__(  # pylint: disable=too-many-arguments
         self,
         resource: KeyValueResource,
         value: Any,
+        codec: KeyValueCodec,
         key: Optional[Union[int, UUID]] = None,
         expires_on: Optional[datetime] = None,
     ):
@@ -50,16 +51,24 @@ def __init__(
 
         :param resource: the resource (dashboard, chart etc)
         :param value: the value to persist in the key-value store
+        :param codec: codec used to encode the value
         :param key: id of entry (autogenerated if undefined)
         :param expires_on: entry expiration time
-        :return: the key associated with the persisted value
+        :
         """
         self.resource = resource
         self.value = value
+        self.codec = codec
         self.key = key
         self.expires_on = expires_on
 
     def run(self) -> Key:
+        """
+        Persist the value
+
+        :return: the key associated with the persisted value
+
+        """
         try:
             return self.create()
         except SQLAlchemyError as ex:
@@ -70,9 +79,13 @@ def validate(self) -> None:
         pass
 
     def create(self) -> Key:
+        try:
+            value = self.codec.encode(self.value)
+        except Exception as ex:  # pylint: disable=broad-except
+            raise KeyValueCreateFailedError("Unable to encode value") from ex
         entry = KeyValueEntry(
             resource=self.resource.value,
-            value=pickle.dumps(self.value),
+            value=value,
             created_on=datetime.now(),
             created_by_fk=get_user_id(),
             expires_on=self.expires_on,

superset/key_value/commands/get.py

@@ -16,7 +16,6 @@
 # under the License.
 
 import logging
-import pickle
 from datetime import datetime
 from typing import Any, Optional, Union
 from uuid import UUID
@@ -27,7 +26,7 @@
 from superset.commands.base import BaseCommand
 from superset.key_value.exceptions import KeyValueGetFailedError
 from superset.key_value.models import KeyValueEntry
-from superset.key_value.types import KeyValueResource
+from superset.key_value.types import KeyValueCodec, KeyValueResource
 from superset.key_value.utils import get_filter
 
 logger = logging.getLogger(__name__)
@@ -36,17 +35,25 @@
 class GetKeyValueCommand(BaseCommand):
     resource: KeyValueResource
     key: Union[int, UUID]
+    codec: KeyValueCodec
 
-    def __init__(self, resource: KeyValueResource, key: Union[int, UUID]):
+    def __init__(
+        self,
+        resource: KeyValueResource,
+        key: Union[int, UUID],
+        codec: KeyValueCodec,
+    ):
         """
         Retrieve a key value entry
 
         :param resource: the resource (dashboard, chart etc)
         :param key: the key to retrieve
+        :param codec: codec used to decode the value
         :return: the value associated with the key if present
         """
         self.resource = resource
         self.key = key
+        self.codec = codec
 
     def run(self) -> Any:
         try:
@@ -66,5 +73,5 @@ def get(self) -> Optional[Any]:
             .first()
         )
         if entry and (entry.expires_on is None or entry.expires_on > datetime.now()):
-            return pickle.loads(entry.value)
+            return self.codec.decode(entry.value)
         return None

superset/key_value/commands/update.py

@@ -16,7 +16,6 @@
 # under the License.
 
 import logging
-import pickle
 from datetime import datetime
 from typing import Any, Optional, Union
 from uuid import UUID
@@ -27,7 +26,7 @@
 from superset.commands.base import BaseCommand
 from superset.key_value.exceptions import KeyValueUpdateFailedError
 from superset.key_value.models import KeyValueEntry
-from superset.key_value.types import Key, KeyValueResource
+from superset.key_value.types import Key, KeyValueCodec, KeyValueResource
 from superset.key_value.utils import get_filter
 from superset.utils.core import get_user_id
 
@@ -37,14 +36,16 @@
 class UpdateKeyValueCommand(BaseCommand):
     resource: KeyValueResource
     value: Any
+    codec: KeyValueCodec
     key: Union[int, UUID]
     expires_on: Optional[datetime]
 
-    def __init__(
+    def __init__(  # pylint: disable=too-many-arguments
         self,
         resource: KeyValueResource,
         key: Union[int, UUID],
         value: Any,
+        codec: KeyValueCodec,
         expires_on: Optional[datetime] = None,
     ):
         """
@@ -53,12 +54,14 @@ def __init__(
         :param resource: the resource (dashboard, chart etc)
         :param key: the key to update
         :param value: the value to persist in the key-value store
+        :param codec: codec used to encode the value
         :param expires_on: entry expiration time
         :return: the key associated with the updated value
         """
         self.resource = resource
         self.key = key
         self.value = value
+        self.codec = codec
         self.expires_on = expires_on
 
     def run(self) -> Optional[Key]:
@@ -80,7 +83,7 @@ def update(self) -> Optional[Key]:
             .first()
         )
         if entry:
-            entry.value = pickle.dumps(self.value)
+            entry.value = self.codec.encode(self.value)
             entry.expires_on = self.expires_on
             entry.changed_on = datetime.now()
             entry.changed_by_fk = get_user_id()

superset/key_value/commands/upsert.py

@@ -16,7 +16,6 @@
 # under the License.
 
 import logging
-import pickle
 from datetime import datetime
 from typing import Any, Optional, Union
 from uuid import UUID
@@ -31,7 +30,7 @@
     KeyValueUpsertFailedError,
 )
 from superset.key_value.models import KeyValueEntry
-from superset.key_value.types import Key, KeyValueResource
+from superset.key_value.types import Key, KeyValueCodec, KeyValueResource
 from superset.key_value.utils import get_filter
 from superset.utils.core import get_user_id
 
@@ -42,13 +41,15 @@ class UpsertKeyValueCommand(BaseCommand):
     resource: KeyValueResource
     value: Any
     key: Union[int, UUID]
+    codec: KeyValueCodec
     expires_on: Optional[datetime]
 
-    def __init__(
+    def __init__(  # pylint: disable=too-many-arguments
         self,
         resource: KeyValueResource,
         key: Union[int, UUID],
         value: Any,
+        codec: KeyValueCodec,
         expires_on: Optional[datetime] = None,
     ):
         """
@@ -57,13 +58,14 @@ def __init__(
         :param resource: the resource (dashboard, chart etc)
         :param key: the key to update
         :param value: the value to persist in the key-value store
-        :param key_type: the type of the key to update
+        :param codec: codec used to encode the value
         :param expires_on: entry expiration time
         :return: the key associated with the updated value
         """
         self.resource = resource
         self.key = key
         self.value = value
+        self.codec = codec
         self.expires_on = expires_on
 
     def run(self) -> Key:
@@ -85,7 +87,7 @@ def upsert(self) -> Key:
             .first()
         )
         if entry:
-            entry.value = pickle.dumps(self.value)
+            entry.value = self.codec.encode(self.value)
             entry.expires_on = self.expires_on
             entry.changed_on = datetime.now()
             entry.changed_by_fk = get_user_id()
@@ -96,6 +98,7 @@ def upsert(self) -> Key:
         return CreateKeyValueCommand(
             resource=self.resource,
             value=self.value,
+            codec=self.codec,
             key=self.key,
             expires_on=self.expires_on,
         ).run()

superset/key_value/shared_entries.py

@@ -18,27 +18,33 @@
 from typing import Any, Optional
 from uuid import uuid3
 
-from superset.key_value.types import KeyValueResource, SharedKey
+from superset.key_value.types import JsonKeyValueCodec, KeyValueResource, SharedKey
 from superset.key_value.utils import get_uuid_namespace, random_key
 
 RESOURCE = KeyValueResource.APP
 NAMESPACE = get_uuid_namespace("")
+CODEC = JsonKeyValueCodec()
 
 
 def get_shared_value(key: SharedKey) -> Optional[Any]:
     # pylint: disable=import-outside-toplevel
     from superset.key_value.commands.get import GetKeyValueCommand
 
     uuid_key = uuid3(NAMESPACE, key)
-    return GetKeyValueCommand(RESOURCE, key=uuid_key).run()
+    return GetKeyValueCommand(RESOURCE, key=uuid_key, codec=CODEC).run()
 
 
 def set_shared_value(key: SharedKey, value: Any) -> None:
     # pylint: disable=import-outside-toplevel
     from superset.key_value.commands.create import CreateKeyValueCommand
 
     uuid_key = uuid3(NAMESPACE, key)
-    CreateKeyValueCommand(resource=RESOURCE, value=value, key=uuid_key).run()
+    CreateKeyValueCommand(
+        resource=RESOURCE,
+        value=value,
+        key=uuid_key,
+        codec=CODEC,
+    ).run()
 
 
 def get_permalink_salt(key: SharedKey) -> str: