# 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. import logging import pickle from datetime import datetime from typing import Any, Optional, Union from uuid import UUID from flask_appbuilder.security.sqla.models import User from sqlalchemy.exc import SQLAlchemyError 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 logger = logging.getLogger(__name__) class CreateKeyValueCommand(BaseCommand): actor: Optional[User] resource: KeyValueResource value: Any key: Optional[Union[int, UUID]] expires_on: Optional[datetime] # pylint: disable=too-many-arguments def __init__( self, resource: KeyValueResource, value: Any, actor: Optional[User] = None, key: Optional[Union[int, UUID]] = None, expires_on: Optional[datetime] = None, ): """ Create a new key-value pair :param resource: the resource (dashboard, chart etc) :param value: the value to persist in the key-value store :param actor: the user performing the command :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.actor = actor self.value = value self.key = key self.expires_on = expires_on def run(self) -> Key: try: return self.create() except SQLAlchemyError as ex: db.session.rollback() logger.exception("Error running create command") raise KeyValueCreateFailedError() from ex def validate(self) -> None: pass def create(self) -> Key: entry = KeyValueEntry( resource=self.resource.value, value=pickle.dumps(self.value), created_on=datetime.now(), created_by_fk=None if self.actor is None or self.actor.is_anonymous else self.actor.id, expires_on=self.expires_on, ) if self.key is not None: try: if isinstance(self.key, UUID): entry.uuid = self.key else: entry.id = self.key except ValueError as ex: raise KeyValueCreateFailedError() from ex db.session.add(entry) db.session.commit() return Key(id=entry.id, uuid=entry.uuid)