aminghadersohi commented on code in PR #36458: URL: https://github.com/apache/superset/pull/36458#discussion_r2599908748
########## superset/mcp_service/system/resources/schema_discovery.py: ########## @@ -0,0 +1,194 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +""" +MCP Resources for schema discovery. + +These resources provide static metadata about available columns, filters, +and sorting options for chart, dataset, and dashboard models. +LLM clients can use these resources to understand what parameters are valid +without making API calls. +""" + +import logging +from typing import Any + +from superset.mcp_service.app import mcp +from superset.mcp_service.auth import mcp_auth_hook + +logger = logging.getLogger(__name__) + + +def _build_schema_resource(model_type: str) -> dict[str, Any]: + """Build schema resource data for a model type.""" + from superset.mcp_service.common.schema_discovery import ( + CHART_ALL_COLUMNS, + CHART_DEFAULT_COLUMNS, + CHART_SEARCH_COLUMNS, + CHART_SELECT_COLUMNS, + CHART_SORTABLE_COLUMNS, + DASHBOARD_ALL_COLUMNS, + DASHBOARD_DEFAULT_COLUMNS, + DASHBOARD_SEARCH_COLUMNS, + DASHBOARD_SELECT_COLUMNS, + DASHBOARD_SORTABLE_COLUMNS, + DATASET_ALL_COLUMNS, + DATASET_DEFAULT_COLUMNS, + DATASET_SEARCH_COLUMNS, + DATASET_SELECT_COLUMNS, + DATASET_SORTABLE_COLUMNS, + ) + + schemas = { + "chart": { + "model_type": "chart", + "select_columns": [ + {"name": col.name, "description": col.description, "type": col.type} + for col in CHART_SELECT_COLUMNS + ], + "all_column_names": CHART_ALL_COLUMNS, + "default_columns": CHART_DEFAULT_COLUMNS, + "sortable_columns": CHART_SORTABLE_COLUMNS, + "search_columns": CHART_SEARCH_COLUMNS, + "default_sort": "changed_on", + "default_sort_direction": "desc", + }, + "dataset": { + "model_type": "dataset", + "select_columns": [ + {"name": col.name, "description": col.description, "type": col.type} + for col in DATASET_SELECT_COLUMNS + ], + "all_column_names": DATASET_ALL_COLUMNS, + "default_columns": DATASET_DEFAULT_COLUMNS, + "sortable_columns": DATASET_SORTABLE_COLUMNS, + "search_columns": DATASET_SEARCH_COLUMNS, + "default_sort": "changed_on", + "default_sort_direction": "desc", + }, + "dashboard": { + "model_type": "dashboard", + "select_columns": [ + {"name": col.name, "description": col.description, "type": col.type} Review Comment: Addressed - Same fix applied using the shared `_safe_col_dict()` helper for JSON-safe serialization. ########## tests/unit_tests/mcp_service/system/tool/test_get_schema.py: ########## @@ -0,0 +1,376 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +""" +Tests for the get_schema unified schema discovery tool. +""" + +from unittest.mock import patch + +import pytest +from fastmcp import Client + +from superset.mcp_service.app import mcp +from superset.mcp_service.common.schema_discovery import ( + CHART_DEFAULT_COLUMNS, + CHART_SEARCH_COLUMNS, + CHART_SORTABLE_COLUMNS, + DASHBOARD_DEFAULT_COLUMNS, + DASHBOARD_SEARCH_COLUMNS, + DASHBOARD_SORTABLE_COLUMNS, + DATASET_DEFAULT_COLUMNS, + DATASET_SEARCH_COLUMNS, + DATASET_SORTABLE_COLUMNS, + GetSchemaRequest, + ModelSchemaInfo, +) +from superset.utils import json + + [email protected] +def mcp_server(): + return mcp + + [email protected](autouse=True) +def mock_auth(): + """Mock authentication for all tests.""" + from unittest.mock import Mock + + with patch("superset.mcp_service.auth.get_user_from_request") as mock_get_user: + mock_user = Mock() + mock_user.id = 1 + mock_user.username = "admin" + mock_get_user.return_value = mock_user + yield mock_get_user + + +class TestGetSchemaRequest: + """Test the GetSchemaRequest schema validation.""" + + def test_chart_model_type(self): + """Test creating request for chart model type.""" + request = GetSchemaRequest(model_type="chart") + assert request.model_type == "chart" + + def test_dataset_model_type(self): + """Test creating request for dataset model type.""" + request = GetSchemaRequest(model_type="dataset") + assert request.model_type == "dataset" + + def test_dashboard_model_type(self): + """Test creating request for dashboard model type.""" + request = GetSchemaRequest(model_type="dashboard") + assert request.model_type == "dashboard" + + def test_invalid_model_type(self): + """Test that invalid model types raise validation error.""" + with pytest.raises(ValueError, match="Input should be"): + GetSchemaRequest(model_type="invalid") + + +class TestModelSchemaInfo: + """Test the ModelSchemaInfo response structure.""" + + def test_chart_schema_info(self): + """Test chart schema info structure.""" + from superset.mcp_service.common.schema_discovery import CHART_SELECT_COLUMNS + + info = ModelSchemaInfo( + model_type="chart", + select_columns=CHART_SELECT_COLUMNS, + filter_columns={"slice_name": ["eq", "sw"]}, + sortable_columns=CHART_SORTABLE_COLUMNS, + default_select=CHART_DEFAULT_COLUMNS, + default_sort="changed_on", + default_sort_direction="desc", + search_columns=CHART_SEARCH_COLUMNS, + ) + + assert info.model_type == "chart" + assert len(info.select_columns) > 0 + assert len(info.default_select) == 4 # id, slice_name, viz_type, uuid + assert "slice_name" in info.filter_columns + assert info.default_sort == "changed_on" + + def test_dataset_default_columns(self): + """Test dataset default columns are minimal set.""" + assert len(DATASET_DEFAULT_COLUMNS) == 4 + assert "id" in DATASET_DEFAULT_COLUMNS + assert "table_name" in DATASET_DEFAULT_COLUMNS + assert "schema" in DATASET_DEFAULT_COLUMNS + assert "uuid" in DATASET_DEFAULT_COLUMNS + # These should NOT be in defaults + assert "columns" not in DATASET_DEFAULT_COLUMNS + assert "metrics" not in DATASET_DEFAULT_COLUMNS + + def test_dashboard_default_columns(self): + """Test dashboard default columns are minimal set.""" + assert len(DASHBOARD_DEFAULT_COLUMNS) == 4 + assert "id" in DASHBOARD_DEFAULT_COLUMNS + assert "dashboard_title" in DASHBOARD_DEFAULT_COLUMNS + assert "slug" in DASHBOARD_DEFAULT_COLUMNS + assert "uuid" in DASHBOARD_DEFAULT_COLUMNS + # These should NOT be in defaults + assert "published" not in DASHBOARD_DEFAULT_COLUMNS + assert "charts" not in DASHBOARD_DEFAULT_COLUMNS + + def test_chart_default_columns(self): + """Test chart default columns are minimal set.""" + assert len(CHART_DEFAULT_COLUMNS) == 4 + assert "id" in CHART_DEFAULT_COLUMNS + assert "slice_name" in CHART_DEFAULT_COLUMNS + assert "viz_type" in CHART_DEFAULT_COLUMNS + assert "uuid" in CHART_DEFAULT_COLUMNS + # These should NOT be in defaults + assert "description" not in CHART_DEFAULT_COLUMNS + assert "form_data" not in CHART_DEFAULT_COLUMNS + + +class TestGetSchemaToolViaClient: + """Test the get_schema tool via MCP client.""" + + @patch("superset.daos.chart.ChartDAO.get_filterable_columns_and_operators") + @pytest.mark.asyncio + async def test_get_schema_chart(self, mock_filters, mcp_server): + """Test get_schema for chart model type.""" + mock_filters.return_value = { + "slice_name": ["eq", "sw", "ilike"], + "viz_type": ["eq", "in"], + } + + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_schema", {"request": {"model_type": "chart"}} + ) + + assert result.content is not None + data = json.loads(result.content[0].text) + + # Check structure + assert "schema_info" in data + info = data["schema_info"] + assert info["model_type"] == "chart" + assert "select_columns" in info + assert "filter_columns" in info + assert "sortable_columns" in info + assert "default_select" in info + assert "search_columns" in info + + # Check default columns are minimal + assert len(info["default_select"]) == 4 + assert set(info["default_select"]) == { + "id", + "slice_name", + "viz_type", + "uuid", + } + + @patch("superset.daos.dataset.DatasetDAO.get_filterable_columns_and_operators") + @pytest.mark.asyncio + async def test_get_schema_dataset(self, mock_filters, mcp_server): + """Test get_schema for dataset model type.""" + mock_filters.return_value = { + "table_name": ["eq", "sw"], + "schema": ["eq"], + } + + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_schema", {"request": {"model_type": "dataset"}} + ) + + assert result.content is not None + data = json.loads(result.content[0].text) + + info = data["schema_info"] + assert info["model_type"] == "dataset" + + # Check default columns are minimal + assert len(info["default_select"]) == 4 + assert set(info["default_select"]) == { + "id", + "table_name", + "schema", + "uuid", + } + + # Check search columns + assert "table_name" in info["search_columns"] + assert "description" in info["search_columns"] + + @patch("superset.daos.dashboard.DashboardDAO.get_filterable_columns_and_operators") + @pytest.mark.asyncio + async def test_get_schema_dashboard(self, mock_filters, mcp_server): + """Test get_schema for dashboard model type.""" + mock_filters.return_value = { + "dashboard_title": ["eq", "ilike"], + "published": ["eq"], + } + + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_schema", {"request": {"model_type": "dashboard"}} + ) + + assert result.content is not None + data = json.loads(result.content[0].text) + + info = data["schema_info"] + assert info["model_type"] == "dashboard" + + # Check default columns are minimal + assert len(info["default_select"]) == 4 + assert set(info["default_select"]) == { + "id", + "dashboard_title", + "slug", + "uuid", + } + + # Check sortable columns include expected values + assert "dashboard_title" in info["sortable_columns"] + assert "changed_on" in info["sortable_columns"] + + @patch("superset.daos.chart.ChartDAO.get_filterable_columns_and_operators") + @pytest.mark.asyncio + async def test_get_schema_with_json_string_request(self, mock_filters, mcp_server): + """Test get_schema accepts JSON string request (Claude Code compatibility).""" + mock_filters.return_value = {"slice_name": ["eq"]} + + async with Client(mcp_server) as client: + # Send request as JSON string (Claude Code bug workaround) + result = await client.call_tool( + "get_schema", {"request": '{"model_type": "chart"}'} + ) + + assert result.content is not None + data = json.loads(result.content[0].text) + assert data["schema_info"]["model_type"] == "chart" + + @patch("superset.daos.chart.ChartDAO.get_filterable_columns_and_operators") + @pytest.mark.asyncio + async def test_get_schema_select_columns_have_metadata( + self, mock_filters, mcp_server + ): + """Test that select_columns include metadata (description, type).""" + mock_filters.return_value = {} + + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_schema", {"request": {"model_type": "chart"}} + ) + + data = json.loads(result.content[0].text) Review Comment: Already addressed - All tests have `assert result.content is not None` before accessing `result.content[0].text`. This assertion exists in every test that accesses content. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
