codeant-ai-for-open-source[bot] commented on code in PR #38414: URL: https://github.com/apache/superset/pull/38414#discussion_r2894748394
########## 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: **Suggestion:** If `SavedQueryDAO.create` or `db.session.commit()` raises an exception, the SQLAlchemy session is left in a failed transaction state because no rollback is performed, which can cause subsequent database operations in the same request context to fail until the session is manually cleaned up. [resource leak] <details> <summary><b>Severity Level:</b> Major ⚠️</summary> ```mdx - ❌ Subsequent DB operations after failed save can consistently fail. - ⚠️ MCP service stability degraded when persistence errors occur. ``` </details> ```suggestion except (SupersetErrorException, SupersetSecurityException): raise except Exception 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 ``` <details> <summary><b>Steps of Reproduction ✅ </b></summary> ```mdx 1. Run the Superset MCP service and invoke the `save_sql_query` MCP tool defined in `superset/mcp_service/sql_lab/tool/save_sql_query.py:48` with parameters that cause a database error during persistence (for example, in a unit test patch `SavedQueryDAO.create` at `superset/mcp_service/sql_lab/tool/save_sql_query.py:94` or `db.session.commit()` at line 105 to raise a SQLAlchemy exception). 2. The call enters the `try` block in `save_sql_query` (`save_sql_query.py:61-125`), executes until `SavedQueryDAO.create` or `db.session.commit()` fails, and the exception is caught by the generic `except Exception as e:` handler at `save_sql_query.py:128`. 3. In this handler (`save_sql_query.py:128-132`), the code logs the error via `await ctx.error(...)` but does not call `db.session.rollback()`, leaving the SQLAlchemy `db.session` (imported from `superset` at line 64) in a failed transaction state as per SQLAlchemy session semantics. 4. Within the same application context (e.g., in the same test case, after the failed `save_sql_query` call) perform another database operation using `from superset import db; db.session.execute("SELECT 1")`; this operation will raise a `PendingRollbackError`/similar because the session from the previous failure was never rolled back, demonstrating that subsequent DB operations are impacted. ``` </details> <details> <summary><b>Prompt for AI Agent 🤖 </b></summary> ```mdx This is a comment left during a code review. **Path:** superset/mcp_service/sql_lab/tool/save_sql_query.py **Line:** 126:133 **Comment:** *Resource Leak: If `SavedQueryDAO.create` or `db.session.commit()` raises an exception, the SQLAlchemy session is left in a failed transaction state because no rollback is performed, which can cause subsequent database operations in the same request context to fail until the session is manually cleaned up. 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. ``` </details> <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F38414&comment_hash=06fae4e258f2d5c11421f2d2af5fee2c0dc0b3342d34694edc5746724f0d5d90&reaction=like'>👍</a> | <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F38414&comment_hash=06fae4e258f2d5c11421f2d2af5fee2c0dc0b3342d34694edc5746724f0d5d90&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]
