feat(mcp): add save_sql_query tool for SQL Lab saved queries (#38414)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

superset/mcp_service/sql_lab/tool/save_sql_query.py

@@ -0,0 +1,137 @@
+# 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.
+
+"""
+Save SQL Query MCP Tool
+
+Tool for saving a SQL query as a named SavedQuery in Superset,
+so it appears in SQL Lab's "Saved Queries" list and can be
+reloaded/shared via URL.
+"""
+
+from __future__ import annotations
+
+import logging
+
+from fastmcp import Context
+from sqlalchemy.exc import SQLAlchemyError
+from superset_core.mcp.decorators import tool
+
+from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
+from superset.exceptions import SupersetErrorException, SupersetSecurityException
+from superset.extensions import event_logger
+from superset.mcp_service.sql_lab.schemas import (
+    SaveSqlQueryRequest,
+    SaveSqlQueryResponse,
+)
+from superset.mcp_service.utils.schema_utils import parse_request
+
+logger = logging.getLogger(__name__)
+
+
+@tool(tags=["mutate"])
+@parse_request(SaveSqlQueryRequest)
+async def save_sql_query(
+    request: SaveSqlQueryRequest, ctx: Context
+) -> SaveSqlQueryResponse:
+    """Save a SQL query so it appears in SQL Lab's Saved Queries list.
+
+    Creates a persistent SavedQuery that the user can reload from
+    SQL Lab, share via URL, and find in the Saved Queries page.
+    Requires a database_id, a label (name), and the SQL text.
+    """
+    await ctx.info(
+        "Saving SQL query: database_id=%s, label=%r"
+        % (request.database_id, request.label)
+    )
+
+    try:
+        from flask import g
+
+        from superset import db, security_manager
+        from superset.daos.query import SavedQueryDAO
+        from superset.mcp_service.utils.url_utils import get_superset_base_url
+        from superset.models.core import Database
+
+        # 1. Validate database exists and user has access
+        with event_logger.log_context(action="mcp.save_sql_query.db_validation"):
+            database = (
+                db.session.query(Database).filter_by(id=request.database_id).first()
+            )
+            if not database:
+                raise SupersetErrorException(
+                    SupersetError(
+                        message=(f"Database with ID {request.database_id} not found"),
+                        error_type=SupersetErrorType.DATABASE_NOT_FOUND_ERROR,
+                        level=ErrorLevel.ERROR,
+                    )
+                )
+
+            if not security_manager.can_access_database(database):
+                raise SupersetSecurityException(
+                    SupersetError(
+                        message=(f"Access denied to database {database.database_name}"),
+                        error_type=(SupersetErrorType.DATABASE_SECURITY_ACCESS_ERROR),
+                        level=ErrorLevel.ERROR,
+                    )
+                )
+
+        # 2. Create the saved query
+        with event_logger.log_context(action="mcp.save_sql_query.create"):
+            saved_query = SavedQueryDAO.create(
+                attributes={
+                    "user_id": g.user.id,
+                    "db_id": request.database_id,
+                    "label": request.label,
+                    "sql": request.sql,
+                    "schema": request.schema_name or "",
+                    "catalog": request.catalog,
+                    "description": request.description or "",
+                }
+            )
+            db.session.commit()  # pylint: disable=consider-using-transaction
+
+        # 3. Build response
+        base_url = get_superset_base_url()
+        saved_query_url = f"{base_url}/sqllab?savedQueryId={saved_query.id}"
+
+        await ctx.info(
+            "Saved query created: id=%s, url=%s" % (saved_query.id, saved_query_url)
+        )
+
+        return SaveSqlQueryResponse(
+            id=saved_query.id,
+            label=saved_query.label,
+            sql=saved_query.sql,
+            database_id=request.database_id,
+            schema_name=request.schema_name,
+            catalog=getattr(saved_query, "catalog", None),
+            description=request.description,
+            url=saved_query_url,
+        )
+
+    except (SupersetErrorException, SupersetSecurityException):
+        raise
+    except SQLAlchemyError as e:
+        from superset import db
+
+        db.session.rollback()
+        await ctx.error(
+            "Failed to save SQL query: error=%s, database_id=%s"
+            % (str(e), request.database_id)
+        )
+        raise