mirror of
https://github.com/apache/superset.git
synced 2026-04-09 19:35:21 +00:00
138 lines
4.9 KiB
Python
138 lines
4.9 KiB
Python
# 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
|