Compare commits

...
Author SHA1 Message Date
Diego PucciandClaude Opus 5 47f0e5f28f fix(mcp): snapshot chart fields in get_chart_data before session expiry
get_chart_data reads the chart Slice throughout a long async function.
Exiting an event_logger.log_context() commits the session
(DBEventLogger.log -> db.session.commit), and with the default
expire_on_commit=True that commit expires every loaded attribute. If the
instance is also detached before the next read, that read raises
DetachedInstanceError, which the broad SQLAlchemyError handler turns into
a confusing internal-session error instead of chart data.

Copy the values the function needs into plain locals immediately after
the lookup, while the instance is still attached. Locals are immune to
both expiry and detachment. Chart reads here are interleaved with four
committing log_context blocks, so keeping the ORM object alive across
them is not workable; db.session.refresh() in particular does not help,
because the next commit expires exactly what it just loaded.

The export helpers and the two chart_helpers query builders also read
chart attributes long after those commits, so they are handed a
_ChartFacts NamedTuple instead of the ORM instance. Its field names match
the Slice columns, so helpers reading chart.viz_type or
getattr(chart, "datasource_id", None) behave identically. Audited: those
consumers touch only id, slice_name, viz_type, datasource_id and
datasource_type. guest_scope still receives the real Slice, since
query_context.slice_ needs the mapped object.

The subqueryload(Slice.table).subqueryload(SqlaTable.metrics) eager-load
from 39483 is untouched.

The not-found return moves inside the lookup block so the snapshot can be
unconditional; the log context still records the action on that path.

