aminghadersohi commented on code in PR #44146: URL: https://github.com/apache/superset/pull/44146#discussion_r4043353978
########## superset/mcp_service/dataset_scope.py: ########## @@ -0,0 +1,223 @@ +# 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. + +"""Optional MCP dataset routing constraints, independent of authorization.""" + +import inspect +from collections.abc import Mapping +from typing import Any, TYPE_CHECKING +from uuid import UUID + +from fastmcp.exceptions import ToolError +from flask import current_app, g, has_app_context +from flask_appbuilder.models.filters import BaseFilter + +from superset import security_manager + +if TYPE_CHECKING: + from sqlalchemy.orm.query import Query + +CONFIG_KEY = "MCP_DATASET_ROLE_ALLOWLIST" + +_NO_SUBSTITUTE = ( + "No query was run. Do not substitute another dataset; explain the scope " + "limitation and ask an administrator to review the routing configuration." +) + +# Distinct refusals so a misrouted request is distinguishable from a tool that +# scoped mode does not support at all — both for callers and for regression +# tests, which would otherwise pass against an implementation that refuses +# unconditionally. +UNSUPPORTED_TOOL_ERROR = ( + "This tool is unavailable while MCP is running with a configured dataset " + f"scope, because its results cannot be attributed to a specific registered " + f"dataset. {_NO_SUBSTITUTE}" +) +NO_DATASET_IDENTITY_ERROR = ( + "This request does not identify a registered dataset, which MCP requires " + f"while running with a configured dataset scope. {_NO_SUBSTITUTE}" +) +OUT_OF_SCOPE_ERROR = ( + "The requested dataset is outside the configured MCP dataset scope. " + f"{_NO_SUBSTITUTE}" +) + +# Only these tools can operate in dataset-scoped mode. Other paths can read data +# through SQL, cached results, screenshots, or external semantic sources without +# a registered dataset identity. Refuse them rather than guess at their lineage. +# +# This deliberately gates execution only, not ``tools/list`` visibility. Tool +# listings are assembled by the tool-search transform, which synthesizes its own +# meta tools; filtering that listing on this set would hide the very tools a +# client needs to reach the scoped ones. A refusal the model can read is a +# better failure than a tool surface that silently disappears. +SCOPED_TOOLS = frozenset( + { + "health_check", + "get_schema", + "list_datasets", + "get_dataset_info", + "query_dataset", + "get_table", + # The metric/dimension discovery tools get_table's own documented + # workflow starts with. They name a dataset, so they can be scoped, and + # refusing them would block the allowed tool they lead into. + "list_metrics", + "get_compatible_dimensions", + "get_compatible_metrics", + } +) + +# Tools whose request names a single dataset, mapped to the field that names it. +DATASET_IDENTIFIER_FIELDS = { + "get_dataset_info": "identifier", + "query_dataset": "dataset_id", + "get_table": "dataset_id", + "list_metrics": "dataset_id", + "get_compatible_dimensions": "dataset_id", + "get_compatible_metrics": "dataset_id", +} + + +class MCPDatasetScopeError(ToolError): + """Raised when a call is refused by, or the config of, the routing allowlist. + + Subclasses ``ToolError`` so the explanation reaches the caller verbatim + instead of being flattened into a generic internal error. Surfacing a + configuration problem rather than quietly ignoring the setting keeps a typo + from silently restoring the unrestricted tool surface an operator opted out + of. + """ + + +class DatasetScopeFilter(BaseFilter): # pylint: disable=too-few-public-methods + """Restrict a dataset query to the caller's routing allowlist. + + Applied as a custom filter rather than a caller-visible column operator so + the resolved allowlist — which may name datasets the caller cannot access — + is not echoed back in the ``filters_applied`` section of a list response. + """ + + name = "MCP dataset scope" + arg_name = "mcp_dataset_scope" + + def apply(self, query: "Query", value: frozenset[UUID]) -> "Query": + from superset.connectors.sqla.models import SqlaTable + + return query.filter(SqlaTable.uuid.in_(value)) + + +def parse_dataset_role_allowlist(config: Any) -> dict[str, set[UUID]] | None: + """Validate the allowlist mapping; None means routing constraints are off. + + Split out from scope resolution so a deployment can fail at startup on a + malformed mapping rather than on every subsequent tool call. + """ + if config is None: + return None + if not isinstance(config, dict): + raise MCPDatasetScopeError( + f"{CONFIG_KEY} must map role names to lists of dataset UUIDs." + ) + normalized: dict[str, set[UUID]] = {} + try: + for role, identifiers in config.items(): + if not isinstance(role, str) or not isinstance(identifiers, (list, tuple)): + raise ValueError("Expected role names and UUID lists") + normalized[role] = {UUID(str(identifier)) for identifier in identifiers} Review Comment: Partly acted on, with one correction to the premise. The specific examples don't hold: `UUID("True")` raises `ValueError` ("badly formed hexadecimal UUID string"), as do `UUID(str(True))`, `UUID(str(123))` and `UUID(str(None))`. All of those were already caught and re-raised as a startup `MCPDatasetScopeError`, so they failed loudly rather than loading a "valid-looking UUID". A `uuid.UUID` object round-tripping through `str()` is correct behavior, not a defect. The underlying suggestion — validate entry types before conversion, and fail with a precise message — is still worth doing, so a0f7522 does that: - `_coerce_dataset_uuid` accepts only `UUID` and `str`; anything else (int dataset IDs, bytes, nested lists) is rejected by type with a message saying names and numeric IDs are not accepted, instead of being funneled through `str()`. - Non-string role keys are rejected separately — they can never match a role name. - Errors now name the offending role and entry (`MCP_DATASET_ROLE_ALLOWLIST['Finance'] contains 'not-a-uuid' ...`) rather than the previous generic "contains an entry that is not a list of dataset UUIDs". Validation still runs at startup via `flask_singleton.py`, so a typo fails at boot, not per tool call. Covered by `test_startup_validation_rejects_non_string_entries`, `test_startup_validation_names_the_offending_role_and_entry`, and `test_startup_validation_rejects_non_string_role_keys`. -- 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]
