You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by ju...@apache.org on 2026/08/20 21:55:09 UTC
(superset) 02/02: fix(mcp): populate user_id in MCP audit logs (#42767)
This is an automated email from the ASF dual-hosted git repository.
justinpark pushed a commit to branch 6.2
in repository https://gitbox.apache.org/repos/asf/superset.git
commit 6188daa99878617b739085bedb341cc094c98857
Author: JUST.in DO IT <ju...@airbnb.com>
AuthorDate: Fri Aug 14 13:28:01 2026 -0700
fix(mcp): populate user_id in MCP audit logs (#42767)
(cherry picked from commit f4587218dd19d046c3e4d00063e7d27f8a2ed354)
---
superset/mcp_service/auth.py | 53 ++++++++---
superset/mcp_service/middleware.py | 79 +++++++++++-----
.../mcp_service/test_auth_user_resolution.py | 85 +++++++++++++++++
.../mcp_service/test_middleware_logging.py | 103 +++++++++++++++++++++
4 files changed, 284 insertions(+), 36 deletions(-)
diff --git a/superset/mcp_service/auth.py b/superset/mcp_service/auth.py
index 13cbc7e8e5e..4e55fa5eb8b 100644
--- a/superset/mcp_service/auth.py
+++ b/superset/mcp_service/auth.py
@@ -30,6 +30,7 @@ ensuring MCP tools respect the same permission model as the REST API.
import logging
from contextlib import AbstractContextManager
+from contextvars import ContextVar
from typing import Any, Callable, TYPE_CHECKING, TypeVar
from flask import g
@@ -44,6 +45,15 @@ F = TypeVar("F", bound=Callable[..., Any])
logger = logging.getLogger(__name__)
+# The resolved user's id for the current tool call, set by
+# _setup_user_context(). g.user is only valid for the lifetime of the
+# per-call app context pushed by _get_app_context_manager() (see its
+# docstring) and is gone by the time LoggingMiddleware's on_call_tool/
+# on_message finally-blocks run after call_next() returns. This ContextVar
+# survives that context pop (mirrors _mcp_call_id_var in middleware.py) so
+# audit logging can attribute the call to the right user.
+_mcp_user_id_var: ContextVar[int | None] = ContextVar("mcp_user_id", default=None)
+
# Constants for RBAC permission attributes (mirrors FAB conventions)
PERMISSION_PREFIX = "can_"
CLASS_PERMISSION_ATTR = "_class_permission_name"
@@ -280,9 +290,9 @@ def check_chart_data_access(chart: Any) -> "DatasetValidationResult":
return validate_chart_dataset(chart, check_access=True)
-def _setup_user_context() -> User | None:
+def _resolve_user_with_retry() -> User | None:
"""
- Set up user context for MCP tool execution.
+ Resolve the current user, retrying once on a stale DB connection.
Includes retry logic for stale database connections (e.g., SSL dropped
by proxy/load balancer after idle periods). On OperationalError, the
@@ -291,15 +301,7 @@ def _setup_user_context() -> User | None:
Returns:
User object with roles and groups loaded, or None if no Flask context
"""
- # Clear stale g.user to prevent user impersonation across
- # tool calls when no per-request middleware refreshes it.
- # Only clear in app-context-only mode; preserve g.user when
- # a request context is active (external middleware set it).
from flask import has_request_context
-
- if not has_request_context():
- g.pop("user", None)
-
from sqlalchemy.exc import OperationalError
user = None # Ensure defined before loop in case of unexpected exit
@@ -316,7 +318,7 @@ def _setup_user_context() -> User | None:
if hasattr(user, "groups"):
user_groups = user.groups # noqa: F841
- break
+ return user
except RuntimeError as e:
# No Flask application context (e.g., prompts before middleware runs)
if "application context" in str(e):
@@ -350,7 +352,36 @@ def _setup_user_context() -> User | None:
g.pop("user", None)
raise
+ return user
+
+
+def _setup_user_context() -> User | None:
+ """
+ Set up user context for MCP tool execution.
+
+ Returns:
+ User object with roles and groups loaded, or None if no Flask context
+ """
+ from flask import has_request_context
+
+ # Clear stale g.user to prevent user impersonation across
+ # tool calls when no per-request middleware refreshes it.
+ # Only clear in app-context-only mode; preserve g.user when
+ # a request context is active (external middleware set it).
+ if not has_request_context():
+ g.pop("user", None)
+ # Clear any user_id left over from a previous call in this context
+ # (e.g. sequential calls sharing one asyncio task) so a failed/
+ # unauthenticated lookup below doesn't inherit a stale value.
+ _mcp_user_id_var.set(None)
+
+ user = _resolve_user_with_retry()
+ if user is None:
+ return None
+
g.user = user
+ if (user_id := getattr(user, "id", None)) is not None:
+ _mcp_user_id_var.set(user_id)
return user
diff --git a/superset/mcp_service/middleware.py b/superset/mcp_service/middleware.py
index ab191a866f7..add71f20afe 100644
--- a/superset/mcp_service/middleware.py
+++ b/superset/mcp_service/middleware.py
@@ -31,6 +31,7 @@ from sqlalchemy.exc import OperationalError, TimeoutError
from starlette.exceptions import HTTPException
from superset.extensions import event_logger
+from superset.mcp_service.auth import _mcp_user_id_var
from superset.mcp_service.constants import (
DEFAULT_TOKEN_LIMIT,
DEFAULT_WARN_THRESHOLD_PCT,
@@ -403,6 +404,24 @@ class LoggingMiddleware(Middleware):
success = False
raise
finally:
+ # user_id was captured before call_next() ran the tool, i.e.
+ # before the @mcp_auth_hook decorator (superset/mcp_service/
+ # auth.py) resolves the user. It sets g.user on a per-call app
+ # context that mcp_auth_hook pushes and pops around the tool's
+ # execution, so g.user/get_user_id() are back to their pre-call
+ # state by the time we get here — re-reading get_user_id() would
+ # still yield the stale value. _mcp_user_id_var is a plain
+ # ContextVar (not tied to that Flask app-context lifecycle) that
+ # _setup_user_context() sets before the context pops, so it
+ # survives to this point.
+ resolved_user_id = _mcp_user_id_var.get(None)
+ if resolved_user_id is not None:
+ user_id = resolved_user_id
+ # Reset so a later on_call_tool/on_message in the same asyncio
+ # task (e.g. an unprotected tool or resource/prompt read that
+ # never calls _setup_user_context()) doesn't inherit this call's
+ # resolved user id.
+ _mcp_user_id_var.set(None)
self._log_call_tool_result(
context=context,
tool_name=tool_name,
@@ -426,32 +445,42 @@ class LoggingMiddleware(Middleware):
agent_id, user_id, dashboard_id, slice_id, dataset_id, params = (
self._extract_context_info(context)
)
- if has_app_context():
- event_logger.log(
- user_id=user_id,
- action="mcp_message",
- dashboard_id=dashboard_id,
- duration_ms=None,
- slice_id=slice_id,
- referrer=None,
- curated_payload={
- "tool": getattr(context.message, "name", None),
- "agent_id": agent_id,
- "params": _sanitize_params(params),
- "method": context.method,
- "dashboard_id": dashboard_id,
- "slice_id": slice_id,
- "dataset_id": dataset_id,
- },
+ try:
+ return await call_next(context)
+ finally:
+ # See the matching comment in on_call_tool: g.user/get_user_id()
+ # are stale here because the per-call app context has already
+ # been popped. _mcp_user_id_var survives it.
+ resolved_user_id = _mcp_user_id_var.get(None)
+ if resolved_user_id is not None:
+ user_id = resolved_user_id
+ # See the matching reset in on_call_tool.
+ _mcp_user_id_var.set(None)
+ if has_app_context():
+ event_logger.log(
+ user_id=user_id,
+ action="mcp_message",
+ dashboard_id=dashboard_id,
+ duration_ms=None,
+ slice_id=slice_id,
+ referrer=None,
+ curated_payload={
+ "tool": getattr(context.message, "name", None),
+ "agent_id": agent_id,
+ "params": _sanitize_params(params),
+ "method": context.method,
+ "dashboard_id": dashboard_id,
+ "slice_id": slice_id,
+ "dataset_id": dataset_id,
+ },
+ )
+ logger.info(
+ "MCP message: tool=%s, agent_id=%s, user_id=%s, method=%s",
+ getattr(context.message, "name", None),
+ agent_id,
+ user_id,
+ context.method,
)
- logger.info(
- "MCP message: tool=%s, agent_id=%s, user_id=%s, method=%s",
- getattr(context.message, "name", None),
- agent_id,
- user_id,
- context.method,
- )
- return await call_next(context)
class PrivateToolMiddleware(Middleware):
diff --git a/tests/unit_tests/mcp_service/test_auth_user_resolution.py b/tests/unit_tests/mcp_service/test_auth_user_resolution.py
new file mode 100644
index 00000000000..cc3cff8aff8
--- /dev/null
+++ b/tests/unit_tests/mcp_service/test_auth_user_resolution.py
@@ -0,0 +1,85 @@
+# 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.
+
+"""Tests for _mcp_user_id_var, the ContextVar that carries the resolved
+MCP user's id past the per-call app context pop (see auth.py docstring).
+"""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+
+def _make_mock_user(username: str = "testuser") -> MagicMock:
+ """Create a mock User with required attributes."""
+ user = MagicMock()
+ user.username = username
+ user.roles = []
+ user.groups = []
+ return user
+
+
+def test_setup_user_context_sets_contextvar_for_active_user(app) -> None:
+ """_mcp_user_id_var carries the resolved user's id past this call."""
+ from superset.mcp_service.auth import _mcp_user_id_var, _setup_user_context
+
+ active_user = _make_mock_user("active_user")
+ active_user.id = 321
+
+ with app.test_request_context():
+ with patch(
+ "superset.mcp_service.auth.get_user_from_request",
+ return_value=active_user,
+ ):
+ _setup_user_context()
+ assert _mcp_user_id_var.get() == 321
+
+
+def test_setup_user_context_clears_stale_contextvar_on_failure(app) -> None:
+ """A previous call's user_id must not leak into a call that fails to
+ resolve a user (e.g. sequential calls sharing one asyncio task)."""
+ from superset.mcp_service.auth import _mcp_user_id_var, _setup_user_context
+
+ with app.test_request_context():
+ _mcp_user_id_var.set(999)
+ with patch(
+ "superset.mcp_service.auth.get_user_from_request",
+ side_effect=ValueError("no user"),
+ ):
+ with pytest.raises(ValueError, match="no user"):
+ _setup_user_context()
+ assert _mcp_user_id_var.get() is None
+
+
+def test_setup_user_context_leaves_contextvar_unset_without_numeric_id(
+ app,
+) -> None:
+ """A resolved principal with no numeric id (e.g. an anonymous/guest
+ user) must leave the ContextVar cleared rather than store a bogus
+ value."""
+ from superset.mcp_service.auth import _mcp_user_id_var, _setup_user_context
+
+ no_id_user = _make_mock_user("no_id_user")
+ no_id_user.id = None
+
+ with app.test_request_context():
+ with patch(
+ "superset.mcp_service.auth.get_user_from_request",
+ return_value=no_id_user,
+ ):
+ _setup_user_context()
+ assert _mcp_user_id_var.get() is None
diff --git a/tests/unit_tests/mcp_service/test_middleware_logging.py b/tests/unit_tests/mcp_service/test_middleware_logging.py
index 46ae4a771b9..ec9651ff55a 100644
--- a/tests/unit_tests/mcp_service/test_middleware_logging.py
+++ b/tests/unit_tests/mcp_service/test_middleware_logging.py
@@ -31,6 +31,7 @@ import mcp.types as mt
import pytest
from fastmcp.tools.tool import ToolResult
+from superset.mcp_service.auth import _mcp_user_id_var
from superset.mcp_service.middleware import LoggingMiddleware
@@ -320,6 +321,108 @@ class TestLoggingMiddlewareOnMessage:
assert "success" not in call_kwargs["curated_payload"]
+class TestLoggingMiddlewareUserIdContextVar:
+ """Tests for _mcp_user_id_var overriding the pre-call user_id snapshot.
+
+ _extract_context_info() runs *before* call_next(), i.e. before the
+ tool's auth decorator (superset/mcp_service/auth.py) has resolved a
+ user, so get_user_id() there is always pre-auth (typically None).
+ _setup_user_context() sets _mcp_user_id_var while running inside
+ call_next(), right before the per-call app context (which g.user
+ depends on) is popped. These tests simulate that by setting the
+ ContextVar from within call_next(), mirroring the real auth decorator.
+ """
+
+ @patch("superset.mcp_service.middleware.event_logger")
+ @patch("superset.mcp_service.middleware.get_user_id", return_value=None)
+ @pytest.mark.asyncio
+ async def test_on_call_tool_uses_contextvar_set_during_call_next(
+ self, mock_get_user_id, mock_event_logger
+ ) -> None:
+ """Logged user_id reflects auth resolved during call_next(), not
+ the pre-call snapshot."""
+ middleware = LoggingMiddleware()
+ ctx = _make_context(name="list_charts")
+
+ async def fake_call_next(_ctx):
+ _mcp_user_id_var.set(42)
+ return "tool_result"
+
+ token = _mcp_user_id_var.set(None)
+ try:
+ result = await middleware.on_call_tool(ctx, fake_call_next)
+ finally:
+ _mcp_user_id_var.reset(token)
+
+ assert result == "tool_result"
+ call_kwargs = mock_event_logger.log.call_args[1]
+ assert call_kwargs["user_id"] == 42
+
+ @patch("superset.mcp_service.middleware.event_logger")
+ @patch("superset.mcp_service.middleware.get_user_id", return_value=7)
+ @pytest.mark.asyncio
+ async def test_on_call_tool_falls_back_when_contextvar_unset(
+ self, mock_get_user_id, mock_event_logger
+ ) -> None:
+ """When the call never runs _setup_user_context() (ContextVar
+ stays at its default), the pre-call get_user_id() value is used
+ as-is -- no regression for paths that don't set the ContextVar."""
+ middleware = LoggingMiddleware()
+ ctx = _make_context(name="health_check")
+ call_next = AsyncMock(return_value="ok")
+
+ token = _mcp_user_id_var.set(None)
+ try:
+ await middleware.on_call_tool(ctx, call_next)
+ finally:
+ _mcp_user_id_var.reset(token)
+
+ call_kwargs = mock_event_logger.log.call_args[1]
+ assert call_kwargs["user_id"] == 7
+
+ @patch("superset.mcp_service.middleware.event_logger")
+ @patch("superset.mcp_service.middleware.get_user_id", return_value=None)
+ @pytest.mark.asyncio
+ async def test_on_message_uses_contextvar_set_during_call_next(
+ self, mock_get_user_id, mock_event_logger
+ ) -> None:
+ middleware = LoggingMiddleware()
+ ctx = _make_context(method="resources/read", name="instance/metadata")
+
+ async def fake_call_next(_ctx):
+ _mcp_user_id_var.set(99)
+ return "resource_data"
+
+ token = _mcp_user_id_var.set(None)
+ try:
+ result = await middleware.on_message(ctx, fake_call_next)
+ finally:
+ _mcp_user_id_var.reset(token)
+
+ assert result == "resource_data"
+ call_kwargs = mock_event_logger.log.call_args[1]
+ assert call_kwargs["user_id"] == 99
+
+ @patch("superset.mcp_service.middleware.event_logger")
+ @patch("superset.mcp_service.middleware.get_user_id", return_value=13)
+ @pytest.mark.asyncio
+ async def test_on_message_falls_back_when_contextvar_unset(
+ self, mock_get_user_id, mock_event_logger
+ ) -> None:
+ middleware = LoggingMiddleware()
+ ctx = _make_context(method="resources/read", name="instance/metadata")
+ call_next = AsyncMock(return_value="resource_data")
+
+ token = _mcp_user_id_var.set(None)
+ try:
+ await middleware.on_message(ctx, call_next)
+ finally:
+ _mcp_user_id_var.reset(token)
+
+ call_kwargs = mock_event_logger.log.call_args[1]
+ assert call_kwargs["user_id"] == 13
+
+
class TestExtractContextInfo:
"""Tests for LoggingMiddleware._extract_context_info()."""