# 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 import logging from datetime import datetime from typing import Dict, List from flask_appbuilder.models.sqla.interface import SQLAInterface from sqlalchemy import or_, select from sqlalchemy.orm import Query from superset.charts.filters import ChartFilter from superset.commands.chart.exceptions import ChartNotFoundError from superset.daos.base import BaseDAO, ColumnOperator, ColumnOperatorEnum from superset.extensions import db from superset.models.core import FavStar, FavStarClassName from superset.models.slice import id_or_uuid_filter, Slice, slice_user from superset.utils.core import get_user_id logger = logging.getLogger(__name__) # Custom filterable fields for charts CHART_CUSTOM_FIELDS = { "viz_type": ["eq", "in", "like"], "datasource_name": ["eq", "in", "like"], "owner": ["eq", "in"], } class ChartDAO(BaseDAO[Slice]): base_filter = ChartFilter @classmethod def apply_column_operators( cls, query: Query, column_operators: list[ColumnOperator] | None = None, ) -> Query: """Override to handle owner filter via the slice_user M2M table.""" if not column_operators: return query remaining_operators: list[ColumnOperator] = [] for c in column_operators: if not isinstance(c, ColumnOperator): c = ColumnOperator.model_validate(c) if c.col == "owner": operator_enum = ColumnOperatorEnum(c.opr) subq = select(slice_user.c.slice_id).where( operator_enum.apply(slice_user.c.user_id, c.value) ) query = query.filter( Slice.id.in_(subq) # type: ignore[attr-defined,unused-ignore] ) elif c.col == "created_by_fk_or_owner": if c.opr != "eq": raise ValueError( f"created_by_fk_or_owner only supports 'eq'; got '{c.opr}'" ) owner_subq = select(slice_user.c.slice_id).where( slice_user.c.user_id == c.value ) query = query.filter( or_( Slice.created_by_fk == c.value, # type: ignore[attr-defined,unused-ignore] Slice.id.in_(owner_subq), # type: ignore[attr-defined,unused-ignore] ) ) else: remaining_operators.append(c) if remaining_operators: query = super().apply_column_operators(query, remaining_operators) return query @classmethod def get_filterable_columns_and_operators(cls) -> Dict[str, List[str]]: filterable = super().get_filterable_columns_and_operators() # Add custom fields for charts filterable.update(CHART_CUSTOM_FIELDS) return filterable @staticmethod def get_by_id_or_uuid(id_or_uuid: str) -> Slice: query = db.session.query(Slice).filter(id_or_uuid_filter(id_or_uuid)) # Apply chart base filters query = ChartFilter("id", SQLAInterface(Slice, db.session)).apply(query, None) chart = query.one_or_none() if not chart: raise ChartNotFoundError() return chart @staticmethod def favorited_ids(charts: list[Slice]) -> list[FavStar]: ids = [chart.id for chart in charts] return [ star.obj_id for star in db.session.query(FavStar.obj_id) .filter( FavStar.class_name == FavStarClassName.CHART, FavStar.obj_id.in_(ids), FavStar.user_id == get_user_id(), ) .all() ] @staticmethod def add_favorite(chart: Slice) -> None: ids = ChartDAO.favorited_ids([chart]) if chart.id not in ids: db.session.add( FavStar( class_name=FavStarClassName.CHART, obj_id=chart.id, user_id=get_user_id(), dttm=datetime.now(), ) ) @staticmethod def remove_favorite(chart: Slice) -> None: fav = ( db.session.query(FavStar) .filter( FavStar.class_name == FavStarClassName.CHART, FavStar.obj_id == chart.id, FavStar.user_id == get_user_id(), ) .one_or_none() ) if fav: db.session.delete(fav)