Regression test drives the tool end to end with a chart that detaches at
the end of the lookup block, for json, csv and excel. It fails against
master's get_chart_data on all three and passes here.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-09-09 17:41:55 +03:00
2 changed files with 236 additions and 71 deletions
+114 -71
View File
@@ -22,7 +22,7 @@ MCP tool: get_chart_data
import logging
import math
import time
from typing import Any, Dict, List, TYPE_CHECKING
from typing import Any, Dict, List, NamedTuple
from fastmcp import Context
from flask import current_app
@@ -30,9 +30,6 @@ from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.orm import subqueryload
from superset_core.mcp.decorators import tool, ToolAnnotations
if TYPE_CHECKING:
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
@@ -68,6 +65,24 @@ from superset.utils.core import GenericDataType
logger = logging.getLogger(__name__)
class _ChartFacts(NamedTuple):
"""Chart values copied off the Slice while it is still session-attached.
Downstream helpers run after several commits and await points, so they are
handed these plain values rather than the ORM instance, whose attributes
may be expired or detached by then. Field names match the Slice columns
they come from, so helpers reading ``chart.viz_type`` or
``getattr(chart, "datasource_id", None)`` behave identically.
"""
id: int
slice_name: str | None
viz_type: str | None
# Unset for an unsaved chart, which carries only export metadata.
datasource_id: Any = None
datasource_type: str | None = None
_GENERIC_TYPE_MAP: dict[int, str] = {
GenericDataType.NUMERIC: "numeric",
GenericDataType.STRING: "string",
@@ -416,38 +431,66 @@ async def get_chart_data( # noqa: C901
chart = find_chart_by_identifier(
request.identifier, query_options=chart_query_options
)
if chart is not None:
guest_dashboard_id = guest_scope.guest_dashboard_id(chart)
if not chart:
await ctx.warning(
"Chart not found: identifier=%s" % (request.identifier,)
)
logger.warning(
"get_chart_data: chart not found: identifier=%s", request.identifier
)
display_id = str(request.identifier)[:200]
return ChartError(
error=(
f"No chart found with identifier: {display_id}."
" Use list_charts to get valid chart IDs."
),
error_type="NotFound",
)
if not chart:
await ctx.warning("Chart not found: identifier=%s" % (request.identifier,))
logger.warning(
"get_chart_data: chart not found: identifier=%s", request.identifier
)
display_id = str(request.identifier)[:200]
return ChartError(
error=(
f"No chart found with identifier: {display_id}."
" Use list_charts to get valid chart IDs."
),
error_type="NotFound",
# Copy the values this function needs into plain locals while the
# instance is freshly loaded and still attached.
#
# Reading them off the ORM object later is not safe: exiting an
# event_logger.log_context() commits the session
# (DBEventLogger.log -> db.session.commit), and a commit expires
# every loaded attribute. If the instance is also detached before
# the next read -- this tool is async and crosses many await
# points -- that read raises DetachedInstanceError, which the
# broad SQLAlchemyError handler below turns into a confusing
# internal-session error instead of chart data. Plain locals are
# immune to both expiry and detachment.
chart_id = chart.id
chart_name = chart.slice_name
chart_viz_type = chart.viz_type
chart_datasource_id = chart.datasource_id
chart_datasource_type = chart.datasource_type
chart_params = chart.params
chart_query_context = chart.query_context
chart_facts = _ChartFacts(
chart_id,
chart_name,
chart_viz_type,
chart_datasource_id,
chart_datasource_type,
)
guest_dashboard_id = guest_scope.guest_dashboard_id(chart)
await ctx.info(
"Chart found successfully: chart_id=%s, chart_name=%s, viz_type=%s"
% (
chart.id,
chart.slice_name,
chart.viz_type,
chart_id,
chart_name,
chart_viz_type,
)
)
logger.info("Getting data for chart %s: %s", chart.id, chart.slice_name)
logger.info("Getting data for chart %s: %s", chart_id, chart_name)
# Guests skip the RBAC check (authorize_query covers it) but keep the
# existence check, so a deleted dataset still returns
# DatasetNotAccessible.
validation_result = validate_chart_dataset(
chart.datasource_id, check_access=not guest_scope.is_guest_read()
chart_datasource_id, check_access=not guest_scope.is_guest_read()
)
if not validation_result.is_valid:
await ctx.warning(
@@ -456,7 +499,7 @@ async def get_chart_data( # noqa: C901
)
logger.warning(
"get_chart_data: dataset not accessible for chart_id=%s: %s",
chart.id,
chart_id,
validation_result.error,
)
return ChartError(
@@ -527,7 +570,7 @@ async def get_chart_data( # noqa: C901
else:
try:
parsed_saved_form_data = (
utils_json.loads(chart.params) if chart.params else {}
utils_json.loads(chart_params) if chart_params else {}
)
form_data = (
parsed_saved_form_data
@@ -538,7 +581,7 @@ async def get_chart_data( # noqa: C901
form_data = {}
if not using_unsaved_state:
form_data["viz_type"] = chart.viz_type or form_data.get("viz_type")
form_data["viz_type"] = chart_viz_type or form_data.get("viz_type")
# If using cached form_data, we need to build query_context from it
if using_unsaved_state and cached_form_data_dict is not None:
@@ -553,7 +596,7 @@ async def get_chart_data( # noqa: C901
query_context = build_query_context_from_form_data(
cached_form_data_dict,
chart=chart,
chart=chart_facts,
extra_form_data=request.extra_form_data,
row_limit=row_limit,
order_desc=cached_form_data_dict.get("order_desc", True),
@@ -563,9 +606,9 @@ async def get_chart_data( # noqa: C901
await ctx.debug(
"Built query_context from cached form_data (unsaved state)"
)
elif chart.query_context:
elif chart_query_context:
try:
query_context_json = utils_json.loads(chart.query_context)
query_context_json = utils_json.loads(chart_query_context)
await ctx.debug(
"Using chart's saved query_context for data retrieval"
)
@@ -586,7 +629,7 @@ async def get_chart_data( # noqa: C901
from superset.common.query_context_factory import QueryContextFactory
factory = QueryContextFactory()
# row_limit from chart.params may be a str; coerce for
# row_limit from chart_params may be a str; coerce for
# apply_max_row_limit's int comparison.
row_limit = _coerce_row_limit(
request.limit or form_data.get("row_limit"),
@@ -603,13 +646,13 @@ async def get_chart_data( # noqa: C901
# Bubble charts use x/y/size as separate metric fields.
# Deck.gl charts (deck_arc, deck_scatter, etc.) use spatial
# column configs (lat/lon, geohash, etc.) instead.
viz_type = chart.viz_type or ""
viz_type = chart_viz_type or ""
fallback_queries = build_query_dicts_from_form_data(
form_data,
chart.datasource_id,
chart.datasource_type,
chart=chart,
chart_datasource_id,
chart_datasource_type,
chart=chart_facts,
extra_form_data=request.extra_form_data,
row_limit=row_limit,
order_desc=True,
@@ -627,17 +670,17 @@ async def get_chart_data( # noqa: C901
"(viz_type=%s): no metrics, columns, or groupby "
"could be extracted from form_data. "
"Re-save the chart to populate query_context."
% (chart.id, viz_type)
% (chart_id, viz_type)
)
logger.warning(
"get_chart_data: cannot construct fallback query for "
"chart_id=%s (viz_type=%s): no metrics/columns found",
chart.id,
chart_id,
viz_type,
)
return ChartError(
error=(
f"Chart {chart.id} (type: {viz_type}) has no "
f"Chart {chart_id} (type: {viz_type}) has no "
f"saved query_context and its form_data does "
f"not contain recognizable metrics or columns. "
f"Please open this chart in Superset and "
@@ -648,8 +691,8 @@ async def get_chart_data( # noqa: C901
query_context = factory.create(
datasource={
"id": chart.datasource_id,
"type": chart.datasource_type,
"id": chart_datasource_id,
"type": chart_datasource_type,
},
queries=fallback_queries,
form_data=form_data,
@@ -685,8 +728,8 @@ async def get_chart_data( # noqa: C901
"Query execution parameters: datasource_id=%s, datasource_type=%s, "
"row_limit=%s, force_refresh=%s"
% (
chart.datasource_id,
chart.datasource_type,
chart_datasource_id,
chart_datasource_type,
request.limit or 100,
request.force_refresh,
)
@@ -699,8 +742,8 @@ async def get_chart_data( # noqa: C901
set_query_context_form_data(
query_context,
chart.datasource_id,
chart.datasource_type,
chart_datasource_id,
chart_datasource_type,
)
# Execute the query
@@ -729,16 +772,16 @@ async def get_chart_data( # noqa: C901
if not result or ("queries" not in result) or len(result["queries"]) == 0:
await ctx.warning(
"Empty query results: chart_id=%s, chart_type=%s"
% (chart.id, chart.viz_type)
% (chart_id, chart_viz_type)
)
logger.warning(
"get_chart_data: empty query results for chart_id=%s, "
"chart_type=%s",
chart.id,
chart.viz_type,
chart_id,
chart_viz_type,
)
return ChartError(
error=f"No query results returned for chart {chart.id}. "
error=f"No query results returned for chart {chart_id}. "
f"This may occur with chart types like big_number.",
error_type="EmptyQuery",
)
@@ -760,13 +803,13 @@ async def get_chart_data( # noqa: C901
# Check if we have data to work with
if not any(query.get("data") for query in result["queries"]):
await ctx.warning("No data in query results: chart_id=%s" % (chart.id,))
await ctx.warning("No data in query results: chart_id=%s" % (chart_id,))
logger.warning(
"get_chart_data: no data in query results for chart_id=%s",
chart.id,
chart_id,
)
return ChartError(
error=f"No data available for chart {chart.id}", error_type="NoData"
error=f"No data available for chart {chart_id}", error_type="NoData"
)
# Create rich column metadata
@@ -833,7 +876,7 @@ async def get_chart_data( # noqa: C901
insights.append("Fresh data retrieved from database")
recommended_visualizations = _recommend_visualizations(
viz_type=chart.viz_type or "unknown",
viz_type=chart_viz_type or "unknown",
columns=columns,
row_count=len(data),
)
@@ -873,7 +916,7 @@ async def get_chart_data( # noqa: C901
cache_info = age_info
summary_parts = [
f"Chart '{chart.slice_name}' ({chart.viz_type})",
f"Chart '{chart_name}' ({chart_viz_type})",
f"Contains {len(data)} rows across {len(raw_columns)} columns"
f"{cache_info}",
]
@@ -891,7 +934,7 @@ async def get_chart_data( # noqa: C901
action="mcp.get_chart_data.format_conversion"
):
return _export_data_as_csv(
chart,
chart_facts,
data[: request.limit] if request.limit else data,
raw_columns,
cache_status,
@@ -902,7 +945,7 @@ async def get_chart_data( # noqa: C901
action="mcp.get_chart_data.format_conversion"
):
return _export_data_as_excel(
chart,
chart_facts,
data[: request.limit] if request.limit else data,
raw_columns,
cache_status,
@@ -922,7 +965,7 @@ async def get_chart_data( # noqa: C901
"rows_returned=%s, columns_returned=%s, execution_time_ms=%s, "
"cache_hit=%s, data_completeness=%s"
% (
chart.id,
chart_id,
len(data),
len(raw_columns),
execution_time,
@@ -933,9 +976,9 @@ async def get_chart_data( # noqa: C901
# Default JSON format
return ChartData(
chart_id=chart.id,
chart_name=chart.slice_name or f"Chart {chart.id}",
chart_type=chart.viz_type or "unknown",
chart_id=chart_id,
chart_name=chart_name or f"Chart {chart_id}",
chart_type=chart_viz_type or "unknown",
columns=columns,
data=data[: request.limit] if request.limit else data,
query_results=_build_query_results(result["queries"], request.limit),
@@ -960,12 +1003,12 @@ async def get_chart_data( # noqa: C901
await ctx.error(
"Data retrieval failed: chart_id=%s, error=%s, error_type=%s"
% (
chart.id,
chart_id,
str(data_error),
type(data_error).__name__,
)
)
logger.error("Data retrieval error for chart %s: %s", chart.id, data_error)
logger.error("Data retrieval error for chart %s: %s", chart_id, data_error)
return ChartError(
error=f"Error retrieving chart data: {str(data_error)}",
error_type="DataError",
@@ -1142,10 +1185,10 @@ async def _query_from_form_data( # noqa: C901
chart_name = form_data.get("slice_name", "Unsaved chart")
if request.format in {"csv", "excel"}:
from superset.models.slice import Slice
# A transient chart supplies export metadata without saving anything.
chart = Slice(id=0, slice_name=chart_name, viz_type=viz_type)
# Export metadata for an unsaved chart. A plain record rather than
# a transient Slice: nothing here is persisted, and an unattached
# mapped instance risks being picked up by a later autoflush.
chart = _ChartFacts(0, chart_name, viz_type)
export = (
_export_data_as_csv
if request.format == "csv"
@@ -1205,7 +1248,7 @@ async def _query_from_form_data( # noqa: C901
def _export_data_as_csv(
chart: "Slice",
chart: "_ChartFacts",
data: List[Dict[str, Any]],
columns: List[str],
cache_status: Any,
@@ -1264,7 +1307,7 @@ def _export_data_as_csv(
def _export_data_as_excel(
chart: "Slice",
chart: "_ChartFacts",
data: List[Dict[str, Any]],
columns: List[str],
cache_status: Any,
@@ -1281,7 +1324,7 @@ def _export_data_as_excel(
def _create_excel_with_openpyxl(
chart: "Slice", data: List[Dict[str, Any]], columns: List[str]
chart: "_ChartFacts", data: List[Dict[str, Any]], columns: List[str]
) -> str:
"""Create Excel file using openpyxl."""
import base64
@@ -1325,7 +1368,7 @@ def _write_excel_data(ws: Any, data: List[Dict[str, Any]], columns: List[str]) -
def _try_xlsxwriter_fallback(
chart: "Slice",
chart: "_ChartFacts",
data: List[Dict[str, Any]],
columns: List[str],
cache_status: Any,
@@ -1352,7 +1395,7 @@ def _try_xlsxwriter_fallback(
def _create_excel_with_xlsxwriter(
chart: "Slice", data: List[Dict[str, Any]], columns: List[str]
chart: "_ChartFacts", data: List[Dict[str, Any]], columns: List[str]
) -> str:
"""Create Excel file using xlsxwriter."""
import base64
@@ -1395,7 +1438,7 @@ def _write_xlsxwriter_data(
def _create_excel_chart_data(
chart: "Slice",
chart: "_ChartFacts",
data: List[Dict[str, Any]],
excel_b64: str,
performance: Any,
@@ -1428,7 +1471,7 @@ def _create_excel_chart_data(
def _create_excel_chart_data_xlsxwriter(
chart: "Slice",
chart: "_ChartFacts",
data: List[Dict[str, Any]],
excel_b64: str,
performance: Any,
@@ -2897,3 +2897,125 @@ def test_xlsxwriter_preserves_nonfinite_group_rows() -> None:
assert [row[0] for row in list(workbook.active.values)[1:]] == [
row["team"] for row in rows
]
class _DetachAfterLookupChart:
"""Slice stand-in that starts attached and detaches on demand.
After ``detach()`` every attribute read raises ``DetachedInstanceError``,
which is what a real Slice does once the session has committed (expiring
its attributes) and then been torn down.
"""
_COLUMNS = {
"id": 9,
"slice_name": "Sales",
"viz_type": "table",
"datasource_id": 1,
"datasource_type": "table",
"params": None,
"query_context": (
'{"datasource": {"id": 1, "type": "table"},'
' "queries": [{"columns": ["country"], "metrics": ["count"],'
' "filters": [], "row_limit": 100}],'
' "result_format": "json", "result_type": "full"}'
),
}
def __init__(self) -> None:
object.__setattr__(self, "_detached", False)
def detach(self) -> None:
object.__setattr__(self, "_detached", True)
def __getattr__(self, name: str) -> Any:
from sqlalchemy.orm.exc import DetachedInstanceError
if object.__getattribute__(self, "_detached"):
raise DetachedInstanceError(
"Instance <Slice at 0x0> is not bound to a Session; "
f"attribute refresh operation cannot proceed (attribute: {name})"
)
try:
return self._COLUMNS[name]
except KeyError:
raise AttributeError(name) from None
@pytest.mark.parametrize("export_format", ["json", "csv", "excel"])
@pytest.mark.asyncio
async def test_chart_data_survives_chart_detached_after_lookup(
export_format: str, mcp_server: Any, mock_auth: Any
) -> None:
"""The tool must still return data when the Slice detaches after lookup.
Reproduces the reported failure: the session commits and is torn down
partway through the request, so every later read on the chart instance
raises DetachedInstanceError and the broad SQLAlchemyError handler returns
an internal-session error instead of chart data. The chart is detached at
the end of the lookup block, right after its last legitimate ORM use.
"""
from unittest.mock import patch
from fastmcp import Client
module = importlib.import_module("superset.mcp_service.chart.tool.get_chart_data")
chart = _DetachAfterLookupChart()
def _detach_at_end_of_lookup(instance: Any) -> None:
instance.detach()
return None
def fake_load(self: Any, data: dict[str, Any]) -> Any:
queries = [
SimpleNamespace(
filter=query.get("filters", []),
time_range=query.get("time_range"),
to_dict=lambda query=query: dict(query),
)
for query in data.get("queries", [])
]
return SimpleNamespace(queries=queries, form_data=data.get("form_data", {}))
class _Command:
def __init__(self, query_context: Any) -> None: ...
def validate(self) -> None: ...
def run(self) -> dict[str, Any]:
return {
"queries": [
{
"data": [{"country": "USA"}],
"colnames": ["country"],
"rowcount": 1,
}
]
}
with (
patch.object(module, "find_chart_by_identifier", return_value=chart),
patch.object(
module,
"validate_chart_dataset",
return_value=SimpleNamespace(is_valid=True, warnings=[], error=None),
),
patch.object(
module.guest_scope, "guest_dashboard_id", _detach_at_end_of_lookup
),
patch(
"superset.commands.chart.data.get_data_command.ChartDataCommand", _Command
),
patch("superset.charts.schemas.ChartDataQueryContextSchema.load", fake_load),
):
async with Client(mcp_server) as client:
result = await client.call_tool(
"get_chart_data",
{"request": {"identifier": 9, "format": export_format}},
)
data = json.loads(result.content[0].text)
assert "error_type" not in data, (
f"format={export_format}: chart detached after lookup produced "
f"{data.get('error_type')}: {data.get('error')}"
)
assert data["chart_id"] == 9
assert data["chart_name"] == "Sales"