From 1e2d0b5f5bb9043d497eb6110ec18c1c3ff775cf Mon Sep 17 00:00:00 2001 From: Richard Fogaca Nienkotter <63572350+richardfogaca@users.noreply.github.com> Date: Fri, 15 May 2026 12:15:33 -0300 Subject: [PATCH] fix(mcp): defer chart preview command imports (#40164) --- superset/mcp_service/chart/preview_utils.py | 2 +- .../mcp_service/chart/test_preview_utils.py | 41 +++++++++++++++++++ .../chart/tool/test_get_chart_data.py | 8 ++-- 3 files changed, 46 insertions(+), 5 deletions(-) diff --git a/superset/mcp_service/chart/preview_utils.py b/superset/mcp_service/chart/preview_utils.py index a5ae10b3de3..1adddad0daf 100644 --- a/superset/mcp_service/chart/preview_utils.py +++ b/superset/mcp_service/chart/preview_utils.py @@ -26,7 +26,6 @@ import logging import math from typing import Any, Dict, List -from superset.commands.chart.data.get_data_command import ChartDataCommand from superset.mcp_service.chart.schemas import ( ASCIIPreview, ChartError, @@ -78,6 +77,7 @@ def generate_preview_from_form_data( """ try: # Execute query to get data + from superset.commands.chart.data.get_data_command import ChartDataCommand from superset.connectors.sqla.models import SqlaTable from superset.extensions import db diff --git a/tests/unit_tests/mcp_service/chart/test_preview_utils.py b/tests/unit_tests/mcp_service/chart/test_preview_utils.py index 0190203e90e..7d5868cb4cf 100644 --- a/tests/unit_tests/mcp_service/chart/test_preview_utils.py +++ b/tests/unit_tests/mcp_service/chart/test_preview_utils.py @@ -19,6 +19,47 @@ Tests for preview_utils query context column building. """ +import ast +import inspect +from pathlib import Path + +from superset.mcp_service.chart import preview_utils + + +def _imports_chart_data_command(node: ast.Import | ast.ImportFrom) -> bool: + blocked_module = "superset.commands.chart.data.get_data_command" + + if isinstance(node, ast.Import): + return any( + alias.name == blocked_module or alias.name.startswith(f"{blocked_module}.") + for alias in node.names + ) + + module = node.module or "" + return ( + module == blocked_module + or module.startswith(f"{blocked_module}.") + or ( + module == "superset.commands.chart.data" + and any(alias.name == "get_data_command" for alias in node.names) + ) + ) + + +def test_preview_utils_does_not_top_level_import_chart_data_command(): + """preview_utils constants should stay safe to import before app setup.""" + source_path = inspect.getsourcefile(preview_utils) or preview_utils.__file__ + source = Path(source_path).read_text(encoding="utf-8") + tree = ast.parse(source) + top_level_imports = [ + node for node in tree.body if isinstance(node, (ast.Import, ast.ImportFrom)) + ] + + assert preview_utils.SUPPORTED_FORM_DATA_PREVIEW_FORMATS == frozenset( + {"ascii", "table", "vega_lite"} + ) + assert not any(_imports_chart_data_command(node) for node in top_level_imports) + class TestPreviewUtilsColumnBuilding: """Tests for x_axis + groupby column building in generate_preview_from_form_data. diff --git a/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py b/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py index b36e175d1aa..5404f8985b6 100644 --- a/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py +++ b/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py @@ -1007,12 +1007,12 @@ class TestChartDataCommandValidation: mock_dataset = MagicMock() mock_dataset.id = 10 - # ChartDataCommand is module-level import in preview_utils; - # db and QueryContextFactory are local imports inside the function. + # ChartDataCommand, db, and QueryContextFactory are local imports inside + # the function so preview_utils stays safe to import before app setup. with ( patch("superset.extensions.db") as mock_db, patch( - "superset.mcp_service.chart.preview_utils.ChartDataCommand", + "superset.commands.chart.data.get_data_command.ChartDataCommand", return_value=mock_command, ), patch( @@ -1061,7 +1061,7 @@ class TestChartDataCommandValidation: with ( patch("superset.extensions.db") as mock_db, patch( - "superset.mcp_service.chart.preview_utils.ChartDataCommand", + "superset.commands.chart.data.get_data_command.ChartDataCommand", return_value=mock_command, ), patch(