codeant-ai-for-open-source[bot] commented on code in PR #41611: URL: https://github.com/apache/superset/pull/41611#discussion_r3533508482
########## superset/mcp_service/semantic_layer/tool/get_table.py: ########## @@ -0,0 +1,514 @@ +# 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_table + +Query a data source (built-in dataset or external semantic view) using +metric and dimension names, returning tabular results. +""" + +import logging +import time +from dataclasses import dataclass, field +from typing import Any + +from fastmcp import Context +from sqlalchemy.exc import SQLAlchemyError +from superset_core.mcp.decorators import tool, ToolAnnotations + +from superset.commands.exceptions import CommandException +from superset.exceptions import OAuth2Error, OAuth2RedirectError, SupersetException +from superset.extensions import event_logger +from superset.mcp_service.chart.schemas import PerformanceMetadata +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 ( + GetTableRequest, + GetTableResponse, + SemanticLayerError, +) +from superset.mcp_service.utils.cache_utils import get_cache_status_from_result +from superset.mcp_service.utils.oauth2_utils import build_oauth2_redirect_message +from superset.mcp_service.utils.query_utils import validate_names +from superset.mcp_service.utils.response_utils import format_data_columns + +logger = logging.getLogger(__name__) + + +@dataclass +class _ResolvedDatasource: + """Metadata needed to validate and query a resolved data source.""" + + display_name: str + time_col: str | None + valid_columns: set[str] + valid_metrics: set[str] + warnings: list[str] = field(default_factory=list) + + +def _time_column_error( + time_col: str, valid_columns: set[str], display_name: str, kind: str +) -> str: + if time_col in valid_columns: + return ( + f"time_column '{time_col}' on {kind} '{display_name}' is " + "not marked as a datetime column." + ) + return f"Unknown time_column: '{time_col}' on {kind} '{display_name}'." + + +def _resolve_builtin_dataset( + request: GetTableRequest, +) -> _ResolvedDatasource | SemanticLayerError: + """Resolve metadata for a built-in dataset.""" + from sqlalchemy.orm import subqueryload + + from superset.connectors.sqla.models import SqlaTable + from superset.daos.dataset import DatasetDAO + + dataset_id = request.dataset_id + assert dataset_id is not None + + with event_logger.log_context(action="mcp.get_table.resolve_dataset"): + 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: {dataset_id}.", + error_type="NotFound", + ) + + display_name = dataset.table_name + valid_columns = {c.column_name for c in dataset.columns} + valid_dttm_columns = {c.column_name for c in dataset.columns if c.is_dttm} + valid_metrics = {m.metric_name for m in dataset.metrics} + + time_col = request.time_column + if time_col is None and request.time_range: + time_col = getattr(dataset, "main_dttm_col", None) + if not time_col: + return SemanticLayerError.create( + error=( + "time_range was provided but no temporal column is " + "configured. Set time_column explicitly." + ), + error_type="ValidationError", + ) + if time_col is not None and time_col not in valid_dttm_columns: + return SemanticLayerError.create( + error=_time_column_error(time_col, valid_columns, display_name, "dataset"), + error_type="ValidationError", + ) + + return _ResolvedDatasource(display_name, time_col, valid_columns, valid_metrics) + + +def _resolve_external_view( + request: GetTableRequest, +) -> _ResolvedDatasource | SemanticLayerError: + """Resolve metadata for an external semantic view.""" + from superset.daos.semantic_layer import SemanticViewDAO + from superset.exceptions import SupersetSecurityException + + view_id = request.view_id + assert view_id is not None + + with event_logger.log_context(action="mcp.get_table.resolve_view"): + 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", + ) + + display_name = view.name + valid_columns = {c.column_name for c in view.columns} + valid_dttm_columns = {c.column_name for c in view.columns if c.is_dttm} + valid_metrics = {m.metric_name for m in view.metrics} + + warnings: list[str] = [] + time_col = request.time_column + if time_col is None and request.time_range: + dttm_cols = [c for c in view.columns if c.is_dttm] + if dttm_cols: + time_col = dttm_cols[0].column_name + else: + return SemanticLayerError.create( + error=( + f"time_range was provided but view '{display_name}' has " + "no datetime dimension. Set time_column explicitly or " + "omit time_range." + ), + error_type="ValidationError", + ) + if time_col is not None and time_col not in valid_dttm_columns: + return SemanticLayerError.create( + error=_time_column_error(time_col, valid_columns, display_name, "view"), + error_type="ValidationError", + ) + + return _ResolvedDatasource( + display_name, time_col, valid_columns, valid_metrics, warnings + ) + + +def _validate_request_names( + request: GetTableRequest, valid_columns: set[str], valid_metrics: set[str] +) -> list[str]: + """Validate requested dimensions, metrics, filters, and order_by names.""" + validation_errors: list[str] = [] + validation_errors.extend( + validate_names(request.dimensions, valid_columns, "dimension") + ) + validation_errors.extend(validate_names(request.metrics, valid_metrics, "metric")) + filter_cols = [f.col for f in request.filters] + validation_errors.extend( + validate_names(filter_cols, valid_columns, "filter column") + ) + if request.order_by: + valid_orderby = valid_columns | valid_metrics + validation_errors.extend( + validate_names(request.order_by, valid_orderby, "order_by") + ) + return validation_errors + + +def _build_query_dict( + request: GetTableRequest, + time_col: str | None, +) -> dict[str, Any]: + """Assemble the query dict for QueryContextFactory.""" + filters: list[dict[str, Any]] = [ + {"col": f.col, "op": f.op, "val": f.val} for f in request.filters + ] + if request.time_range and time_col: + filters.append( + {"col": time_col, "op": "TEMPORAL_RANGE", "val": request.time_range} + ) + + query_dict: dict[str, Any] = { + "filters": filters, + "columns": request.dimensions, + "metrics": request.metrics, + "row_limit": request.row_limit, + "order_desc": request.order_desc, + } + if time_col: + query_dict["granularity"] = time_col + if request.order_by: + query_dict["orderby"] = [ + (col, not request.order_desc) for col in request.order_by + ] + return query_dict + + +def _build_response( + request: GetTableRequest, + is_builtin: bool, + display_name: str, + query_result: dict[str, Any], + query_duration_ms: int, + warnings: list[str], +) -> GetTableResponse: + """Format the query result into a GetTableResponse.""" + data = query_result.get("data", []) + raw_columns = query_result.get("colnames", []) + cache_status = get_cache_status_from_result( + query_result, force_refresh=request.force_refresh + ) + + if not data: + return GetTableResponse( + columns=[], + data=[], + row_count=0, + total_rows=0, + summary=f"'{display_name}': query returned no data.", + source="builtin" if is_builtin else "external", + dataset_id=request.dataset_id, + dataset_name=display_name if is_builtin else None, + view_id=request.view_id, + view_name=display_name if not is_builtin else None, + performance=PerformanceMetadata( + query_duration_ms=query_duration_ms, + cache_status="no_data", + ), + cache_status=cache_status, + warnings=warnings, + ) + + columns_meta = format_data_columns(data, raw_columns) + cache_label = "cached" if cache_status and cache_status.cache_hit else "fresh" + summary = ( + f"'{display_name}': {len(data)} rows, " + f"{len(raw_columns)} columns ({cache_label})." + ) + + return GetTableResponse( + columns=columns_meta, + data=data, + row_count=len(data), + total_rows=query_result.get("rowcount"), + summary=summary, + source="builtin" if is_builtin else "external", + dataset_id=request.dataset_id, + dataset_name=display_name if is_builtin else None, + view_id=request.view_id, + view_name=display_name if not is_builtin else None, + performance=PerformanceMetadata( + query_duration_ms=query_duration_ms, + cache_status=cache_label, + ), + cache_status=cache_status, + warnings=warnings, + ) + + +async def _run_get_table_query( + request: GetTableRequest, + ctx: Context, + is_builtin: bool, + datasource_id: int, + datasource_type: str, +) -> GetTableResponse | SemanticLayerError: + """Resolve, validate, execute, and format a get_table request.""" + from superset.commands.chart.data.get_data_command import ChartDataCommand + from superset.common.query_context_factory import QueryContextFactory + + await ctx.report_progress(1, 5, "Resolving data source") + resolved = ( + _resolve_builtin_dataset(request) + if is_builtin + else _resolve_external_view(request) + ) + if isinstance(resolved, SemanticLayerError): + return resolved + + await ctx.report_progress(2, 5, "Validating metrics and dimensions") + validation_errors = _validate_request_names( + request, resolved.valid_columns, resolved.valid_metrics + ) + if validation_errors: + error_msg = "; ".join(validation_errors) + await ctx.error("Validation failed: %s" % (error_msg,)) + return SemanticLayerError.create( + error=error_msg, + error_type="ValidationError", + ) + + await ctx.report_progress(3, 5, "Building query") + query_dict = _build_query_dict(request, resolved.time_col) + + await ctx.debug("Query dict: %s" % (sorted(query_dict.keys()),)) + await ctx.report_progress(4, 5, "Executing query") + start_time = time.time() + + with event_logger.log_context(action="mcp.get_table.execute"): + factory = QueryContextFactory() + query_context = factory.create( + datasource={"id": datasource_id, "type": datasource_type}, + queries=[query_dict], + form_data={}, + force=not request.use_cache or request.force_refresh, + ) + command = ChartDataCommand(query_context) Review Comment: **Suggestion:** This query path enables shared result caching without adding a user-scoped cache key for OAuth2-backed datasources, so one user can receive another user's cached data. Enforce per-user cache isolation whenever the resolved datasource uses OAuth2 (or disable shared cache in that case) before executing `ChartDataCommand`. [cache] <details> <summary><b>Severity Level:</b> Critical π¨</summary> ```mdx - β get_table leaks cached rows between OAuth2 users. - β MCP semantic-layer queries bypass per-user cache isolation. - β οΈ Underlying OAuth2 DB permissions can be silently bypassed. - β οΈ LLM clients may see another userβs restricted data. ``` </details> <details> <summary><b>Steps of Reproduction β </b></summary> ```mdx 1. Configure an OAuth2-enabled database and map a SqlaTable dataset to it (Database model in `superset/models/core.py:173` is used as the backing DB for SqlaTable; OAuth2 handling for DB connections is wired via `check_for_oauth2` imported at `superset/models/core.py:91` and token handling in `superset/utils/oauth2.py`). 2. Ensure the OAuth2 database does not have `per_user_caching` enabled in its `extra` JSON (field is exposed in the DB configuration schema at `superset/databases/schemas.py:917`, but nothing in the get_table MCP path enforces it on for OAuth2-backed databases). 3. From an MCP client, call the new `get_table` tool defined at `superset/mcp_service/semantic_layer/tool/get_table.py:406` as user A with a `GetTableRequest` that targets a dataset on this OAuth2 DB and leaves `use_cache` as its default (True). The request flows through `get_table()` into `_run_get_table_query()` (lines 296β374), which builds `query_dict` and then constructs a `QueryContext` via `QueryContextFactory.create()` using the snippet at lines 338β343 (datasource `{"id": datasource_id, "type": "table"}`, `queries=[query_dict]`, `form_data={}`, `force=not request.use_cache or request.force_refresh`). 4. ChartDataCommand (`superset/commands/chart/data/get_data_command.py:34`) executes the query via `_query_context.get_payload(...)`, which delegates caching and cache key generation to `QueryContextProcessor.query_cache_key()` (`superset/common/query_context_processor.py:228`) and `QueryObject.cache_key()` (`superset/common/query_object.py:460-48). Because this MCP path does not add any user-specific cache key when `per_user_caching` is off, the cache key is shared for all users hitting the same dataset/query. When user B later calls `get_table` with the same request on the same OAuth2-backed dataset, ChartDataCommand will reuse the cached payload from user A, returning user A's data to user B even though the underlying DB connection uses per-user OAuth2 tokens. This reproduces cross-user cache leakage for OAuth2-backed datasources on the new `get_table` MCP tool. ``` </details> [](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=e37690131b1f497f9be29a6874cd257e&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset) [](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt_id=e37690131b1f497f9be29a6874cd257e&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset) *(Use Cmd/Ctrl + Click for best experience)* <details> <summary><b>Prompt for AI Agent π€ </b></summary> ```mdx This is a comment left during a code review. **Path:** superset/mcp_service/semantic_layer/tool/get_table.py **Line:** 338:343 **Comment:** *Cache: This query path enables shared result caching without adding a user-scoped cache key for OAuth2-backed datasources, so one user can receive another user's cached data. Enforce per-user cache isolation whenever the resolved datasource uses OAuth2 (or disable shared cache in that case) before executing `ChartDataCommand`. Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise. Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix ``` </details> <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F41611&comment_hash=f28161ddf8ff43bda1ae735901ac1e0dc9eb7c1f7903e6cdc01a293caded3068&reaction=like'>π</a> | <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F41611&comment_hash=f28161ddf8ff43bda1ae735901ac1e0dc9eb7c1f7903e6cdc01a293caded3068&reaction=dislike'>π</a> -- 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]
