aminghadersohi commented on code in PR #38414: URL: https://github.com/apache/superset/pull/38414#discussion_r2895038253
########## superset/mcp_service/sql_lab/tool/save_sql_query.py: ########## @@ -0,0 +1,133 @@ +# 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 superset_core.api.mcp 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 Exception as e: + await ctx.error( + "Failed to save SQL query: error=%s, database_id=%s" + % (str(e), request.database_id) + ) + raise Review Comment: Good catch. If SavedQueryDAO.create or the commit fails, the session can be left in a bad state. I'll add a try/except with db.session.rollback() to handle that properly. Thanks for the review. ########## superset/mcp_service/sql_lab/schemas.py: ########## @@ -115,6 +115,64 @@ class ExecuteSqlResponse(BaseModel): ) +class SaveSqlQueryRequest(BaseModel): + """Request schema for saving a SQL query.""" + + database_id: int = Field( + ..., description="Database connection ID the query runs against" + ) + label: str = Field( + ..., + description="Name for the saved query (shown in Saved Queries list)", + min_length=1, + max_length=256, + ) + sql: str = Field( + ..., + description="SQL query text to save", + ) + schema_name: str | None = Field( + None, + description="Schema the query targets", + alias="schema", + ) + catalog: str | None = Field(None, description="Catalog name (if applicable)") + description: str | None = Field( + None, description="Optional description of the query" + ) + + @field_validator("sql") + @classmethod + def sql_not_empty(cls, v: str) -> str: + if not v or not v.strip(): + raise ValueError("SQL query cannot be empty") + return v.strip() + + @field_validator("label") + @classmethod + def label_not_empty(cls, v: str) -> str: + if not v or not v.strip(): + raise ValueError("Label cannot be empty") + return v.strip() + + +class SaveSqlQueryResponse(BaseModel): + """Response schema for a saved SQL query.""" + + id: int = Field(..., description="Saved query ID") + label: str = Field(..., description="Query name") + sql: str = Field(..., description="SQL query text") + database_id: int = Field(..., description="Database ID") + schema_name: str | None = Field(None, description="Schema name", alias="schema") + description: str | None = Field(None, description="Query description") Review Comment: Thanks for confirming — I'll add the catalog field to the response schema. -- 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]
