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

tests/integration_tests/explore/permalink/api_tests.py

@@ -15,7 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 import json
-import pickle
 from typing import Any, Dict, Iterator
 from uuid import uuid3
 
@@ -24,7 +23,7 @@ from sqlalchemy.orm import Session
 
 from superset import db
 from superset.key_value.models import KeyValueEntry
-from superset.key_value.types import KeyValueResource
+from superset.key_value.types import JsonKeyValueCodec, KeyValueResource
 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 +94,7 @@ def test_get_missing_chart(
     chart_id = 1234
     entry = KeyValueEntry(
         resource=KeyValueResource.EXPLORE_PERMALINK,
-        value=pickle.dumps(
+        value=JsonKeyValueCodec().encode(
             {
                 "chartId": chart_id,
                 "datasourceId": chart.datasource.id,

tests/integration_tests/key_value/commands/create_test.py

@@ -16,20 +16,23 @@
 # under the License.
 from __future__ import annotations
 
+import json
 import pickle
-from uuid import UUID
 
+import pytest
 from flask.ctx import AppContext
 from flask_appbuilder.security.sqla.models import User
 
 from superset.extensions import db
+from superset.key_value.exceptions import KeyValueCreateFailedError
 from superset.utils.core import override_user
 from tests.integration_tests.key_value.commands.fixtures import (
     admin,
-    ID_KEY,
+    JSON_CODEC,
+    JSON_VALUE,
+    PICKLE_CODEC,
+    PICKLE_VALUE,
     RESOURCE,
-    UUID_KEY,
-    VALUE,
 )
 
 
@@ -38,11 +41,15 @@ def test_create_id_entry(app_context: AppContext, admin: User) -> None:
     from superset.key_value.models import KeyValueEntry
 
     with override_user(admin):
-        key = CreateKeyValueCommand(resource=RESOURCE, value=VALUE).run()
+        key = CreateKeyValueCommand(
+            resource=RESOURCE,
+            value=JSON_VALUE,
+            codec=JSON_CODEC,
+        ).run()
         entry = (
             db.session.query(KeyValueEntry).filter_by(id=key.id).autoflush(False).one()
         )
-        assert pickle.loads(entry.value) == VALUE
+        assert json.loads(entry.value) == JSON_VALUE
         assert entry.created_by_fk == admin.id
         db.session.delete(entry)
         db.session.commit()
@@ -53,11 +60,43 @@ def test_create_uuid_entry(app_context: AppContext, admin: User) -> None:
     from superset.key_value.models import KeyValueEntry
 
     with override_user(admin):
-        key = CreateKeyValueCommand(resource=RESOURCE, value=VALUE).run()
+        key = CreateKeyValueCommand(
+            resource=RESOURCE, value=JSON_VALUE, codec=JSON_CODEC
+        ).run()
     entry = (
         db.session.query(KeyValueEntry).filter_by(uuid=key.uuid).autoflush(False).one()
     )
-    assert pickle.loads(entry.value) == VALUE
+    assert json.loads(entry.value) == JSON_VALUE
     assert entry.created_by_fk == admin.id
     db.session.delete(entry)
     db.session.commit()
+
+
+def test_create_fail_json_entry(app_context: AppContext, admin: User) -> None:
+    from superset.key_value.commands.create import CreateKeyValueCommand
+
+    with pytest.raises(KeyValueCreateFailedError):
+        CreateKeyValueCommand(
+            resource=RESOURCE,
+            value=PICKLE_VALUE,
+            codec=JSON_CODEC,
+        ).run()
+
+
+def test_create_pickle_entry(app_context: AppContext, admin: User) -> None:
+    from superset.key_value.commands.create import CreateKeyValueCommand
+    from superset.key_value.models import KeyValueEntry
+
+    with override_user(admin):
+        key = CreateKeyValueCommand(
+            resource=RESOURCE,
+            value=PICKLE_VALUE,
+            codec=PICKLE_CODEC,
+        ).run()
+        entry = (
+            db.session.query(KeyValueEntry).filter_by(id=key.id).autoflush(False).one()
+        )
+        assert type(pickle.loads(entry.value)) == type(PICKLE_VALUE)
+        assert entry.created_by_fk == admin.id
+        db.session.delete(entry)
+        db.session.commit()

tests/integration_tests/key_value/commands/delete_test.py

@@ -16,7 +16,7 @@
 # under the License.
 from __future__ import annotations
 
-import pickle
+import json
 from typing import TYPE_CHECKING
 from uuid import UUID
 
@@ -25,7 +25,11 @@ from flask.ctx import AppContext
 from flask_appbuilder.security.sqla.models import User
 
 from superset.extensions import db
-from tests.integration_tests.key_value.commands.fixtures import admin, RESOURCE, VALUE
+from tests.integration_tests.key_value.commands.fixtures import (
+    admin,
+    JSON_VALUE,
+    RESOURCE,
+)
 
 if TYPE_CHECKING:
     from superset.key_value.models import KeyValueEntry
@@ -42,7 +46,7 @@ def key_value_entry() -> KeyValueEntry:
         id=ID_KEY,
         uuid=UUID_KEY,
         resource=RESOURCE,
-        value=pickle.dumps(VALUE),
+        value=bytes(json.dumps(JSON_VALUE), encoding="utf-8"),
     )
     db.session.add(entry)
     db.session.commit()
@@ -55,7 +59,6 @@ def test_delete_id_entry(
     key_value_entry: KeyValueEntry,
 ) -> None:
     from superset.key_value.commands.delete import DeleteKeyValueCommand
-    from superset.key_value.models import KeyValueEntry
 
     assert DeleteKeyValueCommand(resource=RESOURCE, key=ID_KEY).run() is True
 
@@ -66,7 +69,6 @@ def test_delete_uuid_entry(
     key_value_entry: KeyValueEntry,
 ) -> None:
     from superset.key_value.commands.delete import DeleteKeyValueCommand
-    from superset.key_value.models import KeyValueEntry
 
     assert DeleteKeyValueCommand(resource=RESOURCE, key=UUID_KEY).run() is True
 
@@ -77,6 +79,5 @@ def test_delete_entry_missing(
     key_value_entry: KeyValueEntry,
 ) -> None:
     from superset.key_value.commands.delete import DeleteKeyValueCommand
-    from superset.key_value.models import KeyValueEntry
 
     assert DeleteKeyValueCommand(resource=RESOURCE, key=456).run() is False

tests/integration_tests/key_value/commands/fixtures.py

@@ -17,7 +17,7 @@
 
 from __future__ import annotations
 
-import pickle
+import json
 from typing import Generator, TYPE_CHECKING
 from uuid import UUID
 
@@ -26,7 +26,11 @@ from flask_appbuilder.security.sqla.models import User
 from sqlalchemy.orm import Session
 
 from superset.extensions import db
-from superset.key_value.types import KeyValueResource
+from superset.key_value.types import (
+    JsonKeyValueCodec,
+    KeyValueResource,
+    PickleKeyValueCodec,
+)
 from tests.integration_tests.test_app import app
 
 if TYPE_CHECKING:
@@ -35,7 +39,10 @@ if TYPE_CHECKING:
 ID_KEY = 123
 UUID_KEY = UUID("3e7a2ab8-bcaf-49b0-a5df-dfb432f291cc")
 RESOURCE = KeyValueResource.APP
-VALUE = {"foo": "bar"}
+JSON_VALUE = {"foo": "bar"}
+PICKLE_VALUE = object()
+JSON_CODEC = JsonKeyValueCodec()
+PICKLE_CODEC = PickleKeyValueCodec()
 
 
 @pytest.fixture
@@ -46,7 +53,7 @@ def key_value_entry() -> Generator[KeyValueEntry, None, None]:
         id=ID_KEY,
         uuid=UUID_KEY,
         resource=RESOURCE,
-        value=pickle.dumps(VALUE),
+        value=bytes(json.dumps(JSON_VALUE), encoding="utf-8"),
     )
     db.session.add(entry)
     db.session.commit()

tests/integration_tests/key_value/commands/get_test.py

@@ -16,7 +16,7 @@
 # under the License.
 from __future__ import annotations
 
-import pickle
+import json
 import uuid
 from datetime import datetime, timedelta
 from typing import TYPE_CHECKING
@@ -26,10 +26,11 @@ from flask.ctx import AppContext
 from superset.extensions import db
 from tests.integration_tests.key_value.commands.fixtures import (
     ID_KEY,
+    JSON_CODEC,
+    JSON_VALUE,
     key_value_entry,
     RESOURCE,
     UUID_KEY,
-    VALUE,
 )
 
 if TYPE_CHECKING:
@@ -39,8 +40,8 @@ if TYPE_CHECKING:
 def test_get_id_entry(app_context: AppContext, key_value_entry: KeyValueEntry) -> None:
     from superset.key_value.commands.get import GetKeyValueCommand
 
-    value = GetKeyValueCommand(resource=RESOURCE, key=ID_KEY).run()
-    assert value == VALUE
+    value = GetKeyValueCommand(resource=RESOURCE, key=ID_KEY, codec=JSON_CODEC).run()
+    assert value == JSON_VALUE
 
 
 def test_get_uuid_entry(
@@ -48,8 +49,8 @@ def test_get_uuid_entry(
 ) -> None:
     from superset.key_value.commands.get import GetKeyValueCommand
 
-    value = GetKeyValueCommand(resource=RESOURCE, key=UUID_KEY).run()
-    assert value == VALUE
+    value = GetKeyValueCommand(resource=RESOURCE, key=UUID_KEY, codec=JSON_CODEC).run()
+    assert value == JSON_VALUE
 
 
 def test_get_id_entry_missing(
@@ -58,7 +59,7 @@ def test_get_id_entry_missing(
 ) -> None:
     from superset.key_value.commands.get import GetKeyValueCommand
 
-    value = GetKeyValueCommand(resource=RESOURCE, key=456).run()
+    value = GetKeyValueCommand(resource=RESOURCE, key=456, codec=JSON_CODEC).run()
     assert value is None
 
 
@@ -70,12 +71,12 @@ def test_get_expired_entry(app_context: AppContext) -> None:
         id=678,
         uuid=uuid.uuid4(),
         resource=RESOURCE,
-        value=pickle.dumps(VALUE),
+        value=bytes(json.dumps(JSON_VALUE), encoding="utf-8"),
         expires_on=datetime.now() - timedelta(days=1),
     )
     db.session.add(entry)
     db.session.commit()
-    value = GetKeyValueCommand(resource=RESOURCE, key=ID_KEY).run()
+    value = GetKeyValueCommand(resource=RESOURCE, key=ID_KEY, codec=JSON_CODEC).run()
     assert value is None
     db.session.delete(entry)
     db.session.commit()
@@ -90,12 +91,12 @@ def test_get_future_expiring_entry(app_context: AppContext) -> None:
         id=id_,
         uuid=uuid.uuid4(),
         resource=RESOURCE,
-        value=pickle.dumps(VALUE),
+        value=bytes(json.dumps(JSON_VALUE), encoding="utf-8"),
         expires_on=datetime.now() + timedelta(days=1),
     )
     db.session.add(entry)
     db.session.commit()
-    value = GetKeyValueCommand(resource=RESOURCE, key=id_).run()
-    assert value == VALUE
+    value = GetKeyValueCommand(resource=RESOURCE, key=id_, codec=JSON_CODEC).run()
+    assert value == JSON_VALUE
     db.session.delete(entry)
     db.session.commit()

tests/integration_tests/key_value/commands/update_test.py

@@ -16,9 +16,8 @@
 # under the License.
 from __future__ import annotations
 
-import pickle
+import json
 from typing import TYPE_CHECKING
-from uuid import UUID
 
 from flask.ctx import AppContext
 from flask_appbuilder.security.sqla.models import User
@@ -28,6 +27,7 @@ from superset.utils.core import override_user
 from tests.integration_tests.key_value.commands.fixtures import (
     admin,
     ID_KEY,
+    JSON_CODEC,
     key_value_entry,
     RESOURCE,
     UUID_KEY,
@@ -53,11 +53,12 @@ def test_update_id_entry(
             resource=RESOURCE,
             key=ID_KEY,
             value=NEW_VALUE,
+            codec=JSON_CODEC,
         ).run()
     assert key is not None
     assert key.id == ID_KEY
     entry = db.session.query(KeyValueEntry).filter_by(id=ID_KEY).autoflush(False).one()
-    assert pickle.loads(entry.value) == NEW_VALUE
+    assert json.loads(entry.value) == NEW_VALUE
     assert entry.changed_by_fk == admin.id
 
 
@@ -74,13 +75,14 @@ def test_update_uuid_entry(
             resource=RESOURCE,
             key=UUID_KEY,
             value=NEW_VALUE,
+            codec=JSON_CODEC,
         ).run()
     assert key is not None
     assert key.uuid == UUID_KEY
     entry = (
         db.session.query(KeyValueEntry).filter_by(uuid=UUID_KEY).autoflush(False).one()
     )
-    assert pickle.loads(entry.value) == NEW_VALUE
+    assert json.loads(entry.value) == NEW_VALUE
     assert entry.changed_by_fk == admin.id
 
 
@@ -92,5 +94,6 @@ def test_update_missing_entry(app_context: AppContext, admin: User) -> None:
             resource=RESOURCE,
             key=456,
             value=NEW_VALUE,
+            codec=JSON_CODEC,
         ).run()
     assert key is None

tests/integration_tests/key_value/commands/upsert_test.py

@@ -16,9 +16,8 @@
 # under the License.
 from __future__ import annotations
 
-import pickle
+import json
 from typing import TYPE_CHECKING
-from uuid import UUID
 
 from flask.ctx import AppContext
 from flask_appbuilder.security.sqla.models import User
@@ -28,6 +27,7 @@ from superset.utils.core import override_user
 from tests.integration_tests.key_value.commands.fixtures import (
     admin,
     ID_KEY,
+    JSON_CODEC,
     key_value_entry,
     RESOURCE,
     UUID_KEY,
@@ -53,13 +53,14 @@ def test_upsert_id_entry(
             resource=RESOURCE,
             key=ID_KEY,
             value=NEW_VALUE,
+            codec=JSON_CODEC,
         ).run()
     assert key is not None
     assert key.id == ID_KEY
     entry = (
         db.session.query(KeyValueEntry).filter_by(id=int(ID_KEY)).autoflush(False).one()
     )
-    assert pickle.loads(entry.value) == NEW_VALUE
+    assert json.loads(entry.value) == NEW_VALUE
     assert entry.changed_by_fk == admin.id
 
 
@@ -76,13 +77,14 @@ def test_upsert_uuid_entry(
             resource=RESOURCE,
             key=UUID_KEY,
             value=NEW_VALUE,
+            codec=JSON_CODEC,
         ).run()
     assert key is not None
     assert key.uuid == UUID_KEY
     entry = (
         db.session.query(KeyValueEntry).filter_by(uuid=UUID_KEY).autoflush(False).one()
     )
-    assert pickle.loads(entry.value) == NEW_VALUE
+    assert json.loads(entry.value) == NEW_VALUE
     assert entry.changed_by_fk == admin.id
 
 
@@ -95,6 +97,7 @@ def test_upsert_missing_entry(app_context: AppContext, admin: User) -> None:
             resource=RESOURCE,
             key=456,
             value=NEW_VALUE,
+            codec=JSON_CODEC,
         ).run()
     assert key is not None
     assert key.id == 456