diff --git a/superset/mcp_service/chart/schemas.py b/superset/mcp_service/chart/schemas.py index c00adbdcac7..29dddb33fd9 100644 --- a/superset/mcp_service/chart/schemas.py +++ b/superset/mcp_service/chart/schemas.py @@ -31,6 +31,7 @@ from pydantic import ( ConfigDict, Field, field_validator, + model_serializer, model_validator, PositiveInt, ) @@ -79,8 +80,8 @@ class ChartLike(Protocol): class ChartInfo(BaseModel): """Full chart model with all possible attributes.""" - id: int = Field(..., description="Chart ID") - slice_name: str = Field(..., description="Chart name") + id: int | None = Field(None, description="Chart ID") + slice_name: str | None = Field(None, description="Chart name") viz_type: str | None = Field(None, description="Visualization type") datasource_name: str | None = Field(None, description="Datasource name") datasource_type: str | None = Field(None, description="Datasource type") @@ -109,6 +110,26 @@ class ChartInfo(BaseModel): owners: List[UserInfo] = Field(default_factory=list, description="Chart owners") model_config = ConfigDict(from_attributes=True, ser_json_timedelta="iso8601") + @model_serializer(mode="wrap", when_used="json") + def _filter_fields_by_context(self, serializer: Any, info: Any) -> Dict[str, Any]: + """Filter fields based on serialization context. + + If context contains 'select_columns', only include those fields. + Otherwise, include all fields (default behavior). + """ + # Get full serialization + data = serializer(self) + + # Check if we have a context with select_columns + if info.context and isinstance(info.context, dict): + select_columns = info.context.get("select_columns") + if select_columns: + # Filter to only requested fields + return {k: v for k, v in data.items() if k in select_columns} + + # No filtering - return all fields + return data + class GetChartAvailableFiltersRequest(BaseModel): """ diff --git a/superset/mcp_service/chart/tool/list_charts.py b/superset/mcp_service/chart/tool/list_charts.py index c46c27b2c95..2d0ad1e79ed 100644 --- a/superset/mcp_service/chart/tool/list_charts.py +++ b/superset/mcp_service/chart/tool/list_charts.py @@ -20,7 +20,7 @@ MCP tool: list_charts (advanced filtering with metadata cache control) """ import logging -from typing import Any, cast, TYPE_CHECKING +from typing import cast, TYPE_CHECKING from fastmcp import Context from superset_core.mcp import tool @@ -94,8 +94,10 @@ async def list_charts(request: ListChartsRequest, ctx: Context) -> ChartList: from superset.daos.chart import ChartDAO - def _serialize_chart(obj: "Slice | None", cols: Any) -> ChartInfo | None: - """Serialize chart object with proper type casting.""" + def _serialize_chart( + obj: "Slice | None", cols: list[str] | None + ) -> ChartInfo | None: + """Serialize chart object (field filtering handled by model_serializer).""" return serialize_chart_object(cast(ChartLike | None, obj)) tool = ModelListCore( @@ -129,7 +131,21 @@ async def list_charts(request: ListChartsRequest, ctx: Context) -> ChartList: "Charts listed successfully: count=%s, total_pages=%s" % (count, total_pages) ) - return result + + # Apply field filtering via serialization context if select_columns specified + # This triggers ChartInfo._filter_fields_by_context for each chart + if request.select_columns: + await ctx.debug( + "Applying field filtering via serialization context: select_columns=%s" + % (request.select_columns,) + ) + # Return dict with context - FastMCP will serialize it + return result.model_dump( + mode="json", context={"select_columns": request.select_columns} + ) + + # No filtering - return full result as dict + return result.model_dump(mode="json") except Exception as e: await ctx.error("Failed to list charts: %s" % (str(e),)) raise diff --git a/superset/mcp_service/dashboard/schemas.py b/superset/mcp_service/dashboard/schemas.py index 346892bc63d..1aff298b119 100644 --- a/superset/mcp_service/dashboard/schemas.py +++ b/superset/mcp_service/dashboard/schemas.py @@ -73,6 +73,7 @@ from pydantic import ( ConfigDict, Field, field_validator, + model_serializer, model_validator, PositiveInt, ) @@ -286,8 +287,8 @@ class GetDashboardInfoRequest(MetadataCacheControl): class DashboardInfo(BaseModel): - id: int = Field(..., description="Dashboard ID") - dashboard_title: str = Field(..., description="Dashboard title") + id: int | None = Field(None, description="Dashboard ID") + dashboard_title: str | None = Field(None, description="Dashboard title") slug: str | None = Field(None, description="Dashboard slug") description: str | None = Field(None, description="Dashboard description") css: str | None = Field(None, description="Custom CSS for the dashboard") @@ -328,6 +329,26 @@ class DashboardInfo(BaseModel): ) model_config = ConfigDict(from_attributes=True, ser_json_timedelta="iso8601") + @model_serializer(mode="wrap", when_used="json") + def _filter_fields_by_context(self, serializer: Any, info: Any) -> Dict[str, Any]: + """Filter fields based on serialization context. + + If context contains 'select_columns', only include those fields. + Otherwise, include all fields (default behavior). + """ + # Get full serialization + data = serializer(self) + + # Check if we have a context with select_columns + if info.context and isinstance(info.context, dict): + select_columns = info.context.get("select_columns") + if select_columns: + # Filter to only requested fields + return {k: v for k, v in data.items() if k in select_columns} + + # No filtering - return all fields + return data + class DashboardList(BaseModel): dashboards: List[DashboardInfo] diff --git a/superset/mcp_service/dashboard/tool/list_dashboards.py b/superset/mcp_service/dashboard/tool/list_dashboards.py index 6772a6a8e29..19fe69ecba9 100644 --- a/superset/mcp_service/dashboard/tool/list_dashboards.py +++ b/superset/mcp_service/dashboard/tool/list_dashboards.py @@ -23,10 +23,14 @@ advanced filtering with clear, unambiguous request schema and metadata cache con """ import logging +from typing import Any, TYPE_CHECKING from fastmcp import Context from superset_core.mcp import tool +if TYPE_CHECKING: + from superset.models.dashboard import Dashboard + from superset.mcp_service.dashboard.schemas import ( DashboardFilter, DashboardInfo, @@ -63,7 +67,7 @@ SORTABLE_DASHBOARD_COLUMNS = [ @parse_request(ListDashboardsRequest) async def list_dashboards( request: ListDashboardsRequest, ctx: Context -) -> DashboardList: +) -> dict[str, Any]: """List dashboards with filtering and search. Returns dashboard metadata including title, slug, and charts. @@ -72,10 +76,16 @@ async def list_dashboards( """ from superset.daos.dashboard import DashboardDAO + def _serialize_dashboard( + obj: "Dashboard | None", cols: list[str] | None + ) -> DashboardInfo | None: + """Serialize dashboard object (field filtering handled by model_serializer).""" + return serialize_dashboard_object(obj) + tool = ModelListCore( dao_class=DashboardDAO, output_schema=DashboardInfo, - item_serializer=lambda obj, cols: serialize_dashboard_object(obj), + item_serializer=_serialize_dashboard, filter_type=DashboardFilter, default_columns=DEFAULT_DASHBOARD_COLUMNS, search_columns=[ @@ -87,7 +97,8 @@ async def list_dashboards( output_list_schema=DashboardList, logger=logger, ) - return tool.run_tool( + + result = tool.run_tool( filters=request.filters, search=request.search, select_columns=request.select_columns, @@ -96,3 +107,18 @@ async def list_dashboards( page=max(request.page - 1, 0), page_size=request.page_size, ) + + # Apply field filtering via serialization context if select_columns specified + # This triggers DashboardInfo._filter_fields_by_context for each dashboard + if request.select_columns: + await ctx.debug( + "Applying field filtering via serialization context: select_columns=%s" + % (request.select_columns,) + ) + # Return dict with context - FastMCP will serialize it + return result.model_dump( + mode="json", context={"select_columns": request.select_columns} + ) + + # No filtering - return full result as dict + return result.model_dump(mode="json") diff --git a/superset/mcp_service/dataset/schemas.py b/superset/mcp_service/dataset/schemas.py index e7a6adc1300..2cefd04e63d 100644 --- a/superset/mcp_service/dataset/schemas.py +++ b/superset/mcp_service/dataset/schemas.py @@ -24,7 +24,14 @@ from __future__ import annotations from datetime import datetime from typing import Annotated, Any, Dict, List, Literal -from pydantic import BaseModel, ConfigDict, Field, model_validator, PositiveInt +from pydantic import ( + BaseModel, + ConfigDict, + Field, + model_serializer, + model_validator, + PositiveInt, +) from superset.daos.base import ColumnOperator, ColumnOperatorEnum from superset.mcp_service.common.cache_schemas import MetadataCacheControl @@ -156,6 +163,31 @@ class DatasetInfo(BaseModel): populate_by_name=True, # Allow both 'schema' (alias) and 'schema_name' (field) ) + @model_serializer(mode="wrap", when_used="json") + def _filter_fields_by_context(self, serializer: Any, info: Any) -> Dict[str, Any]: + """Filter fields based on serialization context. + + If context contains 'select_columns', only include those fields. + Otherwise, include all fields (default behavior). + """ + # Get full serialization + data = serializer(self) + + # Check if we have a context with select_columns + if info.context and isinstance(info.context, dict): + select_columns = info.context.get("select_columns") + if select_columns: + # Handle alias: 'schema' -> 'schema_name' + requested_fields = set(select_columns) + if "schema" in requested_fields: + requested_fields.add("schema_name") + + # Filter to only requested fields + return {k: v for k, v in data.items() if k in requested_fields} + + # No filtering - return all fields + return data + class DatasetList(BaseModel): datasets: List[DatasetInfo] diff --git a/superset/mcp_service/dataset/tool/list_datasets.py b/superset/mcp_service/dataset/tool/list_datasets.py index 31835d39b78..25da2fdab94 100644 --- a/superset/mcp_service/dataset/tool/list_datasets.py +++ b/superset/mcp_service/dataset/tool/list_datasets.py @@ -23,10 +23,14 @@ advanced filtering with clear, unambiguous request schema and metadata cache con """ import logging +from typing import TYPE_CHECKING from fastmcp import Context from superset_core.mcp import tool +if TYPE_CHECKING: + from superset.connectors.sqla.models import SqlaTable + from superset.mcp_service.dataset.schemas import ( DatasetFilter, DatasetInfo, @@ -102,11 +106,17 @@ async def list_datasets(request: ListDatasetsRequest, ctx: Context) -> DatasetLi try: from superset.daos.dataset import DatasetDAO + def _serialize_dataset( + obj: "SqlaTable | None", cols: list[str] | None + ) -> DatasetInfo | None: + """Serialize dataset (filtering via model_serializer).""" + return serialize_dataset_object(obj) + # Create tool with standard serialization tool = ModelListCore( dao_class=DatasetDAO, output_schema=DatasetInfo, - item_serializer=lambda obj, cols: serialize_dataset_object(obj), + item_serializer=_serialize_dataset, filter_type=DatasetFilter, default_columns=DEFAULT_DATASET_COLUMNS, search_columns=["schema", "sql", "table_name", "uuid"], @@ -134,7 +144,20 @@ async def list_datasets(request: ListDatasetsRequest, ctx: Context) -> DatasetLi ) ) - return result + # Apply field filtering via serialization context if select_columns specified + # This triggers DatasetInfo._filter_fields_by_context for each dataset + if request.select_columns: + await ctx.debug( + "Applying field filtering via serialization context: select_columns=%s" + % (request.select_columns,) + ) + # Return dict with context - FastMCP will serialize it + return result.model_dump( + mode="json", context={"select_columns": request.select_columns} + ) + + # No filtering - return full result as dict + return result.model_dump(mode="json") except Exception as e: await ctx.error( diff --git a/tests/unit_tests/mcp_service/dashboard/tool/test_dashboard_generation.py b/tests/unit_tests/mcp_service/dashboard/tool/test_dashboard_generation.py index b24d21e457a..32d9dda92ee 100644 --- a/tests/unit_tests/mcp_service/dashboard/tool/test_dashboard_generation.py +++ b/tests/unit_tests/mcp_service/dashboard/tool/test_dashboard_generation.py @@ -112,12 +112,17 @@ class TestGenerateDashboard: async with Client(mcp_server) as client: result = await client.call_tool("generate_dashboard", {"request": request}) - assert result.data.error is None - assert result.data.dashboard is not None - assert result.data.dashboard.id == 10 - assert result.data.dashboard.dashboard_title == "Analytics Dashboard" - assert result.data.dashboard.chart_count == 2 - assert "/superset/dashboard/10/" in result.data.dashboard_url + assert result.structured_content["error"] is None + assert result.structured_content["dashboard"] is not None + assert result.structured_content["dashboard"]["id"] == 10 + assert ( + result.structured_content["dashboard"]["dashboard_title"] + == "Analytics Dashboard" + ) + assert result.structured_content["dashboard"]["chart_count"] == 2 + assert ( + "/superset/dashboard/10/" in result.structured_content["dashboard_url"] + ) @patch("superset.db.session") @pytest.mark.asyncio @@ -138,10 +143,10 @@ class TestGenerateDashboard: async with Client(mcp_server) as client: result = await client.call_tool("generate_dashboard", {"request": request}) - assert result.data.error is not None - assert "Charts not found: [2]" in result.data.error - assert result.data.dashboard is None - assert result.data.dashboard_url is None + assert result.structured_content["error"] is not None + assert "Charts not found: [2]" in result.structured_content["error"] + assert result.structured_content["dashboard"] is None + assert result.structured_content["dashboard_url"] is None @patch("superset.commands.dashboard.create.CreateDashboardCommand") @patch("superset.db.session") @@ -169,9 +174,9 @@ class TestGenerateDashboard: async with Client(mcp_server) as client: result = await client.call_tool("generate_dashboard", {"request": request}) - assert result.data.error is None - assert result.data.dashboard.chart_count == 1 - assert result.data.dashboard.published is True # From mock + assert result.structured_content["error"] is None + assert result.structured_content["dashboard"]["chart_count"] == 1 + assert result.structured_content["dashboard"]["published"] is True @patch("superset.commands.dashboard.create.CreateDashboardCommand") @patch("superset.db.session") @@ -198,8 +203,8 @@ class TestGenerateDashboard: async with Client(mcp_server) as client: result = await client.call_tool("generate_dashboard", {"request": request}) - assert result.data.error is None - assert result.data.dashboard.chart_count == 6 + assert result.structured_content["error"] is None + assert result.structured_content["dashboard"]["chart_count"] == 6 # Verify CreateDashboardCommand was called with proper layout mock_create_command.assert_called_once() @@ -257,9 +262,9 @@ class TestGenerateDashboard: async with Client(mcp_server) as client: result = await client.call_tool("generate_dashboard", {"request": request}) - assert result.data.error is not None - assert "Failed to create dashboard" in result.data.error - assert result.data.dashboard is None + assert result.structured_content["error"] is not None + assert "Failed to create dashboard" in result.structured_content["error"] + assert result.structured_content["dashboard"] is None @patch("superset.commands.dashboard.create.CreateDashboardCommand") @patch("superset.db.session") @@ -287,8 +292,11 @@ class TestGenerateDashboard: async with Client(mcp_server) as client: result = await client.call_tool("generate_dashboard", {"request": request}) - assert result.data.error is None - assert result.data.dashboard.dashboard_title == "Minimal Dashboard" + assert result.structured_content["error"] is None + assert ( + result.structured_content["dashboard"]["dashboard_title"] + == "Minimal Dashboard" + ) # Check that description was not included in call call_args = mock_create_command.call_args[0][0] @@ -343,13 +351,15 @@ class TestAddChartToExistingDashboard: "add_chart_to_existing_dashboard", {"request": request} ) - assert result.data.error is None - assert result.data.dashboard is not None - assert result.data.dashboard.chart_count == 3 - assert result.data.position is not None - assert "row" in result.data.position # Should have row info - assert "chart_key" in result.data.position - assert "/superset/dashboard/1/" in result.data.dashboard_url + assert result.structured_content["error"] is None + assert result.structured_content["dashboard"] is not None + assert result.structured_content["dashboard"]["chart_count"] == 3 + assert result.structured_content["position"] is not None + assert "row" in result.structured_content["position"] + assert "chart_key" in result.structured_content["position"] + assert ( + "/superset/dashboard/1/" in result.structured_content["dashboard_url"] + ) @patch("superset.daos.dashboard.DashboardDAO.find_by_id") @pytest.mark.asyncio @@ -364,8 +374,10 @@ class TestAddChartToExistingDashboard: "add_chart_to_existing_dashboard", {"request": request} ) - assert result.data.error is not None - assert "Dashboard with ID 999 not found" in result.data.error + assert result.structured_content["error"] is not None + assert ( + "Dashboard with ID 999 not found" in result.structured_content["error"] + ) @patch("superset.daos.dashboard.DashboardDAO.find_by_id") @patch("superset.db.session") @@ -384,8 +396,8 @@ class TestAddChartToExistingDashboard: "add_chart_to_existing_dashboard", {"request": request} ) - assert result.data.error is not None - assert "Chart with ID 999 not found" in result.data.error + assert result.structured_content["error"] is not None + assert "Chart with ID 999 not found" in result.structured_content["error"] @patch("superset.daos.dashboard.DashboardDAO.find_by_id") @patch("superset.db.session") @@ -407,8 +419,11 @@ class TestAddChartToExistingDashboard: "add_chart_to_existing_dashboard", {"request": request} ) - assert result.data.error is not None - assert "Chart 5 is already in dashboard 1" in result.data.error + assert result.structured_content["error"] is not None + assert ( + "Chart 5 is already in dashboard 1" + in result.structured_content["error"] + ) @patch("superset.commands.dashboard.update.UpdateDashboardCommand") @patch("superset.daos.dashboard.DashboardDAO.find_by_id") @@ -437,9 +452,9 @@ class TestAddChartToExistingDashboard: "add_chart_to_existing_dashboard", {"request": request} ) - assert result.data.error is None - assert "row" in result.data.position # Should have row info - assert result.data.position.get("row") == 0 # First row + assert result.structured_content["error"] is None + assert "row" in result.structured_content["position"] + assert result.structured_content["position"].get("row") == 0 # Verify update was called with proper layout structure call_args = mock_update_command.call_args[0][1] diff --git a/tests/unit_tests/mcp_service/dashboard/tool/test_dashboard_tools.py b/tests/unit_tests/mcp_service/dashboard/tool/test_dashboard_tools.py index 548c546930e..2e05a758eff 100644 --- a/tests/unit_tests/mcp_service/dashboard/tool/test_dashboard_tools.py +++ b/tests/unit_tests/mcp_service/dashboard/tool/test_dashboard_tools.py @@ -103,18 +103,17 @@ async def test_list_dashboards_basic(mock_list, mcp_server): result = await client.call_tool( "list_dashboards", {"request": request.model_dump()} ) - dashboards = result.data.dashboards + dashboards = result.data["dashboards"] assert len(dashboards) == 1 - assert dashboards[0].dashboard_title == "Test Dashboard" - assert dashboards[0].uuid == "test-dashboard-uuid-1" - assert dashboards[0].slug == "test-dashboard" - assert dashboards[0].published is True + assert dashboards[0]["dashboard_title"] == "Test Dashboard" + assert dashboards[0]["uuid"] == "test-dashboard-uuid-1" + assert dashboards[0]["slug"] == "test-dashboard" + assert dashboards[0]["published"] is True - # Verify UUID and slug are in default columns - assert "uuid" in result.data.columns_requested - assert "slug" in result.data.columns_requested - assert "uuid" in result.data.columns_loaded - assert "slug" in result.data.columns_loaded + assert "uuid" in result.data["columns_requested"] + assert "slug" in result.data["columns_requested"] + assert "uuid" in result.data["columns_loaded"] + assert "slug" in result.data["columns_loaded"] @patch("superset.daos.dashboard.DashboardDAO.list") @@ -180,8 +179,8 @@ async def test_list_dashboards_with_filters(mock_list, mcp_server): result = await client.call_tool( "list_dashboards", {"request": request.model_dump()} ) - assert result.data.count == 1 - assert result.data.dashboards[0].dashboard_title == "Filtered Dashboard" + assert result.data["count"] == 1 + assert result.data["dashboards"][0]["dashboard_title"] == "Filtered Dashboard" @patch("superset.daos.dashboard.DashboardDAO.list") @@ -262,8 +261,8 @@ async def test_list_dashboards_with_search(mock_list, mcp_server): result = await client.call_tool( "list_dashboards", {"request": request.model_dump()} ) - assert result.data.count == 1 - assert result.data.dashboards[0].dashboard_title == "search_dashboard" + assert result.data["count"] == 1 + assert result.data["dashboards"][0]["dashboard_title"] == "search_dashboard" args, kwargs = mock_list.call_args assert kwargs["search"] == "search_dashboard" assert "dashboard_title" in kwargs["search_columns"] @@ -283,7 +282,7 @@ async def test_list_dashboards_with_simple_filters(mock_list, mcp_server): result = await client.call_tool( "list_dashboards", {"request": request.model_dump()} ) - assert hasattr(result.data, "count") + assert "count" in result.data @patch("superset.daos.dashboard.DashboardDAO.find_by_id") @@ -503,16 +502,15 @@ async def test_list_dashboards_custom_uuid_slug_columns(mock_list, mcp_server): result = await client.call_tool( "list_dashboards", {"request": request.model_dump()} ) - dashboards = result.data.dashboards + dashboards = result.data["dashboards"] assert len(dashboards) == 1 - assert dashboards[0].uuid == "test-custom-uuid-123" - assert dashboards[0].slug == "custom-dashboard" + assert dashboards[0]["uuid"] == "test-custom-uuid-123" + assert dashboards[0]["slug"] == "custom-dashboard" - # Verify custom columns include UUID and slug - assert "uuid" in result.data.columns_requested - assert "slug" in result.data.columns_requested - assert "uuid" in result.data.columns_loaded - assert "slug" in result.data.columns_loaded + assert "uuid" in result.data["columns_requested"] + assert "slug" in result.data["columns_requested"] + assert "uuid" in result.data["columns_loaded"] + assert "slug" in result.data["columns_loaded"] class TestDashboardSortableColumns: @@ -558,9 +556,8 @@ class TestDashboardSortableColumns: assert call_args["order_column"] == "dashboard_title" assert call_args["order_direction"] == "desc" - # Verify the result - assert result.data.count == 0 - assert result.data.dashboards == [] + assert result.data["count"] == 0 + assert result.data["dashboards"] == [] def test_sortable_columns_in_docstring(self): """Test that sortable columns are documented in tool docstring.""" diff --git a/tests/unit_tests/mcp_service/dataset/tool/test_dataset_tools.py b/tests/unit_tests/mcp_service/dataset/tool/test_dataset_tools.py index 30828b075b8..6f4eb75bf4f 100644 --- a/tests/unit_tests/mcp_service/dataset/tool/test_dataset_tools.py +++ b/tests/unit_tests/mcp_service/dataset/tool/test_dataset_tools.py @@ -358,6 +358,8 @@ async def test_list_datasets_with_filters(mock_list, mcp_server): "params": dataset.params, "template_params": dataset.template_params, "extra": dataset.extra, + "columns": dataset.columns, + "metrics": dataset.metrics, } mock_list.return_value = ([dataset], 1) filters = [ @@ -382,9 +384,6 @@ async def test_list_datasets_with_filters(mock_list, mcp_server): assert len(data["datasets"]) == 1 assert data["datasets"][0]["id"] == 2 assert data["datasets"][0]["table_name"] == "Filtered Dataset" - # Check that columns and metrics are included - assert len(data["datasets"][0]["columns"]) == 1 - assert len(data["datasets"][0]["metrics"]) == 1 @patch("superset.daos.dataset.DatasetDAO.list") @@ -721,6 +720,8 @@ async def test_list_datasets_simple_basic(mock_list, mcp_server): "params": dataset.params, "template_params": dataset.template_params, "extra": dataset.extra, + "columns": dataset.columns, + "metrics": dataset.metrics, } mock_list.return_value = ([dataset], 1) filters = [ @@ -816,6 +817,8 @@ async def test_list_datasets_simple_with_filters(mock_list, mcp_server): "params": dataset.params, "template_params": dataset.template_params, "extra": dataset.extra, + "columns": dataset.columns, + "metrics": dataset.metrics, } mock_list.return_value = ([dataset], 1) filters = [ @@ -1080,17 +1083,19 @@ async def test_list_datasets_includes_columns_and_metrics(mock_list, mcp_server) result = await client.call_tool( "list_datasets", {"request": request.model_dump()} ) - datasets = result.data.datasets + assert result.content is not None + data = json.loads(result.content[0].text) + datasets = data["datasets"] assert len(datasets) == 1 ds = datasets[0] - assert hasattr(ds, "columns") - assert hasattr(ds, "metrics") - assert isinstance(ds.columns, list) - assert isinstance(ds.metrics, list) - assert len(ds.columns) == 1 - assert len(ds.metrics) == 1 - assert ds.columns[0].column_name == "colA" - assert ds.metrics[0].metric_name == "avg_value" + assert "columns" in ds + assert "metrics" in ds + assert isinstance(ds["columns"], list) + assert isinstance(ds["metrics"], list) + assert len(ds["columns"]) == 1 + assert len(ds["metrics"]) == 1 + assert ds["columns"][0]["column_name"] == "colA" + assert ds["metrics"][0]["metric_name"] == "avg_value" @patch("superset.mcp_service.mcp_core.ModelGetInfoCore._find_object")