mirror of
https://github.com/apache/superset.git
synced 2026-04-07 10:31:50 +00:00
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:
@@ -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 (
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user