(superset) 13/30: refactor(mcp): delegate load_user_with_relationships to SecurityManager.find_user_with_relationships

[mi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

aminghadersohi pushed a commit to branch oss-39604
in repository https://gitbox.apache.org/repos/asf/superset.git

commit c4377b524f08e4ae4beb06273bac7bd3b2456e88
Author: Amin Ghadersohi <am...@gmail.com>
AuthorDate: Wed May 13 22:10:36 2026 +0000

    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 <no...@anthropic.com>
---
 superset/mcp_service/auth.py                      | 41 ++++------------
 superset/security/manager.py                      | 57 ++++++++++++++++++++++-
 tests/unit_tests/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 620763e5285..22e82e038fe 100644
--- a/superset/mcp_service/auth.py
+++ b/superset/mcp_service/auth.py
@@ -49,7 +49,7 @@ from contextlib import AbstractContextManager, nullcontext
 from typing import Any, Callable, TYPE_CHECKING, TypeVar
 
 from flask import current_app, g, has_app_context, 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
 
@@ -217,23 +217,14 @@ def is_tool_visible_to_current_user(tool: 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.
-
-    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.
+    """Load a user with roles and group roles eagerly loaded.
 
-    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
@@ -241,21 +232,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.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)
+    from superset import security_manager
 
-    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 5da6e9ced97..982c2c46e8a 100644
--- a/superset/security/manager.py
+++ b/superset/security/manager.py
@@ -52,7 +52,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
@@ -3168,6 +3169,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."
+        )