Copilot commented on code in PR #44337: URL: https://github.com/apache/superset/pull/44337#discussion_r4023802389
########## 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: Review Comment: Because `_find_dataset_for_restore` deliberately bypasses the DAO base filter, this path can reach a dataset that the caller can neither read nor edit. It reads `table_name` and returns `NotDeleted` before any `raise_for_editorship`, and a trashed non-editor receives the name in the `DatasetForbiddenError` response. A custom role with the MCP `Dataset` write permission but no object editorship can therefore enumerate dataset existence/names outside the `SECURITY.md` Read data/own-dataset capability; mirror `restore_chart`/`restore_dashboard`'s pre-response editorship plus visibility check and keep unauthorized rows indistinguishable from `NotFound`. ########## 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 Review Comment: `RefreshDatasetCommand.run()` deliberately catches `SupersetVirtualTableParseException` and returns the unchanged model, so a Jinja/template parse failure does not raise into this block. The code then compares the old column set and sets `columns_synced=True` without a warning, even though no re-sync occurred (and a pending `main_dttm_col` can be applied against stale columns). Expose the skipped-refresh status or propagate that exception here so the response matches the documented warning semantics. ########## tests/unit_tests/mcp_service/dataset/tool/test_dataset_tools.py: ########## @@ -97,6 +97,7 @@ def create_mock_dataset( dataset.template_params = {} dataset.extra = {} dataset.uuid = f"test-dataset-uuid-{dataset_id}" + dataset.deleted_at = None Review Comment: This only updates the shared fixture. The same file still has direct `MagicMock` datasets in the get-info/list tests (for example lines 1029, 1170, 1261, and 1340) that are passed through `serialize_dataset_object` without setting `deleted_at`; `getattr` returns a child `MagicMock`, which Pydantic cannot validate as `str | datetime | None`. Add `deleted_at = None` to each direct fixture, otherwise `pytest tests/unit_tests/mcp_service/dataset` will fail. ########## 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." Review Comment: When `sql` and `main_dttm_col` are supplied together, a target column that is absent after the refresh takes this warning branch and the tool still returns a successful update. That contradicts the stated contract that setting a nonexistent `main_dttm_col` is rejected, and the caller may believe the requested property was applied even though it was not; validate this deferred value as a failure (or redesign the transaction so it can be rejected before saving the SQL). -- 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]
