Compare commits

...
Author SHA1 Message Date
geido e8446b88a1 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.
2026-08-26 16:35:44 +00:00
2 changed files with 213 additions and 3 deletions
@@ -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
@@ -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