mirror of
https://github.com/apache/superset.git
synced 2026-09-09 08:44:32 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c3f16bfe14 | ||
|
|
aa1e940678 |
@@ -21,7 +21,7 @@ MCP tool: get_chart_data
|
||||
|
||||
import logging
|
||||
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
|
||||
@@ -29,13 +29,10 @@ 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
|
||||
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 +60,23 @@ 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
|
||||
datasource_id: Any
|
||||
datasource_type: str | None
|
||||
|
||||
|
||||
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:
|
||||
@@ -470,38 +484,68 @@ 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",
|
||||
# Load every column in one SELECT while the instance is still
|
||||
# attached, then copy the values this function needs into plain
|
||||
# locals.
|
||||
#
|
||||
# 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.
|
||||
db.session.refresh(chart)
|
||||
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(
|
||||
@@ -510,7 +554,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(
|
||||
@@ -590,7 +634,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),
|
||||
@@ -600,9 +644,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"
|
||||
)
|
||||
@@ -620,11 +664,11 @@ async def get_chart_data( # noqa: C901
|
||||
"Consider re-saving the chart to enable full data retrieval."
|
||||
)
|
||||
# Try to construct from form_data as a fallback
|
||||
form_data = utils_json.loads(chart.params) if chart.params else {}
|
||||
form_data = utils_json.loads(chart_params) if chart_params else {}
|
||||
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"),
|
||||
@@ -641,13 +685,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,
|
||||
@@ -665,17 +709,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 "
|
||||
@@ -686,8 +730,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,
|
||||
@@ -723,8 +767,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,
|
||||
)
|
||||
@@ -737,8 +781,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
|
||||
@@ -764,16 +808,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",
|
||||
)
|
||||
@@ -795,13 +839,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
|
||||
@@ -868,7 +912,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),
|
||||
)
|
||||
@@ -908,7 +952,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}",
|
||||
]
|
||||
@@ -926,7 +970,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,
|
||||
@@ -937,7 +981,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,
|
||||
@@ -957,7 +1001,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,
|
||||
@@ -968,9 +1012,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),
|
||||
@@ -995,12 +1039,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",
|
||||
@@ -1220,7 +1264,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,
|
||||
@@ -1279,7 +1323,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,
|
||||
@@ -1296,7 +1340,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
|
||||
@@ -1337,7 +1381,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,
|
||||
@@ -1364,7 +1408,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
|
||||
@@ -1405,7 +1449,7 @@ def _write_xlsxwriter_data(
|
||||
|
||||
|
||||
def _create_excel_chart_data(
|
||||
chart: "Slice",
|
||||
chart: "_ChartFacts",
|
||||
data: List[Dict[str, Any]],
|
||||
excel_b64: str,
|
||||
performance: Any,
|
||||
@@ -1438,7 +1482,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,
|
||||
|
||||
@@ -19,6 +19,7 @@
|
||||
Tests for the get_chart_data request schema and chart type fallback handling.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import importlib
|
||||
from contextlib import nullcontext
|
||||
from types import SimpleNamespace
|
||||
@@ -1331,6 +1332,22 @@ def mock_auth():
|
||||
yield mock_get_user
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def stub_session_refresh():
|
||||
"""No-op db.session.refresh() for tests using an unmapped chart double.
|
||||
|
||||
get_chart_data refreshes the chart right after lookup so its columns are
|
||||
loaded before the session can expire. SQLAlchemy rejects that call on a
|
||||
SimpleNamespace; in production the DAO always returns a persistent Slice.
|
||||
"""
|
||||
from unittest.mock import patch
|
||||
|
||||
from superset.extensions import db
|
||||
|
||||
with patch.object(db.session, "refresh", return_value=None):
|
||||
yield
|
||||
|
||||
|
||||
def _extract_metrics_load_path(load_opt: Any) -> list[str]:
|
||||
"""Walk a SQLAlchemy Load option and return the attr chain.
|
||||
|
||||
@@ -1565,7 +1582,7 @@ class TestSavedChartExtraFormDataFilters:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_filters_key_reaches_executed_query(
|
||||
self, mcp_server: Any, mock_auth: Any
|
||||
self, mcp_server: Any, mock_auth: Any, stub_session_refresh: Any
|
||||
) -> None:
|
||||
"""extra_form_data using the native 'filters' format is applied."""
|
||||
loaded, _ = await self._run(
|
||||
@@ -1576,7 +1593,7 @@ class TestSavedChartExtraFormDataFilters:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_adhoc_filters_key_reaches_executed_query(
|
||||
self, mcp_server: Any, mock_auth: Any
|
||||
self, mcp_server: Any, mock_auth: Any, stub_session_refresh: Any
|
||||
) -> None:
|
||||
"""extra_form_data using the 'adhoc_filters' format is also applied."""
|
||||
loaded, _ = await self._run(
|
||||
@@ -1598,7 +1615,7 @@ class TestSavedChartExtraFormDataFilters:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_temporal_range_filter_reaches_executed_query(
|
||||
self, mcp_server: Any, mock_auth: Any
|
||||
self, mcp_server: Any, mock_auth: Any, stub_session_refresh: Any
|
||||
) -> None:
|
||||
"""A TEMPORAL_RANGE filter narrows the query, not just simple filters."""
|
||||
loaded, _ = await self._run(
|
||||
@@ -1622,7 +1639,7 @@ class TestSavedChartExtraFormDataFilters:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unknown_adhoc_filter_column_returns_validation_error(
|
||||
self, mcp_server: Any, mock_auth: Any
|
||||
self, mcp_server: Any, mock_auth: Any, stub_session_refresh: Any
|
||||
) -> None:
|
||||
"""A rejected request filter must not return plausible unfiltered data."""
|
||||
_, result = await self._run(
|
||||
@@ -1736,6 +1753,7 @@ class TestOAuthErrorRouting:
|
||||
self,
|
||||
mcp_server: Any,
|
||||
mock_auth: Any,
|
||||
stub_session_refresh: Any,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from fastmcp import Client
|
||||
@@ -1767,6 +1785,7 @@ class TestOAuthErrorRouting:
|
||||
self,
|
||||
mcp_server: Any,
|
||||
mock_auth: Any,
|
||||
stub_session_refresh: Any,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from fastmcp import Client
|
||||
@@ -2207,7 +2226,7 @@ class TestGuestScoping:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_guest_query_pinned_to_dashboard_with_existence_only_check(
|
||||
self, mcp_server, mock_auth
|
||||
self, mcp_server, mock_auth, stub_session_refresh
|
||||
) -> None:
|
||||
from unittest.mock import patch
|
||||
|
||||
@@ -2591,3 +2610,249 @@ 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 _AccessRecordingChart:
|
||||
"""Slice stand-in that appends every attribute read to a shared event log.
|
||||
|
||||
Lets a test assert the *order* of session.refresh() relative to the first
|
||||
chart attribute access, which is the property that makes the fix work.
|
||||
"""
|
||||
|
||||
_COLUMNS = {
|
||||
"id": 42,
|
||||
"slice_name": "Sales Chart",
|
||||
"viz_type": "table",
|
||||
"datasource_id": 1,
|
||||
"datasource_type": "table",
|
||||
"params": "{}",
|
||||
"query_context": None,
|
||||
}
|
||||
|
||||
def __init__(self, events: list[str]) -> None:
|
||||
object.__setattr__(self, "_events", events)
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
object.__getattribute__(self, "_events").append(f"attr:{name}")
|
||||
try:
|
||||
return self._COLUMNS[name]
|
||||
except KeyError:
|
||||
raise AttributeError(name) from None
|
||||
|
||||
|
||||
class TestDetachedInstanceError:
|
||||
"""Regression tests for the DetachedInstanceError on the chart Slice.
|
||||
|
||||
get_chart_data is async and crosses many await points; the nested
|
||||
event_logger.log_context() blocks commit the DB session on exit, which
|
||||
expires every loaded ORM attribute. Later attribute access on the chart
|
||||
then raises DetachedInstanceError, which the broad SQLAlchemyError handler
|
||||
turns into a confusing internal-session error instead of chart data.
|
||||
|
||||
Calling db.session.refresh() immediately after the lookup loads all column
|
||||
values into the instance while the session is still live, so a later
|
||||
expiry has nothing left to lazily fetch. Mirrors the fix already applied
|
||||
to get_chart_preview.
|
||||
"""
|
||||
|
||||
def _patches(self, module: Any, chart: Any, refresh: Any) -> tuple[Any, ...]:
|
||||
from unittest.mock import patch
|
||||
|
||||
from superset.extensions import db
|
||||
|
||||
return (
|
||||
patch.object(module, "find_chart_by_identifier", return_value=chart),
|
||||
# Patched on the shared extension session rather than the tool
|
||||
# module, so a missing refresh surfaces as a failed assertion
|
||||
# rather than an AttributeError on the module's db import.
|
||||
patch.object(db.session, "refresh", side_effect=refresh),
|
||||
# Short-circuits the tool right after the lookup so the assertions
|
||||
# target the lookup block rather than the query pipeline.
|
||||
patch.object(
|
||||
module,
|
||||
"validate_chart_dataset",
|
||||
return_value=SimpleNamespace(
|
||||
is_valid=False, error="Dataset gone", warnings=[]
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_refresh_called_after_chart_load(self, mcp_server, mock_auth):
|
||||
"""db.session.refresh() runs once on the chart returned by the lookup."""
|
||||
from fastmcp import Client
|
||||
|
||||
module = importlib.import_module(
|
||||
"superset.mcp_service.chart.tool.get_chart_data"
|
||||
)
|
||||
chart = _AccessRecordingChart([])
|
||||
refreshed: list[object] = []
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
for patcher in self._patches(module, chart, refreshed.append):
|
||||
stack.enter_context(patcher)
|
||||
async with Client(mcp_server) as client:
|
||||
await client.call_tool(
|
||||
"get_chart_data", {"request": {"identifier": 42}}
|
||||
)
|
||||
|
||||
assert refreshed == [chart], (
|
||||
"db.session.refresh() must be called exactly once with the chart "
|
||||
f"loaded by find_chart_by_identifier; got {refreshed!r}"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_refresh_precedes_chart_attribute_access(
|
||||
self, mcp_server, mock_auth
|
||||
):
|
||||
"""The refresh must happen before anything reads a chart attribute.
|
||||
|
||||
A refresh placed after the first attribute read would leave that read
|
||||
exposed to an expired session, which is the bug being fixed.
|
||||
"""
|
||||
from fastmcp import Client
|
||||
|
||||
module = importlib.import_module(
|
||||
"superset.mcp_service.chart.tool.get_chart_data"
|
||||
)
|
||||
events: list[str] = []
|
||||
chart = _AccessRecordingChart(events)
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
for patcher in self._patches(
|
||||
module, chart, lambda _obj: events.append("refresh")
|
||||
):
|
||||
stack.enter_context(patcher)
|
||||
async with Client(mcp_server) as client:
|
||||
await client.call_tool(
|
||||
"get_chart_data", {"request": {"identifier": 42}}
|
||||
)
|
||||
|
||||
assert "refresh" in events, "db.session.refresh() was never called"
|
||||
assert events[0] == "refresh", (
|
||||
"db.session.refresh() must run before the first chart attribute "
|
||||
f"access; event order was {events[:5]}"
|
||||
)
|
||||
|
||||
|
||||
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, stub_session_refresh: 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"
|
||||
|
||||
Reference in New Issue
Block a user