fix(mcp): prevent GRID_ID injection into ROOT_ID on tabbed dashboards (#38890)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Amin Ghadersohi
2026-03-30 11:26:58 +02:00
committed by GitHub
parent d1903afc69
commit aa1a69555b
2 changed files with 265 additions and 177 deletions

View File

@@ -26,8 +26,10 @@ import re
from typing import Any, Dict
from fastmcp import Context
from sqlalchemy.exc import SQLAlchemyError
from superset_core.mcp.decorators import tool, ToolAnnotations
from superset.commands.exceptions import CommandException
from superset.extensions import event_logger
from superset.mcp_service.chart.schemas import serialize_chart_object
from superset.mcp_service.dashboard.constants import (
@@ -289,8 +291,17 @@ def _ensure_layout_structure(
if "ROOT_ID" in layout:
if "children" not in layout["ROOT_ID"]:
layout["ROOT_ID"]["children"] = []
if "GRID_ID" not in layout["ROOT_ID"]["children"]:
layout["ROOT_ID"]["children"].append("GRID_ID")
# Only add GRID_ID to ROOT_ID when TABS are not already a direct
# child of ROOT_ID. Real Superset dashboards with tabs place a
# TABS container directly under ROOT_ID (ROOT_ID → TABS → TABs).
# Adding GRID_ID as a sibling of TABS confuses the frontend layout
# engine and makes charts invisible.
root_children = layout["ROOT_ID"]["children"]
has_tabs_under_root = any(
layout.get(c, {}).get("type") == "TABS" for c in root_children
)
if not has_tabs_under_root and "GRID_ID" not in root_children:
root_children.append("GRID_ID")
else:
# Create ROOT_ID if it doesn't exist
layout["ROOT_ID"] = {
@@ -320,10 +331,6 @@ def add_chart_to_existing_dashboard(
Add chart to existing dashboard. Auto-positions in 2-column grid.
Returns updated dashboard info.
"""
from sqlalchemy.exc import SQLAlchemyError
from superset.commands.exceptions import CommandException
try:
from superset.commands.dashboard.update import UpdateDashboardCommand
from superset.daos.dashboard import DashboardDAO
@@ -428,54 +435,25 @@ def add_chart_to_existing_dashboard(
# Re-fetch the dashboard with eager-loaded relationships to avoid
# "Instance is not bound to a Session" errors when serializing
# chart .tags and .owners. The preceding command.run() commit may
# invalidate the session in multi-tenant environments; on failure,
# return a minimal response using only scalar attributes that are
# already loaded — relationship fields (owners, tags, slices) would
# trigger lazy-loading on the same dead session.
# chart .tags and .owners.
from sqlalchemy.orm import subqueryload
from superset.daos.dashboard import DashboardDAO
from superset.models.dashboard import Dashboard
from superset.models.slice import Slice
try:
updated_dashboard = (
DashboardDAO.find_by_id(
updated_dashboard.id,
query_options=[
subqueryload(Dashboard.slices).subqueryload(Slice.owners),
subqueryload(Dashboard.slices).subqueryload(Slice.tags),
subqueryload(Dashboard.owners),
subqueryload(Dashboard.tags),
],
)
or updated_dashboard
)
except SQLAlchemyError:
logger.warning(
"Re-fetch of dashboard %s failed; returning minimal response",
updated_dashboard = (
DashboardDAO.find_by_id(
updated_dashboard.id,
exc_info=True,
)
db.session.rollback()
dashboard_url = (
f"{get_superset_base_url()}/superset/dashboard/{updated_dashboard.id}/"
)
position_info = {
"row": row_key,
"chart_key": chart_key,
"row_key": row_key,
}
return AddChartToDashboardResponse(
dashboard=DashboardInfo(
id=updated_dashboard.id,
dashboard_title=updated_dashboard.dashboard_title,
url=dashboard_url,
),
dashboard_url=dashboard_url,
position=position_info,
error=None,
query_options=[
subqueryload(Dashboard.slices).subqueryload(Slice.owners),
subqueryload(Dashboard.slices).subqueryload(Slice.tags),
subqueryload(Dashboard.owners),
subqueryload(Dashboard.tags),
],
)
or updated_dashboard
)
# Convert to response format
from superset.mcp_service.dashboard.schemas import (

View File

@@ -20,7 +20,7 @@ Unit tests for dashboard generation MCP tools
"""
import logging
from unittest.mock import MagicMock, Mock, patch
from unittest.mock import Mock, patch
import pytest
from fastmcp import Client
@@ -133,8 +133,7 @@ def _setup_generate_dashboard_mocks(
The tool creates dashboards directly via db.session (bypassing
CreateDashboardCommand) and re-queries user/charts in the tool's
own session. The re-fetch uses DashboardDAO.find_by_id() with
query_options for eager loading of slice relationships.
own session. This helper wires up the mock chain for that path.
"""
mock_user = Mock()
mock_user.id = 1
@@ -144,8 +143,8 @@ def _setup_generate_dashboard_mocks(
mock_user.email = "admin@example.com"
mock_user.active = True
mock_query = MagicMock()
mock_filter = MagicMock()
mock_query = Mock()
mock_filter = Mock()
mock_query.filter.return_value = mock_filter
mock_query.filter_by.return_value = mock_filter
mock_filter.order_by.return_value = mock_filter
@@ -154,7 +153,6 @@ def _setup_generate_dashboard_mocks(
mock_db_session.query.return_value = mock_query
mock_dashboard_cls.return_value = dashboard
# DashboardDAO.find_by_id is used for the re-fetch with eager loading
mock_find_by_id.return_value = dashboard
@@ -558,15 +556,15 @@ class TestAddChartToExistingDashboard:
_mock_chart(id=20),
_mock_chart(id=30),
]
# First call: initial validation returns original dashboard
# Second call: re-fetch after update returns updated dashboard
mock_find_dashboard.side_effect = [mock_dashboard, updated_dashboard]
mock_chart = _mock_chart(id=30, slice_name="New Chart")
mock_db_session.get.return_value = mock_chart
mock_update_command.return_value.run.return_value = updated_dashboard
# First DAO call returns initial dashboard (validation),
# second DAO call returns updated dashboard (re-fetch with eager loading)
mock_find_dashboard.side_effect = [mock_dashboard, updated_dashboard]
request = {"dashboard_id": 1, "chart_id": 30}
async with Client(mcp_server) as client:
@@ -702,6 +700,8 @@ class TestAddChartToExistingDashboard:
mock_dashboard = _mock_dashboard(id=2)
mock_dashboard.slices = []
mock_dashboard.position_json = "{}"
mock_find_dashboard.return_value = mock_dashboard
mock_chart = _mock_chart(id=15)
mock_db_session.get.return_value = mock_chart
@@ -709,10 +709,6 @@ class TestAddChartToExistingDashboard:
updated_dashboard.slices = [_mock_chart(id=15)]
mock_update_command.return_value.run.return_value = updated_dashboard
# First DAO call returns initial dashboard (validation),
# second returns updated dashboard (re-fetch)
mock_find_dashboard.side_effect = [mock_dashboard, updated_dashboard]
request = {"dashboard_id": 2, "chart_id": 15}
async with Client(mcp_server) as client:
@@ -813,6 +809,8 @@ class TestAddChartToExistingDashboard:
"DASHBOARD_VERSION_KEY": "v2",
}
)
mock_find_dashboard.return_value = mock_dashboard
mock_chart = _mock_chart(id=25, slice_name="Tab Chart")
mock_db_session.get.return_value = mock_chart
@@ -820,10 +818,6 @@ class TestAddChartToExistingDashboard:
updated_dashboard.slices = [_mock_chart(id=10), _mock_chart(id=25)]
mock_update_command.return_value.run.return_value = updated_dashboard
# First DAO call returns initial dashboard (validation),
# second returns updated dashboard (re-fetch)
mock_find_dashboard.side_effect = [mock_dashboard, updated_dashboard]
request = {"dashboard_id": 3, "chart_id": 25}
async with Client(mcp_server) as client:
@@ -915,6 +909,8 @@ class TestAddChartToExistingDashboard:
"DASHBOARD_VERSION_KEY": "v2",
}
)
mock_find_dashboard.return_value = mock_dashboard
mock_chart = _mock_chart(id=30, slice_name="Customer Chart")
mock_db_session.get.return_value = mock_chart
@@ -922,10 +918,6 @@ class TestAddChartToExistingDashboard:
updated_dashboard.slices = [_mock_chart(id=10), _mock_chart(id=30)]
mock_update_command.return_value.run.return_value = updated_dashboard
# First DAO call returns initial dashboard (validation),
# second returns updated dashboard (re-fetch)
mock_find_dashboard.side_effect = [mock_dashboard, updated_dashboard]
request = {"dashboard_id": 3, "chart_id": 30, "target_tab": "Customers"}
async with Client(mcp_server) as client:
@@ -954,6 +946,118 @@ class TestAddChartToExistingDashboard:
assert "TAB-tab2" in chart_parents
assert "TAB-tab1" not in chart_parents
@patch("superset.commands.dashboard.update.UpdateDashboardCommand")
@patch("superset.daos.dashboard.DashboardDAO.find_by_id")
@patch("superset.db.session")
@pytest.mark.asyncio
async def test_add_chart_to_tabbed_dashboard_tabs_under_root(
self, mock_db_session, mock_find_dashboard, mock_update_command, mcp_server
):
"""Test adding chart when TABS are under ROOT_ID (real-world layout).
Real Superset dashboards place TABS directly under ROOT_ID with an
empty GRID_ID, unlike test fixtures that place TABS under GRID_ID.
The tool must NOT inject GRID_ID into ROOT_ID.children alongside
TABS, as the frontend hides non-TABS content when a TABS container
is a ROOT_ID child.
"""
mock_dashboard = _mock_dashboard(id=7, title="COVID Vaccine Dashboard")
mock_dashboard.slices = [_mock_chart(id=10)]
mock_dashboard.position_json = json.dumps(
{
"ROOT_ID": {
"children": ["TABS-wUKya7eQ0Z"],
"id": "ROOT_ID",
"type": "ROOT",
},
"GRID_ID": {
"children": [],
"id": "GRID_ID",
"parents": ["ROOT_ID"],
"type": "GRID",
},
"TABS-wUKya7eQ0Z": {
"children": ["TAB-BCIJF4NvgQ", "TAB-kl2Hkh2IR"],
"id": "TABS-wUKya7eQ0Z",
"parents": ["ROOT_ID"],
"type": "TABS",
},
"TAB-BCIJF4NvgQ": {
"children": ["ROW-existing"],
"id": "TAB-BCIJF4NvgQ",
"meta": {"text": "Vaccine Candidates"},
"parents": ["ROOT_ID", "TABS-wUKya7eQ0Z"],
"type": "TAB",
},
"TAB-kl2Hkh2IR": {
"children": [],
"id": "TAB-kl2Hkh2IR",
"meta": {"text": "Doses Administered"},
"parents": ["ROOT_ID", "TABS-wUKya7eQ0Z"],
"type": "TAB",
},
"ROW-existing": {
"children": ["CHART-10"],
"id": "ROW-existing",
"meta": {"background": "BACKGROUND_TRANSPARENT"},
"parents": [
"ROOT_ID",
"TABS-wUKya7eQ0Z",
"TAB-BCIJF4NvgQ",
],
"type": "ROW",
},
"CHART-10": {
"id": "CHART-10",
"type": "CHART",
"parents": [
"ROOT_ID",
"TABS-wUKya7eQ0Z",
"TAB-BCIJF4NvgQ",
"ROW-existing",
],
},
"DASHBOARD_VERSION_KEY": "v2",
}
)
mock_chart = _mock_chart(id=91, slice_name="Vaccines by Stage")
mock_db_session.get.return_value = mock_chart
updated_dashboard = _mock_dashboard(id=7, title="COVID Vaccine Dashboard")
updated_dashboard.slices = [_mock_chart(id=10), _mock_chart(id=91)]
mock_update_command.return_value.run.return_value = updated_dashboard
mock_find_dashboard.side_effect = [mock_dashboard, updated_dashboard]
request = {"dashboard_id": 7, "chart_id": 91}
async with Client(mcp_server) as client:
result = await client.call_tool(
"add_chart_to_existing_dashboard", {"request": request}
)
assert result.structured_content["error"] is None
call_args = mock_update_command.call_args[0][1]
layout = json.loads(call_args["position_json"])
row_key = result.structured_content["position"]["row_key"]
assert row_key in layout
# Chart must be inside the first tab, not GRID_ID
assert row_key in layout["TAB-BCIJF4NvgQ"]["children"]
assert row_key not in layout["GRID_ID"]["children"]
# GRID_ID must NOT be added to ROOT_ID.children alongside TABS
assert "GRID_ID" not in layout["ROOT_ID"]["children"]
assert layout["ROOT_ID"]["children"] == ["TABS-wUKya7eQ0Z"]
# Parent chain must include the tab hierarchy, not GRID_ID
chart_parents = layout["CHART-91"]["parents"]
assert "TABS-wUKya7eQ0Z" in chart_parents
assert "TAB-BCIJF4NvgQ" in chart_parents
assert "GRID_ID" not in chart_parents
@patch("superset.commands.dashboard.update.UpdateDashboardCommand")
@patch("superset.daos.dashboard.DashboardDAO.find_by_id")
@patch("superset.db.session")
@@ -992,6 +1096,8 @@ class TestAddChartToExistingDashboard:
"DASHBOARD_VERSION_KEY": "v2",
}
)
mock_find_dashboard.return_value = mock_dashboard
mock_chart = _mock_chart(id=50, slice_name="New Nanoid Chart")
mock_db_session.get.return_value = mock_chart
@@ -999,10 +1105,6 @@ class TestAddChartToExistingDashboard:
updated_dashboard.slices = [_mock_chart(id=10), _mock_chart(id=50)]
mock_update_command.return_value.run.return_value = updated_dashboard
# First DAO call returns initial dashboard (validation),
# second returns updated dashboard (re-fetch)
mock_find_dashboard.side_effect = [mock_dashboard, updated_dashboard]
request = {"dashboard_id": 4, "chart_id": 50}
async with Client(mcp_server) as client:
@@ -1117,6 +1219,28 @@ class TestLayoutHelpers:
_find_tab_insert_target(layout, target_tab="Nonexistent Tab") == "TAB-first"
)
def test_find_tab_insert_target_tabs_under_root(self):
"""Test _find_tab_insert_target when TABS are under ROOT_ID (real layout)."""
layout = {
"ROOT_ID": {"children": ["TABS-xxx"], "type": "ROOT"},
"GRID_ID": {"children": [], "type": "GRID", "parents": ["ROOT_ID"]},
"TABS-xxx": {"children": ["TAB-a", "TAB-b"], "type": "TABS"},
"TAB-a": {"children": [], "type": "TAB", "meta": {"text": "Overview"}},
"TAB-b": {"children": [], "type": "TAB", "meta": {"text": "Details"}},
}
assert _find_tab_insert_target(layout) == "TAB-a"
def test_find_tab_insert_target_tabs_under_root_by_name(self):
"""Test _find_tab_insert_target matches tab name when TABS under ROOT_ID."""
layout = {
"ROOT_ID": {"children": ["TABS-xxx"], "type": "ROOT"},
"GRID_ID": {"children": [], "type": "GRID", "parents": ["ROOT_ID"]},
"TABS-xxx": {"children": ["TAB-a", "TAB-b"], "type": "TABS"},
"TAB-a": {"children": [], "type": "TAB", "meta": {"text": "Overview"}},
"TAB-b": {"children": [], "type": "TAB", "meta": {"text": "Details"}},
}
assert _find_tab_insert_target(layout, target_tab="Details") == "TAB-b"
def test_find_tab_insert_target_no_grid(self):
"""Test _find_tab_insert_target with missing GRID_ID."""
assert _find_tab_insert_target({"ROOT_ID": {"type": "ROOT"}}) is None
@@ -1174,6 +1298,56 @@ class TestLayoutHelpers:
assert "ROW-new" in layout["TAB-first"]["children"]
assert "ROW-new" not in layout["GRID_ID"]["children"]
def test_ensure_layout_structure_tabs_under_root_no_grid_added(self):
"""Test _ensure_layout_structure does NOT add GRID_ID to ROOT_ID
when TABS already exists as a ROOT_ID child.
Real Superset tabbed dashboards place TABS under ROOT_ID, not
GRID_ID. Adding GRID_ID as a sibling of TABS confuses the
frontend and makes charts invisible.
"""
layout = {
"ROOT_ID": {"children": ["TABS-xxx"], "type": "ROOT"},
"GRID_ID": {"children": [], "type": "GRID", "parents": ["ROOT_ID"]},
"TABS-xxx": {
"children": ["TAB-a", "TAB-b"],
"type": "TABS",
"parents": ["ROOT_ID"],
},
"TAB-a": {
"children": ["ROW-existing"],
"type": "TAB",
"meta": {"text": "Overview"},
"parents": ["ROOT_ID", "TABS-xxx"],
},
"TAB-b": {
"children": [],
"type": "TAB",
"meta": {"text": "Details"},
"parents": ["ROOT_ID", "TABS-xxx"],
},
}
_ensure_layout_structure(layout, "ROW-new", "TAB-a")
# Row added to the correct tab
assert "ROW-new" in layout["TAB-a"]["children"]
# GRID_ID must NOT be injected into ROOT_ID alongside TABS
assert "GRID_ID" not in layout["ROOT_ID"]["children"]
assert layout["ROOT_ID"]["children"] == ["TABS-xxx"]
def test_ensure_layout_structure_no_tabs_adds_grid_to_root(self):
"""Test _ensure_layout_structure still adds GRID_ID to ROOT_ID
when the dashboard has no tabs (non-tabbed dashboard regression check).
"""
layout = {
"ROOT_ID": {"children": [], "type": "ROOT"},
"GRID_ID": {"children": [], "type": "GRID", "parents": ["ROOT_ID"]},
}
_ensure_layout_structure(layout, "ROW-new", "GRID_ID")
assert "GRID_ID" in layout["ROOT_ID"]["children"]
assert "ROW-new" in layout["GRID_ID"]["children"]
class TestGenerateTitleFromCharts:
"""Tests for _generate_title_from_charts helper."""
@@ -1228,122 +1402,58 @@ class TestGenerateTitleFromCharts:
class TestDashboardSerializationEagerLoading:
"""Tests for eager loading fix in dashboard serialization paths.
"""Tests for eager loading fix in dashboard serialization paths."""
The re-fetch uses DashboardDAO.find_by_id() with query_options for
eager loading. A try/except around the DAO call handles "Can't
reconnect until invalid transaction is rolled back" errors in
multi-tenant environments by rolling back and falling back to the
original dashboard object.
"""
@patch("superset.models.dashboard.Dashboard")
@patch("superset.daos.dashboard.DashboardDAO.find_by_id")
@patch("superset.db.session")
@pytest.mark.asyncio
async def test_generate_dashboard_refetches_via_dao(
self, mock_db_session, mock_find_by_id, mock_dashboard_cls, mcp_server
):
"""generate_dashboard re-fetches dashboard via DashboardDAO.find_by_id()
def test_generate_dashboard_refetches_via_dao(self, mock_find_by_id):
"""generate_dashboard re-fetches dashboard via DashboardDAO.find_by_id
with eager-loaded slice relationships before serialization."""
charts = [_mock_chart(id=1, slice_name="Chart 1")]
dashboard = _mock_dashboard(id=10, title="Refetch Test")
_setup_generate_dashboard_mocks(
mock_db_session, mock_find_by_id, mock_dashboard_cls, charts, dashboard
refetched_dashboard = _mock_dashboard()
refetched_chart = _mock_chart(id=1, slice_name="Refetched Chart")
refetched_dashboard.slices = [refetched_chart]
mock_find_by_id.return_value = refetched_dashboard
from superset.daos.dashboard import DashboardDAO
result = (
DashboardDAO.find_by_id(1, query_options=["dummy"]) or _mock_dashboard()
)
request = {"chart_ids": [1], "dashboard_title": "Refetch Test"}
async with Client(mcp_server) as client:
result = await client.call_tool("generate_dashboard", {"request": request})
assert result is refetched_dashboard
mock_find_by_id.assert_called_once_with(1, query_options=["dummy"])
assert result.structured_content["error"] is None
# Verify DashboardDAO.find_by_id was called for re-fetch
mock_find_by_id.assert_called()
@patch("superset.models.dashboard.Dashboard")
@patch("superset.daos.dashboard.DashboardDAO.find_by_id")
@patch("superset.db.session")
@pytest.mark.asyncio
async def test_generate_dashboard_refetch_sqlalchemy_error_rollback(
self, mock_db_session, mock_find_by_id, mock_dashboard_cls, mcp_server
):
"""When the DAO re-fetch raises SQLAlchemyError, the session is
rolled back and a minimal response is returned with only scalar
attributes (no owners/tags/charts that would trigger lazy-loading)."""
from sqlalchemy.exc import SQLAlchemyError
def test_add_chart_refetches_dashboard_via_dao(self, mock_find_by_id):
"""add_chart_to_existing_dashboard re-fetches dashboard via
DashboardDAO.find_by_id with eager-loaded slice relationships."""
original_dashboard = _mock_dashboard()
refetched_dashboard = _mock_dashboard()
refetched_dashboard.slices = [_mock_chart()]
charts = [_mock_chart(id=1, slice_name="Chart 1")]
dashboard = _mock_dashboard(id=10, title="Rollback Test")
_setup_generate_dashboard_mocks(
mock_db_session, mock_find_by_id, mock_dashboard_cls, charts, dashboard
mock_find_by_id.return_value = refetched_dashboard
from superset.daos.dashboard import DashboardDAO
result = (
DashboardDAO.find_by_id(original_dashboard.id, query_options=["dummy"])
or original_dashboard
)
# Make the DAO re-fetch raise SQLAlchemyError
mock_find_by_id.side_effect = SQLAlchemyError("session error")
request = {"chart_ids": [1], "dashboard_title": "Rollback Test"}
async with Client(mcp_server) as client:
result = await client.call_tool("generate_dashboard", {"request": request})
assert result is refetched_dashboard
data = result.structured_content
assert data["error"] is None
mock_db_session.rollback.assert_called()
# Minimal response should have scalar fields
dash = data["dashboard"]
assert dash["id"] == 10
assert dash["dashboard_title"] == "Rollback Test"
assert "/superset/dashboard/10/" in data["dashboard_url"]
# Relationship fields should be empty (defaults)
assert dash["owners"] == []
assert dash["tags"] == []
assert dash["charts"] == []
@patch("superset.commands.dashboard.update.UpdateDashboardCommand")
@patch("superset.daos.dashboard.DashboardDAO.find_by_id")
@patch("superset.db.session")
@pytest.mark.asyncio
async def test_add_chart_refetch_sqlalchemy_error_rollback(
self, mock_db_session, mock_find_dashboard, mock_update_command, mcp_server
):
"""When the DAO re-fetch raises SQLAlchemyError after adding a chart,
the session is rolled back and a minimal response is returned with
only scalar attributes and position info."""
from sqlalchemy.exc import SQLAlchemyError
def test_add_chart_falls_back_on_refetch_failure(self, mock_find_by_id):
"""add_chart_to_existing_dashboard falls back to original dashboard
if DashboardDAO.find_by_id returns None."""
original_dashboard = _mock_dashboard()
mock_find_by_id.return_value = None
mock_dashboard = _mock_dashboard(id=1, title="Dashboard")
mock_dashboard.slices = []
mock_dashboard.position_json = "{}"
from superset.daos.dashboard import DashboardDAO
mock_chart = _mock_chart(id=15)
mock_db_session.get.return_value = mock_chart
result = (
DashboardDAO.find_by_id(original_dashboard.id, query_options=["dummy"])
or original_dashboard
)
updated = _mock_dashboard(id=1, title="Dashboard")
updated.slices = [_mock_chart(id=15)]
mock_update_command.return_value.run.return_value = updated
# First call returns dashboard (validation), second raises (re-fetch)
mock_find_dashboard.side_effect = [
mock_dashboard,
SQLAlchemyError("session error"),
]
request = {"dashboard_id": 1, "chart_id": 15}
async with Client(mcp_server) as client:
result = await client.call_tool(
"add_chart_to_existing_dashboard", {"request": request}
)
data = result.structured_content
assert data["error"] is None
mock_db_session.rollback.assert_called()
# Minimal response should have scalar fields
dash = data["dashboard"]
assert dash["id"] == 1
assert dash["dashboard_title"] == "Dashboard"
assert "/superset/dashboard/1/" in data["dashboard_url"]
# Position info should still be returned
assert data["position"] is not None
assert "chart_key" in data["position"]
# Relationship fields should be empty (defaults)
assert dash["owners"] == []
assert dash["tags"] == []
assert dash["charts"] == []
assert result is original_dashboard