aminghadersohi commented on code in PR #40961: URL: https://github.com/apache/superset/pull/40961#discussion_r3482394959
########## superset/mcp_service/dashboard/tool/get_dashboard_datasets.py: ########## @@ -0,0 +1,128 @@ +# 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. + +""" +Get dashboard datasets FastMCP tool + +Returns the datasets used by a dashboard's charts, including columns and +metrics. This is the prerequisite context an agent needs before configuring +native filters on a dashboard (e.g. picking filter target columns). +""" + +import logging +from datetime import datetime, timezone + +from fastmcp import Context +from sqlalchemy.orm import subqueryload +from superset_core.mcp.decorators import tool, ToolAnnotations + +from superset.extensions import event_logger +from superset.mcp_service.dashboard.schemas import ( + dashboard_datasets_serializer, + DashboardDatasets, + DashboardError, + GetDashboardDatasetsRequest, +) +from superset.mcp_service.mcp_core import ModelGetInfoCore + +logger = logging.getLogger(__name__) + + +@tool( + tags=["core"], + class_permission_name="Dashboard", + annotations=ToolAnnotations( + title="Get dashboard datasets", + readOnlyHint=True, + destructiveHint=False, + ), +) +async def get_dashboard_datasets( + request: GetDashboardDatasetsRequest, ctx: Context +) -> DashboardDatasets | DashboardError: Review Comment: Good catch — fixed in b5f2d71. Added `@requires_data_model_metadata_access` plus the body-level `user_can_view_data_model_metadata()` gate, mirroring `get_dataset_info`/`list_datasets`. Dashboard-only viewers without dataset-metadata permission now get a structured privacy denial and the tool is hidden from tool search for them. Added a unit test covering the denied path. ########## superset/mcp_service/dashboard/tool/get_dashboard_datasets.py: ########## @@ -0,0 +1,128 @@ +# 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. + +""" +Get dashboard datasets FastMCP tool + +Returns the datasets used by a dashboard's charts, including columns and +metrics. This is the prerequisite context an agent needs before configuring +native filters on a dashboard (e.g. picking filter target columns). +""" + +import logging +from datetime import datetime, timezone + +from fastmcp import Context +from sqlalchemy.orm import subqueryload +from superset_core.mcp.decorators import tool, ToolAnnotations + +from superset.extensions import event_logger +from superset.mcp_service.dashboard.schemas import ( + dashboard_datasets_serializer, + DashboardDatasets, + DashboardError, + GetDashboardDatasetsRequest, +) +from superset.mcp_service.mcp_core import ModelGetInfoCore + +logger = logging.getLogger(__name__) + + +@tool( + tags=["core"], + class_permission_name="Dashboard", + annotations=ToolAnnotations( + title="Get dashboard datasets", + readOnlyHint=True, + destructiveHint=False, + ), +) +async def get_dashboard_datasets( + request: GetDashboardDatasetsRequest, ctx: Context +) -> DashboardDatasets | DashboardError: + """ + List the datasets used by a dashboard's charts, by ID, UUID, or slug. + + Each dataset includes its table name, schema, database connection + (id, name, backend), columns (name, type, is_dttm, verbose_name) and + metrics (name, expression, verbose_name). Use this to understand which + columns and metrics are available before configuring native filters or + analyzing a dashboard's data model. + + Datasets the current user cannot access are excluded from the response + and reported via inaccessible_dataset_count. Column and metric lists are + capped per dataset; when truncated, columns_truncated/metrics_truncated + are set and total counts are reported. + + Example usage: + ```json + { + "identifier": 123 + } + ``` + """ + await ctx.info( + "Retrieving dashboard datasets: identifier=%s" % (request.identifier,) + ) + + try: + from superset.daos.dashboard import DashboardDAO + from superset.models.dashboard import Dashboard + + # Eager load slices to avoid N+1 queries when grouping by datasource. + eager_options = [subqueryload(Dashboard.slices)] Review Comment: Fixed in b5f2d71. Eager-load each slice's dataset relationships (`subqueryload(Dashboard.slices).subqueryload(Slice.table)` → `SqlaTable.columns`/`metrics` + `joinedload(SqlaTable.database)`), matching `get_dataset_info`, so serialization no longer triggers per-dataset lazy loads. ########## superset/mcp_service/dashboard/schemas.py: ########## @@ -1512,3 +1523,225 @@ def dashboard_layout_serializer(dashboard: "Dashboard") -> DashboardLayout: has_layout=bool(position_json_str), ) ) + + +# Per-dataset caps keep responses small enough for LLM context: wide +# datasets can have hundreds of columns, which would dwarf the fields an +# agent actually needs to configure native filters. +MAX_DASHBOARD_DATASET_COLUMNS = 100 +MAX_DASHBOARD_DATASET_METRICS = 50 + + +class DashboardDatasetColumn(BaseModel): + """Lean column representation for dashboard dataset context.""" + + column_name: str = Field(..., description="Column name") + verbose_name: str | None = Field(None, description="Verbose (display) name") + type: str | None = Field(None, description="Column data type") + is_dttm: bool | None = Field(None, description="Is datetime column") + + +class DashboardDatasetMetric(BaseModel): + """Lean metric representation for dashboard dataset context.""" + + metric_name: str = Field(..., description="Saved metric name") + verbose_name: str | None = Field(None, description="Verbose (display) name") + expression: str | None = Field(None, description="SQL expression") + + +class DashboardDatasetDatabaseInfo(BaseModel): + """Database connection summary for a dashboard dataset.""" + + id: int | None = Field(None, description="Database ID") + name: str | None = Field(None, description="Database name") + backend: str | None = Field(None, description="Database backend (engine)") + + +class DashboardDatasetSummary(BaseModel): + """A dataset used by a dashboard's charts, with columns and metrics.""" + + id: int | None = Field(None, description="Dataset ID") + uuid: str | None = Field(None, description="Dataset UUID") Review Comment: Keeping the integer `id`. It's returned alongside `uuid` consistently with every other MCP schema (`DatasetInfo`, `DashboardInfo`), and the integer dataset id is exactly what native-filter targets reference — this tool exists to feed that workflow (the sibling `manage_native_filters` tool keys datasets by integer id). `uuid` is also provided for callers that prefer it. ########## superset/mcp_service/dashboard/schemas.py: ########## @@ -1512,3 +1523,225 @@ def dashboard_layout_serializer(dashboard: "Dashboard") -> DashboardLayout: has_layout=bool(position_json_str), ) ) + + +# Per-dataset caps keep responses small enough for LLM context: wide +# datasets can have hundreds of columns, which would dwarf the fields an +# agent actually needs to configure native filters. +MAX_DASHBOARD_DATASET_COLUMNS = 100 +MAX_DASHBOARD_DATASET_METRICS = 50 + + +class DashboardDatasetColumn(BaseModel): + """Lean column representation for dashboard dataset context.""" + + column_name: str = Field(..., description="Column name") + verbose_name: str | None = Field(None, description="Verbose (display) name") + type: str | None = Field(None, description="Column data type") + is_dttm: bool | None = Field(None, description="Is datetime column") + + +class DashboardDatasetMetric(BaseModel): + """Lean metric representation for dashboard dataset context.""" + + metric_name: str = Field(..., description="Saved metric name") + verbose_name: str | None = Field(None, description="Verbose (display) name") + expression: str | None = Field(None, description="SQL expression") + + +class DashboardDatasetDatabaseInfo(BaseModel): + """Database connection summary for a dashboard dataset.""" + + id: int | None = Field(None, description="Database ID") + name: str | None = Field(None, description="Database name") + backend: str | None = Field(None, description="Database backend (engine)") + + +class DashboardDatasetSummary(BaseModel): + """A dataset used by a dashboard's charts, with columns and metrics.""" + + id: int | None = Field(None, description="Dataset ID") + uuid: str | None = Field(None, description="Dataset UUID") + table_name: str | None = Field(None, description="Table name") + schema_name: str | None = Field(None, description="Schema name") + database: DashboardDatasetDatabaseInfo | None = Field( + None, description="Database the dataset belongs to" + ) + chart_count: int = Field( + 0, description="Number of charts on the dashboard using this dataset" + ) + columns: List[DashboardDatasetColumn] = Field( + default_factory=list, description="Dataset columns" + ) + metrics: List[DashboardDatasetMetric] = Field( + default_factory=list, description="Dataset metrics" + ) + total_column_count: int = Field( + 0, description="Total number of columns on the dataset" + ) + total_metric_count: int = Field( + 0, description="Total number of metrics on the dataset" + ) + columns_truncated: bool = Field( + False, + description=( + "True when the columns list was truncated to keep the response small" + ), + ) + metrics_truncated: bool = Field( + False, + description=( + "True when the metrics list was truncated to keep the response small" + ), + ) + + @model_serializer(mode="wrap") + def _rename_schema_field(self, serializer: Any, info: Any) -> Dict[str, Any]: + """Serialize 'schema_name' as 'schema' to match API conventions.""" + data = serializer(self) + if "schema_name" in data: + data["schema"] = data.pop("schema_name") + return data + + +class DashboardDatasets(BaseModel): + """Response schema for get_dashboard_datasets.""" + + id: int | None = Field(None, description="Dashboard ID") + dashboard_title: str | None = Field(None, description="Dashboard title") + uuid: str | None = Field(None, description="Dashboard UUID") Review Comment: Same rationale as the dataset-id thread: the integer `id` is kept for consistency with `DashboardInfo` and the other MCP schemas, and because callers chain it into id-based dashboard tools. `uuid` is provided alongside. ########## tests/unit_tests/mcp_service/dashboard/tool/test_get_dashboard_datasets.py: ########## @@ -0,0 +1,355 @@ +# 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 MCP get_dashboard_datasets tool.""" + +from unittest.mock import Mock, patch + +import pytest +from fastmcp import Client + +from superset.mcp_service.app import mcp +from superset.mcp_service.utils.sanitization import ( + LLM_CONTEXT_CLOSE_DELIMITER, + LLM_CONTEXT_OPEN_DELIMITER, +) +from superset.utils import json + + +def _wrapped(value: str) -> str: + return f"{LLM_CONTEXT_OPEN_DELIMITER}\n{value}\n{LLM_CONTEXT_CLOSE_DELIMITER}" + + +def _build_column_mock( + name: str, + *, + verbose_name: str | None = None, + type_: str | None = "VARCHAR", + is_dttm: bool = False, +) -> Mock: + column = Mock() + column.column_name = name + column.verbose_name = verbose_name + column.type = type_ + column.is_dttm = is_dttm + return column + + +def _build_metric_mock( + name: str, + *, + verbose_name: str | None = None, + expression: str | None = None, +) -> Mock: + metric = Mock() + metric.metric_name = name + metric.verbose_name = verbose_name + metric.expression = expression + return metric + + +def _build_database_mock( + *, database_id: int = 7, name: str = "examples", backend: str = "postgresql" +) -> Mock: + database = Mock() + database.id = database_id + database.database_name = name + database.backend = backend + return database + + +def _build_datasource_mock( + *, + dataset_id: int, + uuid: str | None = None, + table_name: str = "my_table", + schema: str | None = "public", + database: Mock | None = None, + columns: list[Mock] | None = None, + metrics: list[Mock] | None = None, +) -> Mock: + datasource = Mock() + datasource.id = dataset_id + datasource.uuid = uuid + datasource.table_name = table_name + datasource.schema = schema + datasource.database = database + datasource.columns = columns or [] + datasource.metrics = metrics or [] + return datasource + + +def _build_slice_mock(datasource: Mock) -> Mock: + slc = Mock() + slc.datasource_id = datasource.id + slc.datasource = datasource + return slc Review Comment: Leaving the test helpers/fixtures as-is to match the existing convention in the sibling mcp_service test files (e.g. `test_get_dashboard_layout.py`), where pytest fixtures and test functions are unannotated and helper builders carry only return types. mypy passes on these files. (Applies to the related fixture/inner-function typing nits on this file too.) -- 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]
