From d8db1c9230ea54cb3539b8e51fea3d816f151f2c Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Wed, 13 May 2026 22:10:36 +0000 Subject: [PATCH] refactor(mcp): delegate load_user_with_relationships to SecurityManager.find_user_with_relationships Fixes a gap identified in code review: the standalone load_user_with_relationships() in auth.py duplicated SecurityManager.find_user() logic but dropped two FAB behaviors: - auth_username_ci (case-insensitive username lookup) - MultipleResultsFound guard (username uniqueness not guaranteed at DB level in all FAB versions) It also hard-coded User/Group models instead of sm.user_model. Changes: - Add SupersetSecurityManager.find_user_with_relationships() to security/manager.py, mirroring FAB's find_user() (auth_username_ci, MultipleResultsFound handling, self.user_model) and adding eager loading of roles and group.roles via joinedload - Simplify load_user_with_relationships() in auth.py to a thin delegate to the new method, removing the duplicated query logic and raw Group/User imports - Add regression test asserting find_user_with_relationships() exists on the SM Co-Authored-By: Claude Sonnet 4.6 --- superset/mcp_service/auth.py | 41 ++++--------- superset/security/manager.py | 57 ++++++++++++++++++- .../mcp_service/test_auth_api_key.py | 14 +++++ 3 files changed, 80 insertions(+), 32 deletions(-) diff --git a/superset/mcp_service/auth.py b/superset/mcp_service/auth.py index 69c4b75aa72..cb49dc67332 100644 --- a/superset/mcp_service/auth.py +++ b/superset/mcp_service/auth.py @@ -49,7 +49,7 @@ from contextlib import AbstractContextManager from typing import Any, Callable, TYPE_CHECKING, TypeVar from flask import g, has_request_context -from flask_appbuilder.security.sqla.models import Group, User +from flask_appbuilder.security.sqla.models import User from superset.mcp_service.composite_token_verifier import API_KEY_PASSTHROUGH_CLAIM @@ -150,23 +150,14 @@ def check_tool_permission(func: Callable[..., Any]) -> bool: def load_user_with_relationships( username: str | None = None, email: str | None = None ) -> User | None: - """ - Load a user with all relationships needed for permission checks. + """Load a user with roles and group roles eagerly loaded. - This function eagerly loads User.roles, User.groups, and Group.roles - to prevent detached instance errors when the session is closed/rolled back. - - IMPORTANT: Always use this function instead of security_manager.find_user() - when loading users for MCP tool execution. The find_user() method doesn't - eagerly load Group.roles, causing "detached instance" errors when permission - checks access group.roles after the session is rolled back. - - Args: - username: The username to look up (optional if email provided) - email: The email to look up (optional if username provided) - - Returns: - User object with relationships loaded, or None if not found + Delegates to :meth:`SupersetSecurityManager.find_user_with_relationships`, + which mirrors FAB's ``find_user`` (including ``auth_username_ci`` and + ``MultipleResultsFound`` handling) while adding eager loading of + ``User.roles`` and ``User.groups.roles`` to prevent detached-instance + errors when the SQLAlchemy session is closed or rolled back after the + lookup — as happens in MCP tool-execution contexts. Raises: ValueError: If neither username nor email is provided @@ -174,21 +165,9 @@ def load_user_with_relationships( if not username and not email: raise ValueError("Either username or email must be provided") - from sqlalchemy.orm import joinedload + from superset import security_manager - from superset.extensions import db - - query = db.session.query(User).options( - joinedload(User.roles), - joinedload(User.groups).joinedload(Group.roles), - ) - - if username: - query = query.filter(User.username == username) - else: - query = query.filter(User.email == email) - - return query.first() + return security_manager.find_user_with_relationships(username=username, email=email) def _resolve_user_from_jwt_context(app: Any) -> User | None: diff --git a/superset/security/manager.py b/superset/security/manager.py index 6647c59da1f..65e834d92e1 100644 --- a/superset/security/manager.py +++ b/superset/security/manager.py @@ -49,7 +49,8 @@ from flask_login import AnonymousUserMixin, LoginManager from jwt.api_jwt import _jwt_global_obj from sqlalchemy import and_, inspect, or_ from sqlalchemy.engine.base import Connection -from sqlalchemy.orm import eagerload +from sqlalchemy.orm import eagerload, joinedload +from sqlalchemy.orm.exc import MultipleResultsFound from sqlalchemy.orm.mapper import Mapper from sqlalchemy.orm.query import Query as SqlaQuery from sqlalchemy.sql import exists @@ -2866,6 +2867,60 @@ class SupersetSecurityManager( # pylint: disable=too-many-public-methods .one_or_none() ) + def find_user_with_relationships( + self, + username: Optional[str] = None, + email: Optional[str] = None, + ) -> Optional[User]: + """Find a user with roles and group roles eagerly loaded. + + Mirrors FAB's ``SecurityManager.find_user`` + (including ``auth_username_ci`` case-insensitive handling and + ``MultipleResultsFound`` guard) and additionally eager-loads + ``User.roles`` and ``User.groups.roles`` to prevent detached-instance + errors when the SQLAlchemy session is closed or rolled back after the + lookup — as happens in MCP tool-execution contexts. + """ + eager = [ + joinedload(self.user_model.roles), + joinedload(self.user_model.groups).joinedload("roles"), + ] + if username: + try: + if self.auth_username_ci: + from sqlalchemy import func as sa_func + + return ( + self.session.query(self.user_model) + .options(*eager) + .filter( + sa_func.lower(self.user_model.username) + == sa_func.lower(username) + ) + .one_or_none() + ) + return ( + self.session.query(self.user_model) + .options(*eager) + .filter(self.user_model.username == username) + .one_or_none() + ) + except MultipleResultsFound: + logger.error("Multiple results found for user %s", username) + return None + if email: + try: + return ( + self.session.query(self.user_model) + .options(*eager) + .filter_by(email=email) + .one_or_none() + ) + except MultipleResultsFound: + logger.error("Multiple results found for user with email %s", email) + return None + return None + def get_anonymous_user(self) -> User: return AnonymousUserMixin() diff --git a/tests/unit_tests/mcp_service/test_auth_api_key.py b/tests/unit_tests/mcp_service/test_auth_api_key.py index 21d060441ff..961d41fbbc6 100644 --- a/tests/unit_tests/mcp_service/test_auth_api_key.py +++ b/tests/unit_tests/mcp_service/test_auth_api_key.py @@ -351,3 +351,17 @@ def test_security_manager_has_expected_api_key_methods(app: SupersetApp) -> None "auth._resolve_user_from_api_key() references this method by name — " "update auth.py if the FAB API changed." ) + + +def test_security_manager_has_find_user_with_relationships(app: SupersetApp) -> None: + """Regression test: verify SupersetSecurityManager.find_user_with_relationships + exists. load_user_with_relationships() in auth.py delegates to it — a rename + or removal would silently break MCP user resolution at runtime.""" + with app.app_context(): + from superset import security_manager + + assert hasattr(security_manager, "find_user_with_relationships"), ( + "SupersetSecurityManager is missing 'find_user_with_relationships'. " + "auth.load_user_with_relationships() delegates to this method — " + "update auth.py if the method was renamed or removed." + )