[superset] 03/03: chore(key-value): use json serialization for main resources (#23888)

[el...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

elizabeth pushed a commit to branch 2.1
in repository https://gitbox.apache.org/repos/asf/superset.git

commit 893097b68cb9b6dfd10399aeb6234d3f4e266a66
Author: Ville Brofeldt <33...@users.noreply.github.com>
AuthorDate: Thu May 4 08:04:05 2023 +0300

    chore(key-value): use json serialization for main resources (#23888)
---
 superset/extensions/metastore_cache.py              |  1 +
 ...c2a5681ddfd_convert_key_value_entries_to_json.py | 21 +++++++++++++++++++++
 .../explore/permalink/api_tests.py                  |  8 ++++++++
 3 files changed, 30 insertions(+)

diff --git a/superset/extensions/metastore_cache.py b/superset/extensions/metastore_cache.py
index 0dbd42b2b4..4e6f0067dd 100644
--- a/superset/extensions/metastore_cache.py
+++ b/superset/extensions/metastore_cache.py
@@ -31,6 +31,7 @@ from superset.key_value.types import (
 from superset.key_value.utils import get_uuid_namespace
 
 RESOURCE = KeyValueResource.METASTORE_CACHE
+CODEC = PickleKeyValueCodec()
 
 logger = logging.getLogger(__name__)
 
diff --git a/superset/migrations/versions/2023-05-01_12-03_9c2a5681ddfd_convert_key_value_entries_to_json.py b/superset/migrations/versions/2023-05-01_12-03_9c2a5681ddfd_convert_key_value_entries_to_json.py
index 6e55f3ddc9..ea285d0808 100644
--- a/superset/migrations/versions/2023-05-01_12-03_9c2a5681ddfd_convert_key_value_entries_to_json.py
+++ b/superset/migrations/versions/2023-05-01_12-03_9c2a5681ddfd_convert_key_value_entries_to_json.py
@@ -17,14 +17,22 @@
 """convert key-value entries to json
 
 Revision ID: 9c2a5681ddfd
+<<<<<<< HEAD
 Revises: f3c2d8ec8595
+=======
+Revises: 7e67aecbf3f1
+>>>>>>> f1fa1a733... chore(key-value): use json serialization for main resources (#23888)
 Create Date: 2023-05-01 12:03:17.079862
 
 """
 
 # revision identifiers, used by Alembic.
 revision = "9c2a5681ddfd"
+<<<<<<< HEAD
 down_revision = "f3c2d8ec8595"
+=======
+down_revision = "7e67aecbf3f1"
+>>>>>>> f1fa1a733... chore(key-value): use json serialization for main resources (#23888)
 
 import io
 import json
@@ -45,10 +53,14 @@ RESOURCES_TO_MIGRATE = ("app", "dashboard_permalink", "explore_permalink")
 
 class RestrictedUnpickler(pickle.Unpickler):
     def find_class(self, module, name):
+<<<<<<< HEAD
         if not (module == "superset.utils.core" and name == "DatasourceType"):
             raise pickle.UnpicklingError(f"Unpickling of {module}.{name} is forbidden")
 
         return super().find_class(module, name)
+=======
+        raise pickle.UnpicklingError(f"Unpickling of {module}.{name} is forbidden")
+>>>>>>> f1fa1a733... chore(key-value): use json serialization for main resources (#23888)
 
 
 class KeyValueEntry(Base):
@@ -61,12 +73,16 @@ class KeyValueEntry(Base):
 def upgrade():
     bind = op.get_bind()
     session: Session = db.Session(bind=bind)
+<<<<<<< HEAD
     truncated_count = 0
+=======
+>>>>>>> f1fa1a733... chore(key-value): use json serialization for main resources (#23888)
     for entry in paginated_update(
         session.query(KeyValueEntry).filter(
             KeyValueEntry.resource.in_(RESOURCES_TO_MIGRATE)
         )
     ):
+<<<<<<< HEAD
         try:
             value = RestrictedUnpickler(io.BytesIO(entry.value)).load() or {}
         except pickle.UnpicklingError as ex:
@@ -83,6 +99,11 @@ def upgrade():
     if truncated_count:
         print(f"Replaced {truncated_count} corrupted values with an empty value")
 
+=======
+        value = RestrictedUnpickler(io.BytesIO(entry.value)).load() or {}
+        entry.value = bytes(json.dumps(value), encoding="utf-8")
+
+>>>>>>> f1fa1a733... chore(key-value): use json serialization for main resources (#23888)
 
 def downgrade():
     bind = op.get_bind()
diff --git a/tests/integration_tests/explore/permalink/api_tests.py b/tests/integration_tests/explore/permalink/api_tests.py
index 3a07bd977a..c7b631239e 100644
--- a/tests/integration_tests/explore/permalink/api_tests.py
+++ b/tests/integration_tests/explore/permalink/api_tests.py
@@ -24,7 +24,11 @@ from sqlalchemy.orm import Session
 from superset import db
 from superset.explore.permalink.schemas import ExplorePermalinkSchema
 from superset.key_value.models import KeyValueEntry
+<<<<<<< HEAD
 from superset.key_value.types import KeyValueResource, MarshmallowKeyValueCodec
+=======
+from superset.key_value.types import JsonKeyValueCodec, KeyValueResource
+>>>>>>> f1fa1a733... chore(key-value): use json serialization for main resources (#23888)
 from superset.key_value.utils import decode_permalink_id, encode_permalink_key
 from superset.models.slice import Slice
 from superset.utils.core import DatasourceType
@@ -95,7 +99,11 @@ def test_get_missing_chart(
     chart_id = 1234
     entry = KeyValueEntry(
         resource=KeyValueResource.EXPLORE_PERMALINK,
+<<<<<<< HEAD
         value=MarshmallowKeyValueCodec(ExplorePermalinkSchema()).encode(
+=======
+        value=JsonKeyValueCodec().encode(
+>>>>>>> f1fa1a733... chore(key-value): use json serialization for main resources (#23888)
             {
                 "chartId": chart_id,
                 "datasourceId": chart.datasource.id,