aminghadersohi commented on code in PR #41611: URL: https://github.com/apache/superset/pull/41611#discussion_r3537922135
########## superset/mcp_service/semantic_layer/schemas.py: ########## @@ -0,0 +1,302 @@ +# 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. + +"""Pydantic schemas for semantic layer MCP tools.""" + +from __future__ import annotations + +from typing import Any, Literal + +from pydantic import BaseModel, Field + +from superset.mcp_service.chart.schemas import DataColumn, PerformanceMetadata +from superset.mcp_service.common.cache_schemas import CacheStatus +from superset.mcp_service.common.error_schemas import MCPBaseError + +# --------------------------------------------------------------------------- +# Shared error schema +# --------------------------------------------------------------------------- + + +class SemanticLayerError(MCPBaseError): + """Error response returned by semantic layer tools.""" + + success: Literal[False] = False + + @classmethod + def create(cls, *, error: str, error_type: str) -> "SemanticLayerError": Review Comment: Not applying — annotating the `cls` receiver parameter would be inconsistent with every other `.create()` classmethod across `mcp_service/*/schemas.py` (TaskError, ReportError, RlsFilterError, UserError, DatabaseError, TagError, DatasetError, RoleError, DashboardError, AnnotationLayerError, etc.) — none of them annotate `cls`, and this isn't idiomatic mypy usage for classmethods returning `cls(...)`. Left as-is for consistency with the rest of the codebase. ########## superset/daos/semantic_layer.py: ########## @@ -201,6 +201,40 @@ def validate_update_uniqueness( for c in candidates ) + @classmethod + def find_accessible(cls) -> list[SemanticView]: Review Comment: Same as the `SemanticLayerError.create` comment above — leaving `cls` unannotated on `find_accessible` for consistency with every other classmethod in this codebase (none annotate `cls`). ########## superset/mcp_service/chart/schemas.py: ########## @@ -2246,10 +2246,20 @@ class DataColumn(BaseModel): display_name: str = Field(..., description="Human-readable column name") data_type: str = Field(..., description="Inferred data type") sample_values: List[Any] = Field(description="Representative sample values") - null_count: int = Field(description="Number of null values") - unique_count: int = Field(description="Number of unique values") + null_count: int = Field( + description="Number of null values. Approximate — see 'statistics.sampled_rows'" + " if the source result set was larger than the row cap used to compute it." + ) + unique_count: int = Field( + description="Number of unique values. Approximate — see " + "'statistics.sampled_rows' if the source result set was larger than the " + "row cap used to compute it." + ) statistics: Dict[str, Any] | None = Field( - None, description="Column statistics if numeric" + None, + description="Column statistics if numeric. May include 'sampled_rows' " + "when null_count/unique_count were computed on a row-capped sample " + "rather than the full result set.", ) Review Comment: Fixed — reworded the `statistics` field description to drop the numeric-only claim: "Additional column statistics, when available. May include 'sampled_rows' (any column type) when null_count/unique_count were computed on a row-capped sample rather than the full result set." ########## superset/mcp_service/semantic_layer/tool/list_metrics.py: ########## @@ -0,0 +1,336 @@ +# 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 tool: list_metrics + +Unified metric discovery across built-in datasets and external semantic views. +""" + +import logging +from typing import Any + +from fastmcp import Context +from superset_core.mcp.decorators import tool, ToolAnnotations + +from superset.daos.dataset import DatasetDAO +from superset.daos.semantic_layer import SemanticViewDAO +from superset.exceptions import SupersetSecurityException +from superset.extensions import db, event_logger +from superset.mcp_service.privacy import ( + DATA_MODEL_METADATA_ERROR_TYPE, + requires_data_model_metadata_access, + user_can_view_data_model_metadata, +) +from superset.mcp_service.semantic_layer.schemas import ( + DimensionInfo, + ListMetricsRequest, + MetricInfo, + MetricList, + SemanticLayerError, +) +from superset.semantic_layers.models import SemanticView + +logger = logging.getLogger(__name__) + + +def _matches_search(text: str | None, search: str) -> bool: + if not text: + return False + return search.lower() in text.lower() + + +def _builtin_compatible_dims(dataset: Any) -> list[DimensionInfo]: + """Return groupby-enabled columns as compatible dimensions for a built-in metric.""" + return [ + DimensionInfo( + name=col.column_name, + verbose_name=col.verbose_name or None, + description=col.description or None, + type=col.type or None, + is_dttm=bool(col.is_dttm), + groupby=bool(col.groupby), + filterable=bool(col.filterable), + source="builtin", + ) + for col in dataset.columns + if col.groupby + ] + + +def _collect_builtin_metrics(request: ListMetricsRequest) -> list[MetricInfo]: + """Collect metrics from built-in SqlaTable datasets.""" + from sqlalchemy.orm import subqueryload + + from superset.connectors.sqla.models import SqlaTable + + with event_logger.log_context(action="mcp.list_metrics.builtin_query"): + if request.dataset_id is not None: + from sqlalchemy.orm import subqueryload + + from superset.connectors.sqla.models import SqlaTable + + dataset: SqlaTable | None = DatasetDAO.find_by_id( + request.dataset_id, + query_options=[ + subqueryload(SqlaTable.columns), + subqueryload(SqlaTable.metrics), + ], + ) + datasets = [dataset] if dataset else [] + else: + from sqlalchemy.orm import subqueryload + + from superset.connectors.sqla.models import SqlaTable + + # Use _apply_base_filter with explicit eager loading to avoid + # N+1 queries when iterating dataset.metrics / dataset.columns. + query = db.session.query(SqlaTable).options( + subqueryload(SqlaTable.columns), + subqueryload(SqlaTable.metrics), + ) + datasets = DatasetDAO._apply_base_filter(query).all() + + results: list[MetricInfo] = [] + for dataset in datasets: + compat_dims = ( + _builtin_compatible_dims(dataset) + if request.include_compatible_dimensions + else [] + ) + for metric in dataset.metrics: + name = metric.metric_name or "" + desc = metric.description or "" + if request.search and not ( + _matches_search(name, request.search) + or _matches_search(desc, request.search) + ): + continue + results.append( + MetricInfo( + name=name, + verbose_name=metric.verbose_name or None, + description=desc or None, + expression=metric.expression or None, + d3format=metric.d3format or None, + warning_text=metric.warning_text or None, + source="builtin", + dataset_id=dataset.id, + dataset_name=dataset.table_name, + compatible_dimensions=compat_dims, + ) + ) + return results + + +async def _collect_external_metrics( + request: ListMetricsRequest, + ctx: Context, +) -> list[MetricInfo]: + """Collect metrics from external SemanticView models.""" + with event_logger.log_context(action="mcp.list_metrics.external_query"): + if request.view_id is not None: + view: SemanticView | None = SemanticViewDAO.find_by_id(request.view_id) + views = [view] if view else [] + else: + # find_accessible filters at SQL level, avoiding a per-row + # Python permission check and the audit noise of raise_for_access. + views = SemanticViewDAO.find_accessible() + + await ctx.debug("Found %d semantic views to scan for metrics" % len(views)) + + # Explicit single-view lookups aren't pre-filtered like find_accessible(), + # so raise_for_access must run outside the broad except block below so + # access errors are never silently swallowed as "could not load metrics". + check_access = request.view_id is not None + + results: list[MetricInfo] = [] + for view in views: + if check_access: + view.raise_for_access() + try: + raw_metrics = view.metrics + all_cols = { + col.column_name: col + for col in ( + view.columns if request.include_compatible_dimensions else [] + ) + } + for metric in raw_metrics: + name = metric.metric_name or "" + desc = metric.description or "" + if request.search and not ( + _matches_search(name, request.search) + or _matches_search(desc, request.search) + ): + continue + compat_dims: list[DimensionInfo] = [] + if request.include_compatible_dimensions: + # Unlike built-in SQL datasets, external views enforce + # per-metric compatibility, so dimensions must be resolved + # per metric rather than reused across all metrics. + compatible_names = view.get_compatible_dimensions([name], []) + compat_dims = [ + DimensionInfo( + name=dim_name, + verbose_name=all_cols[dim_name].verbose_name + if dim_name in all_cols + else None, + description=all_cols[dim_name].description + if dim_name in all_cols + else None, + type=all_cols[dim_name].type + if dim_name in all_cols + else None, + is_dttm=all_cols[dim_name].is_dttm + if dim_name in all_cols + else False, + groupby=all_cols[dim_name].groupby + if dim_name in all_cols + else True, + filterable=all_cols[dim_name].filterable + if dim_name in all_cols + else True, + source="external", + ) + for dim_name in compatible_names + ] + results.append( + MetricInfo( + name=name, + description=desc or None, + expression=metric.expression or None, + source="external", + view_id=view.id, + view_name=view.name, + compatible_dimensions=compat_dims, + ) Review Comment: Fixed — added `verbose_name=metric.verbose_name or None` to the external `MetricInfo` construction, matching the builtin path. Also added a regression test (test_list_metrics_external_includes_verbose_name). ########## tests/unit_tests/mcp_service/semantic_layer/tool/test_get_compatible_dimensions.py: ########## @@ -0,0 +1,232 @@ +# 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. + +"""Unit tests for the get_compatible_dimensions MCP tool.""" + +from __future__ import annotations + +import importlib +from collections.abc import Generator +from types import ModuleType +from unittest.mock import MagicMock, Mock, patch + +import pytest +from fastmcp import Client, FastMCP + +from superset.errors import ErrorLevel, SupersetError, SupersetErrorType +from superset.exceptions import SupersetSecurityException +from superset.mcp_service.app import mcp +from superset.utils import json + +get_compatible_dimensions_module: ModuleType = importlib.import_module( + "superset.mcp_service.semantic_layer.tool.get_compatible_dimensions" +) + + [email protected] +def mcp_server() -> FastMCP: + return mcp + + [email protected](autouse=True) +def mock_auth() -> Generator[MagicMock, None, None]: + with ( + patch("superset.mcp_service.auth.get_user_from_request") as mock_get_user, + patch.object( + get_compatible_dimensions_module, + "user_can_view_data_model_metadata", + return_value=True, + ), + ): + mock_user = Mock() + mock_user.id = 1 + mock_user.username = "admin" + mock_get_user.return_value = mock_user + yield mock_get_user + + +def _make_column(name: str, groupby: bool = True) -> MagicMock: + col = MagicMock() + col.column_name = name + col.verbose_name = None + col.description = None + col.type = "VARCHAR" + col.is_dttm = False + col.groupby = groupby + col.filterable = True + return col + + +def _make_dataset(dataset_id: int = 42) -> MagicMock: + ds = MagicMock() + ds.id = dataset_id + ds.table_name = f"table_{dataset_id}" + ds.metrics = [] + ds.columns = [ + _make_column("region"), + _make_column("category"), + _make_column("internal_only", groupby=False), + ] + return ds + + +def _make_view(view_id: int = 5) -> MagicMock: + view = MagicMock() + view.id = view_id + view.name = f"view_{view_id}" + view.raise_for_access = MagicMock(return_value=None) + view.columns = [_make_column("country_name")] + view.get_compatible_dimensions = MagicMock(return_value=["country_name"]) + return view + + +def _access_denied_exc(message: str = "Access denied") -> SupersetSecurityException: + return SupersetSecurityException( + SupersetError( + message=message, + error_type=SupersetErrorType.DATASOURCE_SECURITY_ACCESS_ERROR, + level=ErrorLevel.ERROR, + ) + ) + + [email protected] +async def test_get_compatible_dimensions_builtin_happy_path( + mcp_server: FastMCP, +) -> None: + """Builtin datasets return all groupby-enabled columns, ignoring selection.""" + mock_ds = _make_dataset(42) + + with patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=mock_ds): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_compatible_dimensions", + {"request": {"dataset_id": 42, "selected_metrics": ["revenue"]}}, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is True + assert data["source"] == "builtin" + names = {d["name"] for d in data["compatible_dimensions"]} + assert names == {"region", "category"} + + [email protected] +async def test_get_compatible_dimensions_external_happy_path( + mcp_server: FastMCP, +) -> None: + """External views delegate to view.get_compatible_dimensions().""" + mock_view = _make_view(5) + + with patch( + "superset.daos.semantic_layer.SemanticViewDAO.find_by_id", + return_value=mock_view, + ): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_compatible_dimensions", + {"request": {"view_id": 5, "selected_metrics": ["bookings"]}}, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is True + assert data["source"] == "external" + assert [d["name"] for d in data["compatible_dimensions"]] == ["country_name"] + mock_view.get_compatible_dimensions.assert_called_once_with(["bookings"], []) + + [email protected] +async def test_get_compatible_dimensions_mutual_exclusion_validation( + mcp_server: FastMCP, +) -> None: + """Errors when both dataset_id and view_id are provided.""" + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_compatible_dimensions", + {"request": {"dataset_id": 1, "view_id": 2}}, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is False + assert data["error_type"] == "ValidationError" + + [email protected] +async def test_get_compatible_dimensions_requires_one_source( + mcp_server: FastMCP, +) -> None: + """Errors when neither dataset_id nor view_id is provided.""" + async with Client(mcp_server) as client: + result = await client.call_tool("get_compatible_dimensions", {"request": {}}) + data = json.loads(result.content[0].text) + + assert data["success"] is False + assert data["error_type"] == "ValidationError" + + [email protected] +async def test_get_compatible_dimensions_privacy_check(mcp_server: FastMCP) -> None: + """Errors when the user lacks data-model metadata access.""" + with patch.object( + get_compatible_dimensions_module, + "user_can_view_data_model_metadata", + return_value=False, + ): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_compatible_dimensions", {"request": {"dataset_id": 1}} + ) + data = json.loads(result.content[0].text) + + assert data["success"] is False + assert data["error_type"] == "DataModelMetadataRestricted" + + [email protected] +async def test_get_compatible_dimensions_external_access_denied( + mcp_server: FastMCP, +) -> None: + """Returns AccessDenied when raise_for_access rejects the view.""" + mock_view = _make_view(5) + mock_view.raise_for_access.side_effect = _access_denied_exc() + + with patch( + "superset.daos.semantic_layer.SemanticViewDAO.find_by_id", + return_value=mock_view, + ): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_compatible_dimensions", {"request": {"view_id": 5}} + ) + data = json.loads(result.content[0].text) + + assert data["success"] is False + assert data["error_type"] == "AccessDenied" + + [email protected] +async def test_get_compatible_dimensions_not_found(mcp_server: FastMCP) -> None: + """Returns NotFound when the dataset doesn't exist.""" + with patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=None): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_compatible_dimensions", {"request": {"dataset_id": 999}} + ) + data = json.loads(result.content[0].text) + + assert data["success"] is False + assert data["error_type"] == "NotFound" Review Comment: Fixed — added test_get_compatible_dimensions_external_not_found covering the SemanticViewDAO.find_by_id returning None case, symmetric with the existing builtin-path test. -- 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]
