superset2/superset/subjects/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 typing import TYPE_CHECKING

if TYPE_CHECKING:
    from flask_appbuilder.security.sqla.models import Role, User
    from sqlalchemy.sql import CompoundSelect

from sqlalchemy import select, union_all

from superset import db
from superset.subjects.models import Subject
from superset.subjects.types import SubjectType


def get_user_subject_ids_subquery(user_id: int) -> CompoundSelect:
    """Return a SQLAlchemy Select of all Subject IDs for a user (not executed).

    Includes:
    1. The user's own USER-type subject
    2. ROLE-type subjects for all direct and group-derived roles the user has
    3. GROUP-type subjects for all groups the user belongs to
    """
    from flask_appbuilder.security.sqla.models import (
        assoc_group_role,
        assoc_user_group,
        assoc_user_role,
    )

    user_subj = select(Subject.id).where(Subject.user_id == user_id)

    role_subj = (
        select(Subject.id)
        .join(
            assoc_user_role,
            Subject.role_id == assoc_user_role.c.role_id,
        )
        .where(assoc_user_role.c.user_id == user_id)
    )

    group_subj = (
        select(Subject.id)
        .join(
            assoc_user_group,
            Subject.group_id == assoc_user_group.c.group_id,
        )
        .where(assoc_user_group.c.user_id == user_id)
    )

    group_role_subj = (
        select(Subject.id)
        .join(
            assoc_group_role,
            Subject.role_id == assoc_group_role.c.role_id,
        )
        .join(
            assoc_user_group,
            assoc_group_role.c.group_id == assoc_user_group.c.group_id,
        )
        .where(assoc_user_group.c.user_id == user_id)
    )

    return union_all(user_subj, role_subj, group_subj, group_role_subj)


def get_user_subject_ids(user_id: int) -> list[int]:
    """Return all Subject IDs that a user represents.

    Includes:
    1. The user's own USER-type subject
    2. ROLE-type subjects for all direct and group-derived roles the user has
    3. GROUP-type subjects for all groups the user belongs to
    """
    result = db.session.execute(get_user_subject_ids_subquery(user_id))
    return [row[0] for row in result]


def get_current_user_subject_ids() -> list[int]:
    """Return Subject IDs represented by the current request principal."""
    from flask import g

    from superset.utils.core import get_user_id

    user = getattr(g, "user", None)
    if getattr(user, "is_guest_user", False):
        return [
            subject.id for subject in subjects_from_roles(getattr(user, "roles", []))
        ]

    user_id = get_user_id()
    return get_user_subject_ids(user_id) if user_id else []


def get_user_subject(user_id: int) -> Subject | None:
    """Get the USER-type Subject for a given user ID."""
    return (
        db.session.query(Subject)
        .filter_by(
            user_id=user_id,
            type=SubjectType.USER,
        )
        .first()
    )


def get_or_create_user_subject(user_id: int) -> Subject | None:
    """Get or create the USER-type Subject for a given user ID."""
    if subject := get_user_subject(user_id):
        return subject

    from superset import security_manager
    from superset.subjects.sync import sync_user_subject

    user = db.session.get(security_manager.user_model, user_id)
    if not user:
        return None

    sync_user_subject(user)
    db.session.flush()
    return get_user_subject(user_id)


def get_subject(id_: int) -> Subject | None:
    """Look up a Subject by its primary key."""
    return db.session.get(Subject, id_)


def subjects_from_roles(roles: list[Role | int]) -> list[Subject]:
    """Convert a list of Role objects or role IDs to role-type Subjects.

    Bridges the legacy ``roles`` to the Subject model.
    Accepts both Role objects and plain integer IDs.
    Silently skips roles without a matching Subject row.
    """
    subjects = []
    for role in roles:
        role_id = role if isinstance(role, int) else role.id
        subj = (
            db.session.query(Subject)
            .filter_by(role_id=role_id, type=SubjectType.ROLE)
            .first()
        )
        if subj:
            subjects.append(subj)
    return subjects


def get_user_label(user: User) -> str:
    """Return 'First Last' display name, falling back to username."""
    first = getattr(user, "first_name", None) or ""
    last = getattr(user, "last_name", None) or ""
    full = f"{first} {last}".strip()
    return full or user.username