You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by di...@apache.org on 2026/08/26 16:35:55 UTC
(superset) 01/01: fix(mcp): keep get_chart_data's chart readable across mid-call commits
This is an automated email from the ASF dual-hosted git repository.
geido pushed a commit to branch agorithm/mcp-get-chart-data-session-refresh
in repository https://gitbox.apache.org/repos/asf/superset.git
commit e8446b88a1593b6a66b49eaa3c6f8779e8f74c94
Author: geido <di...@gmail.com>
AuthorDate: Wed Aug 26 16:35:44 2026 +0000
fix(mcp): keep get_chart_data's chart readable across mid-call commits
get_chart_data fetches a Slice once and then reads columns off it for the
rest of a long async call. Every event_logger.log_context block in the tool
commits the request session on exit (DBEventLogger.log), and SQLAlchemy
expires an instance's loaded columns on commit. The MCP per-call session can
also be removed while the call is still in flight (see
superset/mcp_service/session_scope.py), which detaches that Slice. Reading an
expired column off a detached instance raises DetachedInstanceError, which the
tool's broad `except (..., SQLAlchemyError, ...)` turns into
"Failed to get chart data: Instance <Slice ...> is not bound to a Session"
instead of the chart's data.
Hold expire_on_commit off for the duration of the call so the columns loaded
at fetch time stay in the instance and survive both the commit and the
detach, and restore the previous setting in a finally.
Note that db.session is a scoped_session proxy which does not forward
expire_on_commit assignment to the Session it wraps, so the flag is set on the
underlying Session; setting it on the proxy is silently ignored.
A refresh() call placed immediately after the lookup, as done in
get_chart_preview, does not help here: it runs inside the chart_lookup
log_context block, so that block's own exit commit expires everything
refresh() just loaded.
---
superset/mcp_service/chart/tool/get_chart_data.py | 34 +++-
.../mcp_service/chart/tool/test_get_chart_data.py | 182 ++++++++++++++++++++-
2 files changed, 213 insertions(+), 3 deletions(-)
diff --git a/superset/mcp_service/chart/tool/get_chart_data.py b/superset/mcp_service/chart/tool/get_chart_data.py
index a80250b7634..437adb0c831 100644
--- a/superset/mcp_service/chart/tool/get_chart_data.py
+++ b/superset/mcp_service/chart/tool/get_chart_data.py
@@ -26,16 +26,18 @@ from typing import Any, Dict, List, TYPE_CHECKING
from fastmcp import Context
from flask import current_app
from sqlalchemy.exc import SQLAlchemyError
-from sqlalchemy.orm import subqueryload
+from sqlalchemy.orm import scoped_session, subqueryload
from superset_core.mcp.decorators import tool, ToolAnnotations
if TYPE_CHECKING:
+ from sqlalchemy.orm import Session
+
from superset.models.slice import Slice
from superset.charts.data.form_data import set_query_context_form_data
from superset.commands.exceptions import CommandException
from superset.exceptions import OAuth2Error, OAuth2RedirectError, SupersetException
-from superset.extensions import event_logger
+from superset.extensions import db, event_logger
from superset.mcp_service import guest_scope
from superset.mcp_service.chart.chart_helpers import (
build_query_context_from_form_data,
@@ -63,6 +65,19 @@ from superset.utils.core import GenericDataType
logger = logging.getLogger(__name__)
+def _request_session() -> "Session":
+ """Return the ``Session`` backing ``db.session``.
+
+ ``db.session`` is a ``scoped_session`` proxy, and assigning a session
+ option such as ``expire_on_commit`` on the proxy does not reach the
+ ``Session`` it wraps -- the assignment lands on the proxy object and is
+ silently ignored. Callers that need to change session behavior therefore
+ have to resolve the underlying ``Session`` first.
+ """
+ session = db.session
+ return session() if isinstance(session, scoped_session) else session
+
+
def _requested_filter_columns(extra_form_data: dict[str, Any] | None) -> set[str]:
"""Return simple column names explicitly requested through extra form data."""
if not extra_form_data:
@@ -358,6 +373,19 @@ async def get_chart_data( # noqa: C901
)
effective_force = _compute_effective_force(request)
+ # The chart is fetched once below and then read from for the rest of this
+ # call. Every event_logger.log_context block here commits the request
+ # session on exit (DBEventLogger.log), and SQLAlchemy expires an
+ # instance's loaded columns on commit. The per-call session can also be
+ # removed while the call is still in flight (see
+ # superset/mcp_service/session_scope.py), which detaches that Slice --
+ # and reading an expired column off a detached instance raises
+ # DetachedInstanceError. Keeping the columns loaded across those commits
+ # means the already-fetched chart stays readable either way.
+ session = _request_session()
+ expire_on_commit = session.expire_on_commit
+ session.expire_on_commit = False
+
try:
await ctx.report_progress(1, 4, "Looking up chart")
from superset.utils import json as utils_json
@@ -1026,6 +1054,8 @@ async def get_chart_data( # noqa: C901
return ChartError(
error=f"Failed to get chart data: {str(e)}", error_type="InternalError"
)
+ finally:
+ session.expire_on_commit = expire_on_commit
async def _query_from_form_data( # noqa: C901
diff --git a/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py b/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py
index 16287494ff8..2fbb797c725 100644
--- a/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py
+++ b/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py
@@ -20,7 +20,7 @@ Tests for the get_chart_data request schema and chart type fallback handling.
"""
import importlib
-from contextlib import nullcontext
+from contextlib import contextmanager, nullcontext
from types import SimpleNamespace
from typing import Any
from unittest.mock import MagicMock
@@ -2532,3 +2532,183 @@ async def test_query_from_form_data_refreshed_reflects_force_refresh_only(
assert isinstance(result, ChartData)
assert result.cache_status is not None
assert result.cache_status.refreshed is expected_refreshed
+
+
+class _CommittingEventLogger:
+ """Stand-in for the default ``DBEventLogger`` plus MCP session teardown.
+
+ Two real behaviors are reproduced, in the order production hits them:
+
+ 1. ``DBEventLogger.log()`` ends every ``event_logger.log_context`` block
+ with ``db.session.commit()`` (``superset/utils/log.py``), and SQLAlchemy
+ expires every loaded attribute of the fetched ``Slice`` on commit.
+ 2. ``superset/mcp_service/session_scope.py`` documents that the request
+ session can be removed while a tool call is still in flight, which
+ detaches that ``Slice``. ``Session.close()`` reproduces exactly that:
+ it expunges every instance and drops the transaction.
+
+ Once both have happened, the next attribute read on the chart raises
+ ``DetachedInstanceError``.
+ """
+
+ def __init__(self, session: Any) -> None:
+ self._session = session
+ self.blocks = 0
+
+ @contextmanager
+ def log_context(self, *_args: Any, **_kwargs: Any) -> Any:
+ yield lambda **_kw: None
+ self._session.commit()
+ self._session.close()
+ self.blocks += 1
+
+
+def _make_persisted_chart(**overrides: Any) -> Any:
+ """Create a real ``Slice`` row in a throwaway in-memory session."""
+ from sqlalchemy import create_engine
+ from sqlalchemy.orm import sessionmaker
+
+ from superset.models.slice import Slice
+
+ engine = create_engine("sqlite://")
+ # Slice is versioned by sqlalchemy-continuum, whose before_flush hook
+ # writes to version_transaction, so the whole metadata has to exist.
+ Slice.metadata.create_all(engine)
+ session = sessionmaker(bind=engine)()
+
+ attrs: dict[str, Any] = {
+ "id": 42,
+ "slice_name": "Sales by region",
+ "viz_type": "table",
+ "datasource_id": 1,
+ "datasource_type": "table",
+ "params": json.dumps({"viz_type": "table"}),
+ "query_context": None,
+ }
+ attrs.update(overrides)
+ session.add(Slice(**attrs))
+ session.commit()
+ session.close()
+
+ return session.get(Slice, attrs["id"]), session
+
+
+@contextmanager
+def _detaching_get_chart_data(session: Any, validation: Any, chart: Any) -> Any:
+ """Patch get_chart_data so its session commits and is removed mid-call."""
+ from unittest.mock import Mock, patch
+
+ module = importlib.import_module("superset.mcp_service.chart.tool.get_chart_data")
+
+ user = Mock()
+ user.id = 1
+ user.username = "admin"
+
+ with (
+ patch("superset.mcp_service.auth.get_user_from_request", return_value=user),
+ patch.object(module, "event_logger", _CommittingEventLogger(session)),
+ patch.object(module, "find_chart_by_identifier", return_value=chart),
+ patch.object(module, "validate_chart_dataset", return_value=validation),
+ # ``db`` is only imported by the fixed module; ``raising=False``
+ # semantics keep this patch usable against the unfixed one too.
+ patch.object(module, "db", SimpleNamespace(session=session), create=True),
+ ):
+ yield module
+
+
+class TestDetachedInstanceError:
+ """``get_chart_data`` must survive its session committing then being removed.
+
+ The tool fetches a ``Slice`` once and then reads columns off it for the
+ rest of a long ``async`` call. Every ``event_logger.log_context`` exit
+ commits the request session, which expires those columns; if the per-call
+ session is then removed while the call is still running, the next read
+ raises ``DetachedInstanceError``. That is caught by the tool's broad
+ ``except (..., SQLAlchemyError, ...)`` handler, so the caller gets
+ "Failed to get chart data: Instance <Slice ...> is not bound to a Session"
+ instead of the chart's data.
+ """
+
+ @staticmethod
+ async def _call(identifier: Any = 42) -> dict[str, Any]:
+ from fastmcp import Client
+
+ from superset.mcp_service.app import mcp
+
+ async with Client(mcp) as client:
+ result = await client.call_tool(
+ "get_chart_data", {"request": {"identifier": identifier}}
+ )
+ return json.loads(result.content[0].text)
+
+ @pytest.mark.asyncio
+ async def test_first_chart_reads_survive_commit_and_removal(self) -> None:
+ """The reads immediately after the lookup must not hit a dead session."""
+ chart, session = _make_persisted_chart()
+ validation = SimpleNamespace(
+ is_valid=False, error="Dataset is gone", warnings=[]
+ )
+
+ with _detaching_get_chart_data(session, validation, chart):
+ data = await self._call()
+
+ assert "not bound to a Session" not in str(data), (
+ f"get_chart_data leaked a SQLAlchemy session error to the caller: {data}"
+ )
+ assert data["error_type"] == "DatasetNotAccessible"
+
+ @pytest.mark.asyncio
+ async def test_later_chart_reads_survive_commit_and_removal(self) -> None:
+ """Reads further down the tool (params, query_context, viz_type) too."""
+ chart, session = _make_persisted_chart(params=json.dumps({}))
+ validation = SimpleNamespace(is_valid=True, error=None, warnings=[])
+
+ with _detaching_get_chart_data(session, validation, chart):
+ data = await self._call()
+
+ assert "not bound to a Session" not in str(data), (
+ f"get_chart_data leaked a SQLAlchemy session error to the caller: {data}"
+ )
+ # Reaching this branch means chart.query_context, chart.params,
+ # chart.viz_type, chart.datasource_id/type and chart.id were all read
+ # successfully after the session had committed and been removed.
+ assert data["error_type"] == "MissingQueryContext"
+
+ @pytest.mark.asyncio
+ async def test_scoped_session_proxy_is_unwrapped(self) -> None:
+ """The guard must reach the real Session, not a scoped_session proxy.
+
+ ``db.session`` is a ``scoped_session``, which does not forward
+ ``expire_on_commit`` assignment to the Session it wraps -- setting the
+ flag on the proxy is a silent no-op. Guarding the proxy instead of the
+ underlying Session would leave the bug fully intact in production.
+ """
+ from sqlalchemy.orm import scoped_session
+
+ chart, session = _make_persisted_chart()
+ proxy = scoped_session(lambda: session)
+ validation = SimpleNamespace(
+ is_valid=False, error="Dataset is gone", warnings=[]
+ )
+
+ with _detaching_get_chart_data(proxy, validation, chart):
+ data = await self._call()
+
+ assert "not bound to a Session" not in str(data), (
+ f"Chart columns expired through the scoped_session proxy: {data}"
+ )
+ assert data["error_type"] == "DatasetNotAccessible"
+
+ @pytest.mark.asyncio
+ async def test_expire_on_commit_is_restored(self) -> None:
+ """The guard must not leak its relaxed setting past the tool call."""
+ chart, session = _make_persisted_chart()
+ validation = SimpleNamespace(
+ is_valid=False, error="Dataset is gone", warnings=[]
+ )
+ assert session.expire_on_commit is True
+
+ with _detaching_get_chart_data(session, validation, chart):
+ await self._call()
+
+ assert session.expire_on_commit is True