mirror of https://github.com/apache/superset.git
chore(key-value): use json serialization for main resources (#23888)
This commit is contained in:
parent
10d640e940
commit
f1fa1a733d
|
@ -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:
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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"])
|
||||
|
|
|
@ -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:
|
||||
|
|
|
@ -52,6 +52,7 @@ class CreateExplorePermalinkCommand(BaseExplorePermalinkCommand):
|
|||
command = CreateKeyValueCommand(
|
||||
resource=self.resource,
|
||||
value=value,
|
||||
codec=self.codec,
|
||||
)
|
||||
key = command.run()
|
||||
if key.id is None:
|
||||
|
|
|
@ -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")
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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,
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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()
|
||||
|
|
|
@ -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()
|
||||
|
|
|
@ -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:
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -0,0 +1,79 @@
|
|||
# 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: 7e67aecbf3f1
|
||||
Create Date: 2023-05-01 12:03:17.079862
|
||||
|
||||
"""
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "9c2a5681ddfd"
|
||||
down_revision = "7e67aecbf3f1"
|
||||
|
||||
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):
|
||||
raise pickle.UnpicklingError(f"Unpickling of {module}.{name} is forbidden")
|
||||
|
||||
|
||||
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)
|
||||
for entry in paginated_update(
|
||||
session.query(KeyValueEntry).filter(
|
||||
KeyValueEntry.resource.in_(RESOURCES_TO_MIGRATE)
|
||||
)
|
||||
):
|
||||
value = RestrictedUnpickler(io.BytesIO(entry.value)).load() or {}
|
||||
entry.value = bytes(json.dumps(value), encoding="utf-8")
|
||||
|
||||
|
||||
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)
|
|
@ -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()
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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,
|
||||
|
|
|
@ -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()
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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()
|
||||
|
|
|
@ -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()
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
|
Loading…
Reference in New Issue