aminghadersohi commented on code in PR #41611: URL: https://github.com/apache/superset/pull/41611#discussion_r3537919369
########## superset/mcp_service/semantic_layer/tool/get_compatible_metrics.py: ########## @@ -0,0 +1,230 @@ +# 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: get_compatible_metrics + +Returns metrics compatible with the current dimension/metric selection. +""" + +import logging + +from fastmcp import Context +from superset_core.mcp.decorators import tool, ToolAnnotations + +from superset.extensions import 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 ( + CompatibleMetricsResponse, + GetCompatibleMetricsRequest, + MetricInfo, + SemanticLayerError, +) + +logger = logging.getLogger(__name__) + + +@tool( + tags=["data", "semantic"], + class_permission_name="Dataset", + annotations=ToolAnnotations( + title="Get compatible metrics", + readOnlyHint=True, + destructiveHint=False, + ), +) +@requires_data_model_metadata_access +async def get_compatible_metrics( + request: GetCompatibleMetricsRequest, + ctx: Context, +) -> CompatibleMetricsResponse | SemanticLayerError: + """Return metrics compatible with the current dimension/metric selection. + + Used to progressively refine a query: given a set of already-selected + metrics and dimensions, returns the additional metrics that can be + combined without breaking the underlying semantic constraints. + + Provide exactly one of ``dataset_id`` (built-in) or ``view_id`` (external). + + For built-in datasets, all metrics from the dataset are considered + compatible (SQL GROUP BY imposes no metric-level constraints). + + For external semantic views, delegates to the view's + ``get_compatible_metrics`` implementation. + + Example: + ```json + { + "selected_metrics": [], + "selected_dimensions": ["region"], + "view_id": 5 + } + ``` + """ + await ctx.info( + "Getting compatible metrics: dataset_id=%s, view_id=%s, " + "metrics=%s, dims=%s" + % ( + request.dataset_id, + request.view_id, + request.selected_metrics, + request.selected_dimensions, + ) + ) + + if not user_can_view_data_model_metadata(): + return SemanticLayerError.create( + error="You don't have permission to access dataset details for your role.", + error_type=DATA_MODEL_METADATA_ERROR_TYPE, + ) + + if request.dataset_id is None and request.view_id is None: + return SemanticLayerError.create( + error="Provide either dataset_id (built-in) or view_id (external).", + error_type="ValidationError", + ) + if request.dataset_id is not None and request.view_id is not None: + return SemanticLayerError.create( + error="Provide only one of dataset_id or view_id, not both.", + error_type="ValidationError", + ) + + try: + # ------------------------------------------------------------------ + # Built-in dataset path + # ------------------------------------------------------------------ + if request.dataset_id is not None: + from sqlalchemy.orm import subqueryload + + from superset.connectors.sqla.models import SqlaTable + from superset.daos.dataset import DatasetDAO + + with event_logger.log_context(action="mcp.get_compatible_metrics.builtin"): + dataset = DatasetDAO.find_by_id( + request.dataset_id, + query_options=[ + subqueryload(SqlaTable.columns), + subqueryload(SqlaTable.metrics), + ], + ) + + if dataset is None: + return SemanticLayerError.create( + error=f"No dataset found with id: {request.dataset_id}.", + error_type="NotFound", + ) + + # All metrics on a SQL dataset are always mutually compatible; + # exclude ones already selected so clients don't get duplicate + # suggestions for metrics they've already added. + selected_metrics = set(request.selected_metrics) + compatible = [ + MetricInfo( + name=m.metric_name, + verbose_name=m.verbose_name or None, + description=m.description or None, + expression=m.expression or None, + d3format=m.d3format or None, + warning_text=m.warning_text or None, + source="builtin", + dataset_id=dataset.id, + dataset_name=dataset.table_name, + ) + for m in dataset.metrics + if m.metric_name not in selected_metrics + ] + + await ctx.info("Compatible metrics (builtin): count=%d" % len(compatible)) + return CompatibleMetricsResponse( + compatible_metrics=compatible, + source="builtin", + ) + + # ------------------------------------------------------------------ + # External semantic view path + # ------------------------------------------------------------------ + from superset.daos.semantic_layer import SemanticViewDAO + from superset.exceptions import SupersetSecurityException + + view_id: int = request.view_id # type: ignore[assignment] + with event_logger.log_context(action="mcp.get_compatible_metrics.external"): + view = SemanticViewDAO.find_by_id(view_id) + + if view is None: + return SemanticLayerError.create( + error=f"No semantic view found with id: {view_id}.", + error_type="NotFound", + ) + + try: + view.raise_for_access() + except SupersetSecurityException as ex: + return SemanticLayerError.create( + error=str(ex.error.message), + error_type="AccessDenied", + ) + + compatible_names = view.get_compatible_metrics( + request.selected_metrics, + request.selected_dimensions, + ) + + # Enrich with full metric metadata + all_metrics_map = {m.metric_name: m for m in view.metrics} Review Comment: Fixed — annotated as `all_metrics_map: dict[str, MetricMetadata]`. ########## 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 [] Review Comment: Fixed — annotated as `datasets: list[SqlaTable]`. ########## tests/unit_tests/mcp_service/semantic_layer/tool/test_list_metrics.py: ########## @@ -0,0 +1,305 @@ +# 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 list_metrics MCP tool.""" + +from __future__ import annotations + +import importlib +from collections.abc import Generator +from types import ModuleType +from unittest.mock import call, 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 + +list_metrics_module: ModuleType = importlib.import_module( + "superset.mcp_service.semantic_layer.tool.list_metrics" +) + + [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( + list_metrics_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_metric(name: str, expression: str = "COUNT(*)") -> MagicMock: + m = MagicMock() + m.metric_name = name + m.verbose_name = None + m.expression = expression + m.description = None + m.d3format = None + m.warning_text = None + return m + + +def _make_column(name: str) -> MagicMock: + col = MagicMock() + col.column_name = name + col.verbose_name = None + col.description = None + col.type = "VARCHAR" + col.is_dttm = False + col.groupby = True + col.filterable = True + return col + + +def _make_dataset(dataset_id: int = 1) -> MagicMock: + ds = MagicMock() + ds.id = dataset_id + ds.table_name = f"table_{dataset_id}" + ds.metrics = [_make_metric("count"), _make_metric("revenue", "SUM(revenue)")] + ds.columns = [_make_column("region"), _make_column("category")] + 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.metrics = [_make_metric("bookings"), _make_metric("revenue", "SUM(revenue)")] + view.columns = [_make_column("listing__country_name"), _make_column("channel")] + view.get_compatible_dimensions = MagicMock(return_value=["listing__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_list_metrics_builtin_happy_path(mcp_server: FastMCP) -> None: + """list_metrics returns builtin metrics when only datasets exist.""" + mock_ds = _make_dataset(42) + + with ( + patch( + "superset.mcp_service.semantic_layer.tool.list_metrics.DatasetDAO" + ) as mock_dao, + patch( + "superset.mcp_service.semantic_layer.tool.list_metrics.SemanticViewDAO" + ) as mock_view_dao, + ): + mock_dao.find_by_id.return_value = mock_ds + mock_view_dao.find_accessible.return_value = [] + + async with Client(mcp_server) as client: + result = await client.call_tool( + "list_metrics", + {"request": {"dataset_id": 42, "include_compatible_dimensions": False}}, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is True + assert data["total_count"] == 2 + metrics = data["metrics"] + assert {m["name"] for m in metrics} == {"count", "revenue"} + assert all(m["source"] == "builtin" for m in metrics) + assert all(m["dataset_id"] == 42 for m in metrics) + + [email protected] +async def test_list_metrics_mutual_exclusion_validation(mcp_server: FastMCP) -> None: + """list_metrics returns a validation error when dataset_id and view_id coexist.""" + async with Client(mcp_server) as client: + result = await client.call_tool( + "list_metrics", + {"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_list_metrics_privacy_check(mcp_server: FastMCP) -> None: + """list_metrics returns an error when the user lacks data-model metadata access.""" + with patch.object( + list_metrics_module, + "user_can_view_data_model_metadata", + return_value=False, + ): + async with Client(mcp_server) as client: + result = await client.call_tool("list_metrics", {}) + data = json.loads(result.content[0].text) + + assert data["success"] is False + assert data["error_type"] == "DataModelMetadataRestricted" + + [email protected] +async def test_list_metrics_search_filter(mcp_server: FastMCP) -> None: + """list_metrics filters metrics by search term.""" + mock_ds = _make_dataset(1) + + with ( + patch( + "superset.mcp_service.semantic_layer.tool.list_metrics.DatasetDAO" + ) as mock_dao, + patch( + "superset.mcp_service.semantic_layer.tool.list_metrics.SemanticViewDAO" + ) as mock_view_dao, + patch("superset.mcp_service.semantic_layer.tool.list_metrics.db") as mock_db, + ): + mock_view_dao.find_accessible.return_value = [] + mock_query = MagicMock() Review Comment: Fixed — annotated as `mock_query: MagicMock`. ########## tests/unit_tests/mcp_service/semantic_layer/tool/test_list_metrics.py: ########## @@ -0,0 +1,305 @@ +# 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 list_metrics MCP tool.""" + +from __future__ import annotations + +import importlib +from collections.abc import Generator +from types import ModuleType +from unittest.mock import call, 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 + +list_metrics_module: ModuleType = importlib.import_module( + "superset.mcp_service.semantic_layer.tool.list_metrics" +) + + [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( + list_metrics_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_metric(name: str, expression: str = "COUNT(*)") -> MagicMock: + m = MagicMock() + m.metric_name = name + m.verbose_name = None + m.expression = expression + m.description = None + m.d3format = None + m.warning_text = None + return m + + +def _make_column(name: str) -> MagicMock: + col = MagicMock() + col.column_name = name + col.verbose_name = None + col.description = None + col.type = "VARCHAR" + col.is_dttm = False + col.groupby = True + col.filterable = True + return col + + +def _make_dataset(dataset_id: int = 1) -> MagicMock: + ds = MagicMock() + ds.id = dataset_id + ds.table_name = f"table_{dataset_id}" + ds.metrics = [_make_metric("count"), _make_metric("revenue", "SUM(revenue)")] + ds.columns = [_make_column("region"), _make_column("category")] + 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.metrics = [_make_metric("bookings"), _make_metric("revenue", "SUM(revenue)")] + view.columns = [_make_column("listing__country_name"), _make_column("channel")] + view.get_compatible_dimensions = MagicMock(return_value=["listing__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_list_metrics_builtin_happy_path(mcp_server: FastMCP) -> None: + """list_metrics returns builtin metrics when only datasets exist.""" + mock_ds = _make_dataset(42) + + with ( + patch( + "superset.mcp_service.semantic_layer.tool.list_metrics.DatasetDAO" + ) as mock_dao, + patch( + "superset.mcp_service.semantic_layer.tool.list_metrics.SemanticViewDAO" + ) as mock_view_dao, + ): + mock_dao.find_by_id.return_value = mock_ds + mock_view_dao.find_accessible.return_value = [] + + async with Client(mcp_server) as client: + result = await client.call_tool( + "list_metrics", + {"request": {"dataset_id": 42, "include_compatible_dimensions": False}}, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is True + assert data["total_count"] == 2 + metrics = data["metrics"] + assert {m["name"] for m in metrics} == {"count", "revenue"} + assert all(m["source"] == "builtin" for m in metrics) + assert all(m["dataset_id"] == 42 for m in metrics) + + [email protected] +async def test_list_metrics_mutual_exclusion_validation(mcp_server: FastMCP) -> None: + """list_metrics returns a validation error when dataset_id and view_id coexist.""" + async with Client(mcp_server) as client: + result = await client.call_tool( + "list_metrics", + {"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_list_metrics_privacy_check(mcp_server: FastMCP) -> None: + """list_metrics returns an error when the user lacks data-model metadata access.""" + with patch.object( + list_metrics_module, + "user_can_view_data_model_metadata", + return_value=False, + ): + async with Client(mcp_server) as client: + result = await client.call_tool("list_metrics", {}) + data = json.loads(result.content[0].text) + + assert data["success"] is False + assert data["error_type"] == "DataModelMetadataRestricted" + + [email protected] +async def test_list_metrics_search_filter(mcp_server: FastMCP) -> None: + """list_metrics filters metrics by search term.""" + mock_ds = _make_dataset(1) Review Comment: Fixed — annotated as `mock_ds: MagicMock`. ########## superset/mcp_service/semantic_layer/tool/get_compatible_dimensions.py: ########## @@ -0,0 +1,230 @@ +# 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: get_compatible_dimensions + +Returns dimensions compatible with the current metric/dimension selection. +""" + +import logging + +from fastmcp import Context +from superset_core.mcp.decorators import tool, ToolAnnotations + +from superset.extensions import 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 ( + CompatibleDimensionsResponse, + DimensionInfo, + GetCompatibleDimensionsRequest, + SemanticLayerError, +) + +logger = logging.getLogger(__name__) + + +@tool( + tags=["data", "semantic"], + class_permission_name="Dataset", + annotations=ToolAnnotations( + title="Get compatible dimensions", + readOnlyHint=True, + destructiveHint=False, + ), +) +@requires_data_model_metadata_access +async def get_compatible_dimensions( + request: GetCompatibleDimensionsRequest, + ctx: Context, +) -> CompatibleDimensionsResponse | SemanticLayerError: + """Return dimensions compatible with the current metric/dimension selection. + + Used to drive progressive disclosure in query builders: after the user + selects one or more metrics (and optionally some dimensions), this tool + returns the dimensions that can validly be added without breaking the + underlying query. + + Provide exactly one of ``dataset_id`` (built-in) or ``view_id`` (external). + + For built-in datasets, returns all groupby-enabled columns from the dataset. + SQL datasets have no semantic compatibility constraints between metrics and + dimensions, so all groupby columns are always returned regardless of the + selected metrics. + + For external semantic views, delegates to the view's + ``get_compatible_dimensions`` implementation. + + Example: + ```json + { + "selected_metrics": ["revenue"], + "selected_dimensions": [], + "view_id": 5 + } + ``` + """ + await ctx.info( + "Getting compatible dimensions: dataset_id=%s, view_id=%s, " + "metrics=%s, dims=%s" + % ( + request.dataset_id, + request.view_id, + request.selected_metrics, + request.selected_dimensions, + ) + ) + + if not user_can_view_data_model_metadata(): + return SemanticLayerError.create( + error="You don't have permission to access dataset details for your role.", + error_type=DATA_MODEL_METADATA_ERROR_TYPE, + ) + + if request.dataset_id is None and request.view_id is None: + return SemanticLayerError.create( + error="Provide either dataset_id (built-in) or view_id (external).", + error_type="ValidationError", + ) + if request.dataset_id is not None and request.view_id is not None: + return SemanticLayerError.create( + error="Provide only one of dataset_id or view_id, not both.", + error_type="ValidationError", + ) + + try: + # ------------------------------------------------------------------ + # Built-in dataset path + # ------------------------------------------------------------------ + if request.dataset_id is not None: + from sqlalchemy.orm import subqueryload + + from superset.connectors.sqla.models import SqlaTable + from superset.daos.dataset import DatasetDAO + + dataset_id: int = request.dataset_id + with event_logger.log_context( + action="mcp.get_compatible_dimensions.builtin" + ): + dataset = DatasetDAO.find_by_id( + dataset_id, + query_options=[ + subqueryload(SqlaTable.columns), + subqueryload(SqlaTable.metrics), + ], + ) + + if dataset is None: + return SemanticLayerError.create( + error=f"No dataset found with id: {request.dataset_id}.", + error_type="NotFound", + ) + + # For built-in datasets all groupby columns are always compatible; + # there's no per-metric compatibility constraint at the SQL level. + dims: list[DimensionInfo] = [ + 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 + ] + + await ctx.info("Compatible dimensions (builtin): count=%d" % len(dims)) + return CompatibleDimensionsResponse( + compatible_dimensions=dims, + source="builtin", + ) + + # ------------------------------------------------------------------ + # External semantic view path + # ------------------------------------------------------------------ + from superset.daos.semantic_layer import SemanticViewDAO + from superset.exceptions import SupersetSecurityException + from superset.semantic_layers.models import ColumnMetadata + + view_id: int = request.view_id # type: ignore[assignment] + with event_logger.log_context(action="mcp.get_compatible_dimensions.external"): + view = SemanticViewDAO.find_by_id(view_id) + + if view is None: + return SemanticLayerError.create( + error=f"No semantic view found with id: {view_id}.", + error_type="NotFound", + ) + + try: + view.raise_for_access() + except SupersetSecurityException as ex: + return SemanticLayerError.create( + error=str(ex.error.message), + error_type="AccessDenied", + ) + + compatible_names: list[str] = view.get_compatible_dimensions( + request.selected_metrics, + request.selected_dimensions, + ) + + # Enrich with full column metadata + all_cols: dict[str, ColumnMetadata] = { + col.column_name: col for col in view.columns + } + dims = [ Review Comment: Attempted, but mypy rejects it: the builtin-path `dims` a few lines above is already annotated as `list[DimensionInfo]` in the same function scope, and mypy's `no-redef` check flags re-annotating a same-named variable even across mutually-exclusive early-return branches. Left this occurrence implicitly typed (mypy still infers `list[DimensionInfo]` correctly from the list comprehension). ########## superset/mcp_service/semantic_layer/tool/get_compatible_dimensions.py: ########## @@ -0,0 +1,230 @@ +# 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: get_compatible_dimensions + +Returns dimensions compatible with the current metric/dimension selection. +""" + +import logging + +from fastmcp import Context +from superset_core.mcp.decorators import tool, ToolAnnotations + +from superset.extensions import 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 ( + CompatibleDimensionsResponse, + DimensionInfo, + GetCompatibleDimensionsRequest, + SemanticLayerError, +) + +logger = logging.getLogger(__name__) + + +@tool( + tags=["data", "semantic"], + class_permission_name="Dataset", + annotations=ToolAnnotations( + title="Get compatible dimensions", + readOnlyHint=True, + destructiveHint=False, + ), +) +@requires_data_model_metadata_access +async def get_compatible_dimensions( + request: GetCompatibleDimensionsRequest, + ctx: Context, +) -> CompatibleDimensionsResponse | SemanticLayerError: + """Return dimensions compatible with the current metric/dimension selection. + + Used to drive progressive disclosure in query builders: after the user + selects one or more metrics (and optionally some dimensions), this tool + returns the dimensions that can validly be added without breaking the + underlying query. + + Provide exactly one of ``dataset_id`` (built-in) or ``view_id`` (external). + + For built-in datasets, returns all groupby-enabled columns from the dataset. + SQL datasets have no semantic compatibility constraints between metrics and + dimensions, so all groupby columns are always returned regardless of the + selected metrics. + + For external semantic views, delegates to the view's + ``get_compatible_dimensions`` implementation. + + Example: + ```json + { + "selected_metrics": ["revenue"], + "selected_dimensions": [], + "view_id": 5 + } + ``` + """ + await ctx.info( + "Getting compatible dimensions: dataset_id=%s, view_id=%s, " + "metrics=%s, dims=%s" + % ( + request.dataset_id, + request.view_id, + request.selected_metrics, + request.selected_dimensions, + ) + ) + + if not user_can_view_data_model_metadata(): + return SemanticLayerError.create( + error="You don't have permission to access dataset details for your role.", + error_type=DATA_MODEL_METADATA_ERROR_TYPE, + ) + + if request.dataset_id is None and request.view_id is None: + return SemanticLayerError.create( + error="Provide either dataset_id (built-in) or view_id (external).", + error_type="ValidationError", + ) + if request.dataset_id is not None and request.view_id is not None: + return SemanticLayerError.create( + error="Provide only one of dataset_id or view_id, not both.", + error_type="ValidationError", + ) + + try: + # ------------------------------------------------------------------ + # Built-in dataset path + # ------------------------------------------------------------------ + if request.dataset_id is not None: + from sqlalchemy.orm import subqueryload + + from superset.connectors.sqla.models import SqlaTable + from superset.daos.dataset import DatasetDAO + + dataset_id: int = request.dataset_id + with event_logger.log_context( + action="mcp.get_compatible_dimensions.builtin" + ): + dataset = DatasetDAO.find_by_id( + dataset_id, + query_options=[ + subqueryload(SqlaTable.columns), + subqueryload(SqlaTable.metrics), + ], + ) Review Comment: Fixed — annotated as `dataset: SqlaTable | None`. ########## superset/mcp_service/semantic_layer/tool/get_compatible_dimensions.py: ########## @@ -0,0 +1,230 @@ +# 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: get_compatible_dimensions + +Returns dimensions compatible with the current metric/dimension selection. +""" + +import logging + +from fastmcp import Context +from superset_core.mcp.decorators import tool, ToolAnnotations + +from superset.extensions import 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 ( + CompatibleDimensionsResponse, + DimensionInfo, + GetCompatibleDimensionsRequest, + SemanticLayerError, +) + +logger = logging.getLogger(__name__) + + +@tool( + tags=["data", "semantic"], + class_permission_name="Dataset", + annotations=ToolAnnotations( + title="Get compatible dimensions", + readOnlyHint=True, + destructiveHint=False, + ), +) +@requires_data_model_metadata_access +async def get_compatible_dimensions( + request: GetCompatibleDimensionsRequest, + ctx: Context, +) -> CompatibleDimensionsResponse | SemanticLayerError: + """Return dimensions compatible with the current metric/dimension selection. + + Used to drive progressive disclosure in query builders: after the user + selects one or more metrics (and optionally some dimensions), this tool + returns the dimensions that can validly be added without breaking the + underlying query. + + Provide exactly one of ``dataset_id`` (built-in) or ``view_id`` (external). + + For built-in datasets, returns all groupby-enabled columns from the dataset. + SQL datasets have no semantic compatibility constraints between metrics and + dimensions, so all groupby columns are always returned regardless of the + selected metrics. + + For external semantic views, delegates to the view's + ``get_compatible_dimensions`` implementation. + + Example: + ```json + { + "selected_metrics": ["revenue"], + "selected_dimensions": [], + "view_id": 5 + } + ``` + """ + await ctx.info( + "Getting compatible dimensions: dataset_id=%s, view_id=%s, " + "metrics=%s, dims=%s" + % ( + request.dataset_id, + request.view_id, + request.selected_metrics, + request.selected_dimensions, + ) + ) + + if not user_can_view_data_model_metadata(): + return SemanticLayerError.create( + error="You don't have permission to access dataset details for your role.", + error_type=DATA_MODEL_METADATA_ERROR_TYPE, + ) + + if request.dataset_id is None and request.view_id is None: + return SemanticLayerError.create( + error="Provide either dataset_id (built-in) or view_id (external).", + error_type="ValidationError", + ) + if request.dataset_id is not None and request.view_id is not None: + return SemanticLayerError.create( + error="Provide only one of dataset_id or view_id, not both.", + error_type="ValidationError", + ) + + try: + # ------------------------------------------------------------------ + # Built-in dataset path + # ------------------------------------------------------------------ + if request.dataset_id is not None: + from sqlalchemy.orm import subqueryload + + from superset.connectors.sqla.models import SqlaTable + from superset.daos.dataset import DatasetDAO + + dataset_id: int = request.dataset_id + with event_logger.log_context( + action="mcp.get_compatible_dimensions.builtin" + ): + dataset = DatasetDAO.find_by_id( + dataset_id, + query_options=[ + subqueryload(SqlaTable.columns), + subqueryload(SqlaTable.metrics), + ], + ) + + if dataset is None: + return SemanticLayerError.create( + error=f"No dataset found with id: {request.dataset_id}.", + error_type="NotFound", + ) + + # For built-in datasets all groupby columns are always compatible; + # there's no per-metric compatibility constraint at the SQL level. + dims: list[DimensionInfo] = [ + 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 + ] + + await ctx.info("Compatible dimensions (builtin): count=%d" % len(dims)) + return CompatibleDimensionsResponse( + compatible_dimensions=dims, + source="builtin", + ) + + # ------------------------------------------------------------------ + # External semantic view path + # ------------------------------------------------------------------ + from superset.daos.semantic_layer import SemanticViewDAO + from superset.exceptions import SupersetSecurityException + from superset.semantic_layers.models import ColumnMetadata + + view_id: int = request.view_id # type: ignore[assignment] + with event_logger.log_context(action="mcp.get_compatible_dimensions.external"): + view = SemanticViewDAO.find_by_id(view_id) Review Comment: Fixed — annotated as `view: SemanticView | None`. ########## tests/unit_tests/mcp_service/semantic_layer/tool/test_get_compatible_metrics.py: ########## @@ -0,0 +1,247 @@ +# 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_metrics 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_metrics_module: ModuleType = importlib.import_module( + "superset.mcp_service.semantic_layer.tool.get_compatible_metrics" +) + + [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_metrics_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_metric(name: str, expression: str = "COUNT(*)") -> MagicMock: + m = MagicMock() + m.metric_name = name + m.verbose_name = None + m.expression = expression + m.description = None + m.d3format = None + m.warning_text = None + return m + + +def _make_dataset(dataset_id: int = 42) -> MagicMock: + ds = MagicMock() Review Comment: Fixed — annotated as `ds: MagicMock`. ########## tests/unit_tests/mcp_service/semantic_layer/tool/test_get_compatible_metrics.py: ########## @@ -0,0 +1,247 @@ +# 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_metrics 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_metrics_module: ModuleType = importlib.import_module( + "superset.mcp_service.semantic_layer.tool.get_compatible_metrics" +) + + [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_metrics_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_metric(name: str, expression: str = "COUNT(*)") -> MagicMock: + m = MagicMock() Review Comment: Fixed — annotated as `m: MagicMock`. -- 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]
