codeant-ai-for-open-source[bot] commented on code in PR #44337: URL: https://github.com/apache/superset/pull/44337#discussion_r4023726798
########## superset/mcp_service/dataset/tool/restore_dataset.py: ########## @@ -0,0 +1,196 @@ +# 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: restore_dataset +""" + +import logging +from typing import Any + +from fastmcp import Context +from sqlalchemy.exc import SQLAlchemyError +from superset_core.mcp.decorators import tool, ToolAnnotations + +from superset.commands.dataset.exceptions import ( + DatasetForbiddenError, + DatasetLogicalDuplicateError, + DatasetNotFoundError, +) +from superset.commands.exceptions import CommandException +from superset.extensions import event_logger +from superset.mcp_service.dataset.schemas import ( + RestoreDatasetRequest, + RestoreDatasetResponse, +) + +logger = logging.getLogger(__name__) + + +def _find_dataset_for_restore(identifier: int | str) -> Any | None: + """Resolve a dataset by numeric ID or UUID, including soft-deleted rows. + + Both bypasses mirror ``BaseRestoreCommand.validate``'s own lookup: + ``skip_visibility_filter`` unhides the soft-deleted row, and + ``skip_base_filter`` keeps an editor's own trash reachable even when the + dataset's datasource-access base_filter would hide it (a lost grant must + not hide a row from the one audience that can restore it). The restore + audience is enforced by ``RestoreDatasetCommand`` via + ``raise_for_editorship``. + """ + from superset.daos.dataset import DatasetDAO + + return DatasetDAO.find_by_id_or_uuid( + str(identifier), + skip_base_filter=True, + skip_visibility_filter=True, + ) + + +def _rollback() -> None: + from superset import db + + try: + db.session.rollback() # pylint: disable=consider-using-transaction + except SQLAlchemyError: + logger.warning("Database rollback failed during restore_dataset error handling") + + +@tool( + tags=["mutate"], + class_permission_name="Dataset", + annotations=ToolAnnotations( + title="Restore dataset", + readOnlyHint=False, + destructiveHint=False, + idempotentHint=False, + openWorldHint=False, + ), +) +async def restore_dataset( + request: RestoreDatasetRequest, ctx: Context +) -> RestoreDatasetResponse: + """Restore a soft-deleted dataset from trash. + + Identify the dataset by numeric ID or UUID string (NOT table name). Only + datasets that were soft-deleted (moved to trash while the ``SOFT_DELETE`` + feature flag was enabled) can be restored; permanently deleted datasets + are unrecoverable. The caller must be an editor of the dataset (owners + and Admins qualify). Use list_datasets with deleted_state='only' to find + trashed datasets. + + Restoring fails with ``error_type`` ``LogicalDuplicate`` when another + active dataset already points at the same physical table; that dataset + must be deleted or renamed first. + + Example: + ```json + {"identifier": 123} + ``` + + Returns success with the restored dataset's id/name, or an error. When the + caller lacks permission, ``permission_denied`` is true โ do not retry; ask + the user. + """ + await ctx.info("Restoring dataset: identifier=%s" % (request.identifier,)) + + try: + dataset = _find_dataset_for_restore(request.identifier) + except SQLAlchemyError: + _rollback() + logger.exception("Dataset lookup failed during restore_dataset") + return RestoreDatasetResponse( + success=False, + error="Dataset lookup failed due to a database error.", + error_type="LookupFailed", + ) + if not dataset: + display_id = str(request.identifier)[:200] + msg = f"No dataset found with identifier: {display_id}." + return RestoreDatasetResponse(success=False, error=msg, error_type="NotFound") + + dataset_id = dataset.id + # Table names are user-controlled and must remain exact in response text. + dataset_name = dataset.table_name + + if dataset.deleted_at is None: + return RestoreDatasetResponse( + success=False, + error=( + f"Dataset '{dataset_name}' (id={dataset_id}) is not in trash; " + "nothing to restore." + ), + error_type="NotDeleted", + ) + + # The try/except sits inside log_context so failed restore attempts are + # recorded in the audit log too โ the context manager does not log when + # an exception propagates through it. + with event_logger.log_context(action="mcp.restore_dataset"): + try: + from superset.commands.dataset.restore import RestoreDatasetCommand + + RestoreDatasetCommand(str(dataset.uuid)).run() + + return RestoreDatasetResponse( + success=True, + restored_id=dataset_id, + restored_name=dataset_name, + message=( + f"Restored dataset '{dataset_name}' (id={dataset_id}) from trash." + ), + ) + except DatasetForbiddenError: + await ctx.warning( + "Permission denied restoring dataset id=%s" % (dataset_id,) + ) + return RestoreDatasetResponse( + success=False, + permission_denied=True, + error=( + "You do not have permission to restore dataset " + f"'{dataset_name}' (id={dataset_id}). Ask the user to restore " + "it or grant access; do not retry." + ), + error_type="Forbidden", + ) + except DatasetLogicalDuplicateError as ex: Review Comment: **Suggestion:** The lookup bypasses visibility filters, then this forbidden response includes the dataset name and ID, allowing unauthorized callers to confirm hidden datasets and disclose table names. [security] **Assessment:** ๐ `Major` ยท ๐ `Occurrence: Sometimes` [](https://docs.codeant.ai/cli/resolve-pr-comments-skill) [](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=daa11480896b413293d390b7d2125ac7&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=daa11480896b413293d390b7d2125ac7&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset) <details> <summary><b>Prompt for AI Agent ๐ค </b></summary> ```mdx This is a comment left during a code review. **Path:** superset/mcp_service/dataset/tool/restore_dataset.py **Line:** 161:171 **Comment:** *Security: The lookup bypasses visibility filters, then this forbidden response includes the dataset name and ID, allowing unauthorized callers to confirm hidden datasets and disclose table names. 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%2F44337&comment_hash=4ba715f72aa8c6b65b93057dbb8a9acae21d93dbb1ad80e8ab5c1ccf7678acff&reaction=like'>๐</a> | <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F44337&comment_hash=4ba715f72aa8c6b65b93057dbb8a9acae21d93dbb1ad80e8ab5c1ccf7678acff&reaction=dislike'>๐</a> ########## superset/mcp_service/dataset/tool/update_dataset.py: ########## @@ -0,0 +1,289 @@ +# 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: update_dataset +""" + +import logging +from typing import Any + +from fastmcp import Context +from sqlalchemy.exc import SQLAlchemyError +from superset_core.mcp.decorators import tool, ToolAnnotations + +from superset import security_manager +from superset.exceptions import SupersetException +from superset.extensions import event_logger +from superset.mcp_service.dataset.schemas import ( + UpdateDatasetRequest, + UpdateDatasetResponse, +) + +logger = logging.getLogger(__name__) + + +def _column_names(dataset: Any) -> set[str]: + return {column.column_name for column in dataset.columns} + + +def _sync_error_message(ex: Exception) -> str: + # Raw SQLAlchemy text can leak SQL or connection details; Superset + # exception messages are user-facing by design. + if isinstance(ex, SQLAlchemyError): + return "a database error occurred" + return str(ex) + + +@tool( + tags=["mutate"], + class_permission_name="Dataset", + method_permission_name="write", + annotations=ToolAnnotations( + title="Update dataset", + readOnlyHint=False, + # Rewriting a virtual dataset's SQL or re-syncing its columns changes + # what every chart built on it queries โ non-additive, like + # update_chart. + destructiveHint=True, + idempotentHint=False, + openWorldHint=False, + ), +) +async def update_dataset( # noqa: C901 + request: UpdateDatasetRequest, ctx: Context +) -> UpdateDatasetResponse: + """Update a dataset's name, SQL, description, default datetime column or + cache timeout, and optionally re-sync its columns from the data source. + + Only the properties you pass are changed. ``sql`` applies to virtual + datasets only. When ``sql`` changes, columns are re-synced from the new + query (like "Sync columns from source" in the dataset editor) unless + ``sync_columns`` is false; pass ``sync_columns=true`` on its own to pick + up schema changes in the underlying table or query. Calculated columns + and saved metrics are kept. Use update_dataset_metric to edit metrics. + Requires ownership of the dataset (or Admin). + + Check ``removed_columns`` in the response: charts that use those columns + fail until they are updated. ``warnings`` reports problems that did not + undo the update, e.g. a column re-sync that failed after the SQL was saved. + + Workflow: + 1. Call get_dataset_info to inspect the dataset + 2. Call this tool with the dataset ID and only the properties to change + + Example usage: + ```json + { + "dataset_id": 123, + "sql": "SELECT region, SUM(revenue) AS revenue FROM sales GROUP BY region", + "description": "Revenue by region" + } + ``` + """ + updates = request.updates() + await ctx.info( + "Updating dataset: dataset_id=%s, properties=%s, sync_columns=%s" + % (request.dataset_id, sorted(updates), request.sync_columns) + ) + + try: + from sqlalchemy.orm import joinedload, subqueryload + + from superset.commands.dataset.exceptions import ( + DatasetForbiddenError, + DatasetInvalidError, + DatasetNotFoundError, + DatasetUpdateFailedError, + ) + from superset.commands.dataset.refresh import RefreshDatasetCommand + from superset.commands.dataset.update import UpdateDatasetCommand + from superset.connectors.sqla.models import SqlaTable + from superset.exceptions import SupersetSecurityException + from superset.mcp_service.dataset.dataset_utils import resolve_dataset + from superset.mcp_service.utils.url_utils import get_superset_base_url + + eager_options = [ + subqueryload(SqlaTable.columns), + joinedload(SqlaTable.database), + ] + + with event_logger.log_context(action="mcp.update_dataset.lookup"): + dataset = resolve_dataset(request.dataset_id, eager_options) + + if dataset is None: + display_id = str(request.dataset_id)[:200] + await ctx.warning("Dataset not found: %s" % (display_id,)) + return UpdateDatasetResponse( + error=( + f"No dataset found with identifier: {display_id}." + " Use list_datasets to get valid dataset IDs." + ), + ) + + dataset_id = dataset.id + + # Enforce editorship before validating against the dataset's columns, + # so a caller without edit rights learns nothing beyond "forbidden". + # UpdateDatasetCommand and RefreshDatasetCommand re-check this. + try: + security_manager.raise_for_editorship(dataset) + except SupersetSecurityException: + await ctx.warning("Dataset update forbidden: dataset_id=%s" % (dataset_id,)) + return UpdateDatasetResponse( + dataset_id=dataset_id, + permission_denied=True, + error="You must be an owner of this dataset (or an Admin) to " + "update it. Ask the user to update it or grant access; do not " + "retry.", + ) + + if "sql" in updates and not dataset.sql: + return UpdateDatasetResponse( + dataset_id=dataset_id, + error="sql can only be set on a virtual dataset; this dataset " + "is a physical table.", + ) + + sync_columns = ( + request.sync_columns + if request.sync_columns is not None + else "sql" in updates and updates["sql"] != dataset.sql + ) + + # A new default datetime column is checked against the columns the + # dataset will have once the update is done: the current ones, or the + # re-synced ones, in which case it is applied after the sync. + pending_dttm_col = None + if updates.get("main_dttm_col") is not None: + if sync_columns: + pending_dttm_col = updates.pop("main_dttm_col") + elif updates["main_dttm_col"] not in _column_names(dataset): + dttm_col = updates["main_dttm_col"] + return UpdateDatasetResponse( + dataset_id=dataset_id, + error=f"main_dttm_col '{dttm_col}' is not a column of this " + "dataset. Use get_dataset_info to list its columns.", + ) + + columns_before = _column_names(dataset) + updated_properties = sorted(updates) + + if updates: + # Same pair of commands as PUT /api/v1/dataset/<pk>?override_columns= + # โ the update commits before the column refresh runs. + with event_logger.log_context(action="mcp.update_dataset.update"): + dataset = UpdateDatasetCommand( + dataset_id, updates, override_columns=sync_columns + ).run() + + warnings: list[str] = [] + added_columns: list[str] = [] + removed_columns: list[str] = [] + columns_synced = False + if sync_columns: + try: + with event_logger.log_context(action="mcp.update_dataset.sync_columns"): + dataset = RefreshDatasetCommand(dataset_id).run() + columns_after = _column_names(dataset) + added_columns = sorted(columns_after - columns_before) + removed_columns = sorted(columns_before - columns_after) + columns_synced = True + except (SupersetException, SQLAlchemyError) as ex: + await ctx.warning( + "Dataset column sync failed: %s: %s" % (type(ex).__name__, ex) + ) + warnings.append( + "The update was saved, but re-syncing columns failed " + f"({_sync_error_message(ex)}). The column list may not " + "match the dataset's SQL; retry with sync_columns=true." + ) + + if pending_dttm_col is not None: + if not columns_synced: + warnings.append( + f"main_dttm_col was not changed to '{pending_dttm_col}' " + "because the columns could not be re-synced." + ) + elif pending_dttm_col not in _column_names(dataset): + warnings.append( + f"main_dttm_col was not changed: '{pending_dttm_col}' is not " + "a column of the dataset after the update." + ) + else: + with event_logger.log_context(action="mcp.update_dataset.update"): + dataset = UpdateDatasetCommand( + dataset_id, {"main_dttm_col": pending_dttm_col} + ).run() Review Comment: **Suggestion:** This second update runs after the SQL update and column refresh have committed; if it fails, the tool reports failure although the dataset is already partially changed. [api mismatch] **Assessment:** ๐ `Major` ยท ๐ `Occurrence: Sometimes` [](https://docs.codeant.ai/cli/resolve-pr-comments-skill) [](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=b194d473cb3d4d6aadb3a4a20a6bb346&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=b194d473cb3d4d6aadb3a4a20a6bb346&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset) <details> <summary><b>Prompt for AI Agent ๐ค </b></summary> ```mdx This is a comment left during a code review. **Path:** superset/mcp_service/dataset/tool/update_dataset.py **Line:** 229:231 **Comment:** *Api Mismatch: This second update runs after the SQL update and column refresh have committed; if it fails, the tool reports failure although the dataset is already partially changed. 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%2F44337&comment_hash=1c987ddfc62c4b321e428f3b5684c0867946b468c87cdf4f435e8f15d19f27d8&reaction=like'>๐</a> | <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F44337&comment_hash=1c987ddfc62c4b321e428f3b5684c0867946b468c87cdf4f435e8f15d19f27d8&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]
