eschutho commented on code in PR #36013:
URL: https://github.com/apache/superset/pull/36013#discussion_r2499569737


##########
superset/mcp_service/auth.py:
##########
@@ -119,61 +154,92 @@ def mcp_auth_hook(tool_func: F) -> F:
     """
     Authentication and authorization decorator for MCP tools.
 
-    This is a minimal implementation that:
-    1. Gets the current user
-    2. Sets g.user for Flask context
-
-    TODO (future PR): Add permission checking
-    TODO (future PR): Add JWT scope validation
-    TODO (future PR): Add comprehensive audit logging
-    TODO (future PR): Add rate limiting integration
+    This hook ensures proper Flask app context and user authentication for all
+    MCP tool calls. Deployers can extend this decorator to add:
+
+    - Permission checks: Validate user has specific roles or tool access rights
+    - JWT scope validation: Check OAuth scopes match required tool permissions
+    - Audit logging: Track tool usage, parameters, and results for compliance
+    - Rate limiting: Enforce per-user or per-tool rate limits
+    - Custom middleware: Inject organization-specific security policies
+
+    Default behavior:
+    1. Ensures Flask app context is active (critical for stateless_http mode)
+    2. Authenticates user via get_user_from_request()
+    3. Sets g.user for Superset's security manager integration
+    4. Manages database session lifecycle
+
+    Example extensions:
+        # Add permission check before tool execution:
+        if tool_func.__name__ in ['execute_sql'] and not user.is_admin():
+            raise PermissionError("SQL execution requires admin role")
+
+        # Add audit logging after successful execution:
+        audit_log.info({
+            'user': user.username,
+            'tool': tool_func.__name__,
+            'timestamp': datetime.now(),
+            'result': 'success'
+        })
     """
     import functools
 
     @functools.wraps(tool_func)
     def wrapper(*args: Any, **kwargs: Any) -> Any:
         from superset.extensions import db
+        from superset.mcp_service.flask_singleton import get_flask_app
 
-        # Get user and set Flask context OUTSIDE try block
-        user = get_user_from_request()
+        # Get Flask app instance
+        app = get_flask_app()
 
-        # Force load relationships NOW while session is definitely active
-        _ = user.roles
-        if hasattr(user, "groups"):
-            _ = user.groups
+        # Ensure Flask app context is active
+        # This is critical for stateless_http mode where context isn't 
preserved
+        with app.app_context():
+            # Get user and set Flask context OUTSIDE try block
+            user = get_user_from_request()
 
-        g.user = user
+            # Force load relationships NOW while session is definitely active
+            _ = user.roles
+            if hasattr(user, "groups"):
+                _ = user.groups
 
-        try:
-            # TODO: Add permission checks here in future PR
-            # TODO: Add audit logging here in future PR
+            g.user = user
 
-            logger.debug(
-                "MCP tool call: user=%s, tool=%s", user.username, 
tool_func.__name__
-            )
+            try:
+                # Deployers: Add custom permission checks here
+                # Example: Check tool-specific permissions, validate JWT 
scopes,
+                # or enforce organizational policies
 
-            result = tool_func(*args, **kwargs)
+                logger.debug(
+                    "MCP tool call: user=%s, tool=%s", user.username, 
tool_func.__name__
+                )
 
-            return result
+                result = tool_func(*args, **kwargs)
 
-        except Exception:
-            # On error, rollback and cleanup session
-            # pylint: disable=consider-using-transaction
-            try:
-                db.session.rollback()
-                db.session.remove()
-            except Exception as e:
-                logger.warning("Error cleaning up session after exception: 
%s", e)
-            raise
-
-        finally:
-            # Only rollback if session is still active (no exception occurred)
-            # Do NOT call remove() on success to avoid detaching user
-            try:
-                if db.session.is_active:
-                    # pylint: disable=consider-using-transaction
+                # Deployers: Add audit logging or post-execution hooks here
+                # Example: Log successful tool execution, track metrics,
+                # or trigger downstream workflows
+
+                return result
+
+            except Exception:
+                # On error, rollback and cleanup session
+                # pylint: disable=consider-using-transaction
+                try:
                     db.session.rollback()
-            except Exception as e:
-                logger.warning("Error in finally block: %s", e)
+                    db.session.remove()
+                except Exception as e:
+                    logger.warning("Error cleaning up session after exception: 
%s", e)

Review Comment:
   Do we need both the warning and the raised error in the logs?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to