You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by el...@apache.org on 2023/07/07 22:41:34 UTC
[superset] 04/12: merge in fix with migration (#24314)
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 9abe28bc09f42ef434e9a1fb1e7ab62f5750076d
Author: Elizabeth Thompson <es...@gmail.com>
AuthorDate: Wed Jun 7 13:41:32 2023 -0700
merge in fix with migration (#24314)
Co-authored-by: Ville Brofeldt <33...@users.noreply.github.com>
Co-authored-by: Ville Brofeldt <vi...@apple.com>
---
superset/dashboards/permalink/commands/base.py | 3 +-
superset/dashboards/permalink/commands/create.py | 1 +
superset/dashboards/permalink/commands/get.py | 6 +-
superset/explore/permalink/commands/base.py | 3 +-
superset/explore/permalink/commands/create.py | 3 +-
superset/explore/permalink/commands/get.py | 1 +
superset/extensions/metastore_cache.py | 11 ++-
superset/key_value/commands/create.py | 23 ++++--
superset/key_value/commands/get.py | 15 +++-
superset/key_value/commands/update.py | 11 ++-
superset/key_value/commands/upsert.py | 13 +--
superset/key_value/shared_entries.py | 12 ++-
superset/key_value/types.py | 33 +++++++-
...2a5681ddfd_convert_key_value_entries_to_json.py | 96 ++++++++++++++++++++++
superset/temporary_cache/api.py | 13 ++-
superset/temporary_cache/commands/parameters.py | 3 +
.../explore/permalink/api_tests.py | 5 +-
.../key_value/commands/create_test.py | 55 +++++++++++--
.../key_value/commands/delete_test.py | 13 +--
.../key_value/commands/fixtures.py | 15 +++-
.../key_value/commands/get_test.py | 25 +++---
.../key_value/commands/update_test.py | 11 ++-
.../key_value/commands/upsert_test.py | 11 ++-
23 files changed, 311 insertions(+), 71 deletions(-)
diff --git a/superset/dashboards/permalink/commands/base.py b/superset/dashboards/permalink/commands/base.py
index f4dc4f0726..82e24264ca 100644
--- a/superset/dashboards/permalink/commands/base.py
+++ b/superset/dashboards/permalink/commands/base.py
@@ -18,11 +18,12 @@ from abc import ABC
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:
diff --git a/superset/dashboards/permalink/commands/create.py b/superset/dashboards/permalink/commands/create.py
index 51dac2d5de..2b6151fbb2 100644
--- a/superset/dashboards/permalink/commands/create.py
+++ b/superset/dashboards/permalink/commands/create.py
@@ -58,6 +58,7 @@ class CreateDashboardPermalinkCommand(BaseDashboardPermalinkCommand):
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)
diff --git a/superset/dashboards/permalink/commands/get.py b/superset/dashboards/permalink/commands/get.py
index f89f9444e7..4206263a37 100644
--- a/superset/dashboards/permalink/commands/get.py
+++ b/superset/dashboards/permalink/commands/get.py
@@ -39,7 +39,11 @@ class GetDashboardPermalinkCommand(BaseDashboardPermalinkCommand):
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"])
diff --git a/superset/explore/permalink/commands/base.py b/superset/explore/permalink/commands/base.py
index bef9546e21..a87183b7e9 100644
--- a/superset/explore/permalink/commands/base.py
+++ b/superset/explore/permalink/commands/base.py
@@ -18,11 +18,12 @@ from abc import ABC
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:
diff --git a/superset/explore/permalink/commands/create.py b/superset/explore/permalink/commands/create.py
index 77ce04c4e4..21c0f4e42f 100644
--- a/superset/explore/permalink/commands/create.py
+++ b/superset/explore/permalink/commands/create.py
@@ -45,13 +45,14 @@ class CreateExplorePermalinkCommand(BaseExplorePermalinkCommand):
value = {
"chartId": self.chart_id,
"datasourceId": datasource_id,
- "datasourceType": datasource_type,
+ "datasourceType": datasource_type.value,
"datasource": self.datasource,
"state": self.state,
}
command = CreateKeyValueCommand(
resource=self.resource,
value=value,
+ codec=self.codec,
)
key = command.run()
if key.id is None:
diff --git a/superset/explore/permalink/commands/get.py b/superset/explore/permalink/commands/get.py
index 3376cab080..4823117ece 100644
--- a/superset/explore/permalink/commands/get.py
+++ b/superset/explore/permalink/commands/get.py
@@ -43,6 +43,7 @@ class GetExplorePermalinkCommand(BaseExplorePermalinkCommand):
value: Optional[ExplorePermalinkValue] = GetKeyValueCommand(
resource=self.resource,
key=key,
+ codec=self.codec,
).run()
if value:
chart_id: Optional[int] = value.get("chartId")
diff --git a/superset/extensions/metastore_cache.py b/superset/extensions/metastore_cache.py
index 1e5cff7ee3..f69276c908 100644
--- a/superset/extensions/metastore_cache.py
+++ b/superset/extensions/metastore_cache.py
@@ -23,10 +23,11 @@ from flask import Flask
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 @@ class SupersetMetastoreCache(BaseCache):
resource=RESOURCE,
key=self.get_key(key),
value=value,
+ codec=CODEC,
expires_on=self._get_expiry(timeout),
).run()
return True
@@ -80,6 +82,7 @@ class SupersetMetastoreCache(BaseCache):
CreateKeyValueCommand(
resource=RESOURCE,
value=value,
+ codec=CODEC,
key=self.get_key(key),
expires_on=self._get_expiry(timeout),
).run()
@@ -92,7 +95,11 @@ class SupersetMetastoreCache(BaseCache):
# 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)
diff --git a/superset/key_value/commands/create.py b/superset/key_value/commands/create.py
index 93e99c223b..d66d99d6e9 100644
--- a/superset/key_value/commands/create.py
+++ b/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 import db
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 @@ logger = logging.getLogger(__name__)
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 @@ class CreateKeyValueCommand(BaseCommand):
: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 @@ class CreateKeyValueCommand(BaseCommand):
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,
diff --git a/superset/key_value/commands/get.py b/superset/key_value/commands/get.py
index 44c02331cc..9d659f3bc7 100644
--- a/superset/key_value/commands/get.py
+++ b/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 import db
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 @@ logger = logging.getLogger(__name__)
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 @@ class GetKeyValueCommand(BaseCommand):
.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
diff --git a/superset/key_value/commands/update.py b/superset/key_value/commands/update.py
index b69ca5e70d..becd6d9ca8 100644
--- a/superset/key_value/commands/update.py
+++ b/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 import db
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 @@ logger = logging.getLogger(__name__)
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 @@ class UpdateKeyValueCommand(BaseCommand):
: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 @@ class UpdateKeyValueCommand(BaseCommand):
.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()
diff --git a/superset/key_value/commands/upsert.py b/superset/key_value/commands/upsert.py
index 06b33c90fc..c5668f1161 100644
--- a/superset/key_value/commands/upsert.py
+++ b/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 @@ from superset.key_value.exceptions import (
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 @@ class UpsertKeyValueCommand(BaseCommand):
: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 @@ class UpsertKeyValueCommand(BaseCommand):
.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 @@ class UpsertKeyValueCommand(BaseCommand):
return CreateKeyValueCommand(
resource=self.resource,
value=self.value,
+ codec=self.codec,
key=self.key,
expires_on=self.expires_on,
).run()
diff --git a/superset/key_value/shared_entries.py b/superset/key_value/shared_entries.py
index 5f4ded9498..7895b75907 100644
--- a/superset/key_value/shared_entries.py
+++ b/superset/key_value/shared_entries.py
@@ -18,11 +18,12 @@
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]:
@@ -30,7 +31,7 @@ def get_shared_value(key: SharedKey) -> Optional[Any]:
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:
@@ -38,7 +39,12 @@ def set_shared_value(key: SharedKey, value: Any) -> None:
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:
diff --git a/superset/key_value/types.py b/superset/key_value/types.py
index c3064fbef4..07d06414f6 100644
--- a/superset/key_value/types.py
+++ b/superset/key_value/types.py
@@ -14,9 +14,14 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+from __future__ import annotations
+
+import json
+import pickle
+from abc import ABC, abstractmethod
from dataclasses import dataclass
from enum import Enum
-from typing import Optional, TypedDict
+from typing import Any, Optional, TypedDict
from uuid import UUID
@@ -42,3 +47,29 @@ class KeyValueResource(str, Enum):
class SharedKey(str, Enum):
DASHBOARD_PERMALINK_SALT = "dashboard_permalink_salt"
EXPLORE_PERMALINK_SALT = "explore_permalink_salt"
+
+
+class KeyValueCodec(ABC):
+ @abstractmethod
+ def encode(self, value: Any) -> bytes:
+ ...
+
+ @abstractmethod
+ 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")
+
+ def decode(self, value: bytes) -> dict[Any, Any]:
+ return json.loads(value)
+
+
+class PickleKeyValueCodec(KeyValueCodec):
+ def encode(self, value: dict[Any, Any]) -> bytes:
+ return pickle.dumps(value)
+
+ def decode(self, value: bytes) -> dict[Any, Any]:
+ return pickle.loads(value)
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
new file mode 100644
index 0000000000..6e55f3ddc9
--- /dev/null
+++ b/superset/migrations/versions/2023-05-01_12-03_9c2a5681ddfd_convert_key_value_entries_to_json.py
@@ -0,0 +1,96 @@
+# 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.
+"""convert key-value entries to json
+
+Revision ID: 9c2a5681ddfd
+Revises: f3c2d8ec8595
+Create Date: 2023-05-01 12:03:17.079862
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = "9c2a5681ddfd"
+down_revision = "f3c2d8ec8595"
+
+import io
+import json
+import pickle
+
+from alembic import op
+from sqlalchemy import Column, Integer, LargeBinary, String
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import Session
+
+from superset import db
+from superset.migrations.shared.utils import paginated_update
+
+Base = declarative_base()
+VALUE_MAX_SIZE = 2**24 - 1
+RESOURCES_TO_MIGRATE = ("app", "dashboard_permalink", "explore_permalink")
+
+
+class RestrictedUnpickler(pickle.Unpickler):
+ def find_class(self, module, name):
+ 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)
+
+
+class KeyValueEntry(Base):
+ __tablename__ = "key_value"
+ id = Column(Integer, primary_key=True)
+ resource = Column(String(32), nullable=False)
+ value = Column(LargeBinary(length=VALUE_MAX_SIZE), nullable=False)
+
+
+def upgrade():
+ bind = op.get_bind()
+ session: Session = db.Session(bind=bind)
+ truncated_count = 0
+ for entry in paginated_update(
+ session.query(KeyValueEntry).filter(
+ KeyValueEntry.resource.in_(RESOURCES_TO_MIGRATE)
+ )
+ ):
+ try:
+ value = RestrictedUnpickler(io.BytesIO(entry.value)).load() or {}
+ except pickle.UnpicklingError as ex:
+ if str(ex) == "pickle data was truncated":
+ # make truncated values that were created prior to #20385 an empty
+ # dict so that downgrading will work properly.
+ truncated_count += 1
+ value = {}
+ else:
+ raise
+
+ entry.value = bytes(json.dumps(value), encoding="utf-8")
+
+ if truncated_count:
+ print(f"Replaced {truncated_count} corrupted values with an empty value")
+
+
+def downgrade():
+ bind = op.get_bind()
+ session: Session = db.Session(bind=bind)
+ for entry in paginated_update(
+ session.query(KeyValueEntry).filter(
+ KeyValueEntry.resource.in_(RESOURCES_TO_MIGRATE)
+ ),
+ ):
+ value = json.loads(entry.value) or {}
+ entry.value = pickle.dumps(value)
diff --git a/superset/temporary_cache/api.py b/superset/temporary_cache/api.py
index b6376c63c3..85db65c62c 100644
--- a/superset/temporary_cache/api.py
+++ b/superset/temporary_cache/api.py
@@ -24,6 +24,7 @@ from flask import request, Response
from marshmallow import ValidationError
from superset.constants import MODEL_API_RW_METHOD_PERMISSION_MAP, RouteMethod
+from superset.key_value.types import JsonKeyValueCodec
from superset.temporary_cache.commands.exceptions import (
TemporaryCacheAccessDeniedError,
TemporaryCacheResourceNotFoundError,
@@ -37,6 +38,8 @@ from superset.views.base_api import BaseSupersetApi, requires_json
logger = logging.getLogger(__name__)
+CODEC = JsonKeyValueCodec()
+
class TemporaryCacheRestApi(BaseSupersetApi, ABC):
add_model_schema = TemporaryCachePostSchema()
@@ -69,7 +72,12 @@ class TemporaryCacheRestApi(BaseSupersetApi, ABC):
try:
item = self.add_model_schema.load(request.json)
tab_id = request.args.get("tab_id")
- args = CommandParameters(resource_id=pk, value=item["value"], tab_id=tab_id)
+ args = CommandParameters(
+ resource_id=pk,
+ value=item["value"],
+ tab_id=tab_id,
+ codec=CODEC,
+ )
key = self.get_create_command()(args).run()
return self.response(201, key=key)
except ValidationError as ex:
@@ -89,6 +97,7 @@ class TemporaryCacheRestApi(BaseSupersetApi, ABC):
key=key,
value=item["value"],
tab_id=tab_id,
+ codec=CODEC,
)
key = self.get_update_command()(args).run()
return self.response(200, key=key)
@@ -101,7 +110,7 @@ class TemporaryCacheRestApi(BaseSupersetApi, ABC):
def get(self, pk: int, key: str) -> Response:
try:
- args = CommandParameters(resource_id=pk, key=key)
+ args = CommandParameters(resource_id=pk, key=key, codec=CODEC)
value = self.get_get_command()(args).run()
if not value:
return self.response_404()
diff --git a/superset/temporary_cache/commands/parameters.py b/superset/temporary_cache/commands/parameters.py
index 74b9c1c632..e4e5b9b06a 100644
--- a/superset/temporary_cache/commands/parameters.py
+++ b/superset/temporary_cache/commands/parameters.py
@@ -17,10 +17,13 @@
from dataclasses import dataclass
from typing import Optional
+from superset.key_value.types import KeyValueCodec
+
@dataclass
class CommandParameters:
resource_id: int
+ codec: Optional[KeyValueCodec] = None
tab_id: Optional[int] = None
key: Optional[str] = None
value: Optional[str] = None
diff --git a/tests/integration_tests/explore/permalink/api_tests.py b/tests/integration_tests/explore/permalink/api_tests.py
index 22a36f41e1..4c6a3c12dd 100644
--- a/tests/integration_tests/explore/permalink/api_tests.py
+++ b/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,
diff --git a/tests/integration_tests/key_value/commands/create_test.py b/tests/integration_tests/key_value/commands/create_test.py
index 0e789026ba..a2ee3d13ae 100644
--- a/tests/integration_tests/key_value/commands/create_test.py
+++ b/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()
diff --git a/tests/integration_tests/key_value/commands/delete_test.py b/tests/integration_tests/key_value/commands/delete_test.py
index 62f9883370..3c4892faa6 100644
--- a/tests/integration_tests/key_value/commands/delete_test.py
+++ b/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
diff --git a/tests/integration_tests/key_value/commands/fixtures.py b/tests/integration_tests/key_value/commands/fixtures.py
index 2fd4fde4e1..66aea8a4ed 100644
--- a/tests/integration_tests/key_value/commands/fixtures.py
+++ b/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()
diff --git a/tests/integration_tests/key_value/commands/get_test.py b/tests/integration_tests/key_value/commands/get_test.py
index b1800a4c3b..28a6dd73d5 100644
--- a/tests/integration_tests/key_value/commands/get_test.py
+++ b/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()
diff --git a/tests/integration_tests/key_value/commands/update_test.py b/tests/integration_tests/key_value/commands/update_test.py
index 8eb03b4eda..2c0fc3e31d 100644
--- a/tests/integration_tests/key_value/commands/update_test.py
+++ b/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
diff --git a/tests/integration_tests/key_value/commands/upsert_test.py b/tests/integration_tests/key_value/commands/upsert_test.py
index e5cd27e3a6..c26b66d02e 100644
--- a/tests/integration_tests/key_value/commands/upsert_test.py
+++ b/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