superset2/superset/commands/utils.py

# 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.
from __future__ import annotations

from collections import Counter
from typing import Any, Optional, TYPE_CHECKING

from flask import current_app, has_app_context
from flask_appbuilder.models.sqla import Model
from marshmallow import ValidationError

from superset import security_manager
from superset.commands.exceptions import (
    DatasourceNotFoundValidationError,
    TagForbiddenError,
    TagNotFoundValidationError,
)
from superset.daos.datasource import DatasourceDAO
from superset.daos.exceptions import DatasourceNotFound
from superset.daos.tag import TagDAO
from superset.exceptions import SupersetSecurityException
from superset.subjects.exceptions import SubjectsNotFoundValidationError
from superset.subjects.models import Subject
from superset.subjects.utils import (
    get_default_viewers_for_new_asset,
    get_or_create_user_subject as get_user_subject,
    get_subject,
    get_user_subject_ids,
)
from superset.tags.models import ObjectType, Tag, TagType
from superset.utils import json
from superset.utils.core import DatasourceType, get_user_id

if TYPE_CHECKING:
    from superset.connectors.sqla.models import BaseDatasource


def _has_extra_editors_resolver() -> bool:
    return bool(has_app_context() and current_app.config.get("EXTRA_EDITORS_RESOLVER"))


def current_user_can_modify_object(model: Any) -> bool:
    """Whether the current user is authorized to create/modify ``model``.

    Delegates to ``security_manager.raise_for_editorship``, which grants
    access to admins and any subject in ``model.editors`` (when that
    relationship exists). For models that don't carry an ``editors``
    relationship, or when the current subject isn't one of them, this falls
    back to allowing the object's creator.

    Callers that need to distinguish "not found" from "no access" should
    look the model up bypassing DAO base filters (e.g.
    ``skip_base_filter=True``) before calling this, so an object the user
    cannot access reaches the check instead of resolving to ``None`` and
    being written without any check.
    """
    try:
        security_manager.raise_for_editorship(model)
        return True
    except SupersetSecurityException:
        return bool(
            model.created_by and model.created_by == security_manager.current_user
        )


def populate_subject_list(
    subject_ids: list[int] | None,
    default_to_user: bool,
    ensure_no_lockout: bool = False,
    field_name: str = "subjects",
) -> list[Subject]:
    """
    Helper function for commands, will fetch all subjects from subject id's.

    :param subject_ids: list of subject id's
    :param default_to_user: add current user's subject if list is empty
    :param ensure_no_lockout: prevent non-admins from removing themselves
    :param field_name: field name for validation errors
    :raises SubjectsNotFoundValidationError: if a subject id can't be resolved
    :returns: Final list of subjects
    """
    subject_ids = subject_ids or []
    subjects: list[Subject] = []
    user_id = get_user_id()

    if not subject_ids and default_to_user:
        if user_id:
            user_subject = get_user_subject(user_id)
            return [user_subject] if user_subject else []
        return []

    if ensure_no_lockout and not security_manager.is_admin() and user_id:
        user_subject = get_user_subject(user_id)
        if (
            user_subject
            and user_subject.id not in subject_ids
            and not _has_extra_editors_resolver()
        ):
            user_subject_ids = set(get_user_subject_ids(user_id))
            if not (user_subject_ids & set(subject_ids)):
                subjects.append(user_subject)

    for sid in subject_ids:
        subject = get_subject(sid)
        if not subject:
            raise SubjectsNotFoundValidationError(field_name)
        subjects.append(subject)
    return subjects


def compute_subject_list(
    current_subjects: list[Subject] | None,
    new_subject_ids: list[int] | None,
    ensure_no_lockout: bool = False,
    field_name: str = "subjects",
) -> list[Subject]:
    """
    Helper function for update commands, to properly handle the subjects list.
    Preserve the previous configuration unless included in the update payload.

    :param current_subjects: list of current subjects
    :param new_subject_ids: list of new subject ids specified in the update payload
    :param ensure_no_lockout: prevent non-admins from removing themselves
    :param field_name: field name for validation errors
    :returns: Final list of subjects
    """
    current_subjects = current_subjects or []
    subject_ids = (
        [s.id for s in current_subjects] if new_subject_ids is None else new_subject_ids
    )
    return populate_subject_list(
        subject_ids,
        default_to_user=False,
        ensure_no_lockout=ensure_no_lockout,
        field_name=field_name,
    )


def populate_subjects(
    properties: dict[str, Any],
    exceptions: list[ValidationError],
    *,
    include_viewers: bool = True,
) -> None:
    """
    Populate editors (and optionally viewers) on a properties dict for
    create commands.

    Appends any validation errors to *exceptions*.
    """
    try:
        properties["editors"] = populate_subject_list(
            properties.get("editors"),
            default_to_user=True,
            ensure_no_lockout=True,
            field_name="editors",
        )
    except ValidationError as ex:
        exceptions.append(ex)

    if not include_viewers:
        return

    if properties.get("viewers") is not None:
        try:
            properties["viewers"] = populate_subject_list(
                properties["viewers"],
                default_to_user=False,
                field_name="viewers",
            )
        except ValidationError as ex:
            exceptions.append(ex)
    elif default_viewers := get_default_viewers_for_new_asset(get_user_id()):
        properties["viewers"] = default_viewers


def compute_subjects(
    model: Model,
    properties: dict[str, Any],
    exceptions: list[ValidationError],
    *,
    include_viewers: bool = True,
) -> None:
    """
    Compute editors (and optionally viewers) on a properties dict for
    update commands, preserving existing values when the payload omits them.

    Appends any validation errors to *exceptions*.
    """
    try:
        properties["editors"] = compute_subject_list(
            model.editors,
            properties.get("editors"),
            ensure_no_lockout=True,
            field_name="editors",
        )
    except ValidationError as ex:
        exceptions.append(ex)

    if include_viewers and hasattr(model, "viewers"):
        try:
            properties["viewers"] = compute_subject_list(
                model.viewers,
                properties.get("viewers"),
                field_name="viewers",
            )
        except ValidationError as ex:
            exceptions.append(ex)


def get_datasource_by_id(datasource_id: int, datasource_type: str) -> BaseDatasource:
    try:
        return DatasourceDAO.get_datasource(
            DatasourceType(datasource_type), datasource_id
        )
    except DatasourceNotFound as ex:
        raise DatasourceNotFoundValidationError() from ex


def validate_tags(
    object_type: ObjectType,
    current_tags: list[Tag],
    new_tag_ids: Optional[list[int]],
) -> None:
    """
    Helper function for update commands, to validate the tags list. Users
    with `can_write` on `Tag` are allowed to both create new tags and manage
    tag association with objects. Users with `can_tag` on `object_type` are
    only allowed to manage existing existing tags' associations with the object.

    :param current_tags: list of current tags
    :param new_tag_ids: list of tags specified in the update payload
    """

    # `tags` not part of the update payload
    if new_tag_ids is None:
        return

    # No changes in the list
    current_custom_tags = [tag.id for tag in current_tags if tag.type == TagType.custom]
    if Counter(current_custom_tags) == Counter(new_tag_ids):
        return

    # No perm to tags assets
    if not (
        security_manager.can_access("can_write", "Tag")
        or security_manager.can_access("can_tag", object_type.name.capitalize())
    ):
        validation_error = (
            f"You do not have permission to manage tags on {object_type.name}s"
        )
        raise TagForbiddenError(validation_error)

    # Validate if new tags already exist
    additional_tags = [tag for tag in new_tag_ids if tag not in current_custom_tags]
    for tag_id in additional_tags:
        if not TagDAO.find_by_id(tag_id):
            validation_error = f"Tag ID {tag_id} not found"
            raise TagNotFoundValidationError(validation_error)

    return


def update_tags(
    object_type: ObjectType,
    object_id: int,
    current_tags: list[Tag],
    new_tag_ids: list[int],
) -> None:
    """
    Helper function for update commands, to update the tag relationship.

    :param object_id: The object (dashboard, chart, etc) ID
    :param object_type: The object type
    :param current_tags: list of current tags
    :param new_tag_ids: list of tags specified in the update payload
    """

    current_custom_tags = [tag for tag in current_tags if tag.type == TagType.custom]
    current_custom_tag_ids = [
        tag.id for tag in current_tags if tag.type == TagType.custom
    ]

    tags_to_delete = [tag for tag in current_custom_tags if tag.id not in new_tag_ids]
    for tag in tags_to_delete:
        TagDAO.delete_tagged_object(object_type, object_id, tag.name)

    tag_ids_to_add = [
        tag_id for tag_id in new_tag_ids if tag_id not in current_custom_tag_ids
    ]
    if tag_ids_to_add:
        tags_to_add = TagDAO.find_by_ids(tag_ids_to_add)
        TagDAO.create_custom_tagged_objects(
            object_type, object_id, [tag.name for tag in tags_to_add]
        )


def update_chart_config_dataset(
    config: dict[str, Any], dataset_info: dict[str, Any]
) -> dict[str, Any]:
    """
    Update the chart configuration and query_context with new dataset information

    :param config: The original chart configuration
    :param dataset_info: Dict with datasource_id, datasource_type, and datasource_name
    :return: The updated chart configuration
    """
    # Update datasource id, type, and name
    config.update(dataset_info)

    dataset_uid = f"{dataset_info['datasource_id']}__{dataset_info['datasource_type']}"
    config["params"].update({"datasource": dataset_uid})

    if "query_context" in config and config["query_context"] is not None:
        try:
            query_context = json.loads(config["query_context"])

            query_context["datasource"] = {
                "id": dataset_info["datasource_id"],
                "type": dataset_info["datasource_type"],
            }

            if "form_data" in query_context:
                query_context["form_data"]["datasource"] = dataset_uid

            if "queries" in query_context:
                for query in query_context["queries"]:
                    if "datasource" in query:
                        query["datasource"] = query_context["datasource"]

            config["query_context"] = json.dumps(query_context)
        except json.JSONDecodeError:
            config["query_context"] = None

    return config