From 6a254a27cc476e21392d005a0b27efa38bddbebf Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Wed, 20 May 2026 22:44:12 +0000 Subject: [PATCH 01/10] feat(mcp): add list and get tools for saved queries and query history Implements list_saved_queries, get_saved_query_info, list_queries, and get_query_info MCP tools in new saved_query/ and query/ domains. Co-Authored-By: Claude Sonnet 4.6 --- superset/mcp_service/app.py | 8 + superset/mcp_service/query/__init__.py | 16 + superset/mcp_service/query/schemas.py | 290 ++++++++++++++++++ superset/mcp_service/query/tool/__init__.py | 24 ++ .../mcp_service/query/tool/get_query_info.py | 122 ++++++++ .../mcp_service/query/tool/list_queries.py | 156 ++++++++++ superset/mcp_service/saved_query/__init__.py | 16 + superset/mcp_service/saved_query/schemas.py | 270 ++++++++++++++++ .../mcp_service/saved_query/tool/__init__.py | 24 ++ .../saved_query/tool/get_saved_query_info.py | 129 ++++++++ .../saved_query/tool/list_saved_queries.py | 159 ++++++++++ .../unit_tests/mcp_service/query/__init__.py | 16 + .../mcp_service/query/tool/__init__.py | 16 + .../query/tool/test_query_tools.py | 271 ++++++++++++++++ .../mcp_service/saved_query/__init__.py | 16 + .../mcp_service/saved_query/tool/__init__.py | 16 + .../tool/test_saved_query_tools.py | 256 ++++++++++++++++ 17 files changed, 1805 insertions(+) create mode 100644 superset/mcp_service/query/__init__.py create mode 100644 superset/mcp_service/query/schemas.py create mode 100644 superset/mcp_service/query/tool/__init__.py create mode 100644 superset/mcp_service/query/tool/get_query_info.py create mode 100644 superset/mcp_service/query/tool/list_queries.py create mode 100644 superset/mcp_service/saved_query/__init__.py create mode 100644 superset/mcp_service/saved_query/schemas.py create mode 100644 superset/mcp_service/saved_query/tool/__init__.py create mode 100644 superset/mcp_service/saved_query/tool/get_saved_query_info.py create mode 100644 superset/mcp_service/saved_query/tool/list_saved_queries.py create mode 100644 tests/unit_tests/mcp_service/query/__init__.py create mode 100644 tests/unit_tests/mcp_service/query/tool/__init__.py create mode 100644 tests/unit_tests/mcp_service/query/tool/test_query_tools.py create mode 100644 tests/unit_tests/mcp_service/saved_query/__init__.py create mode 100644 tests/unit_tests/mcp_service/saved_query/tool/__init__.py create mode 100644 tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py diff --git a/superset/mcp_service/app.py b/superset/mcp_service/app.py index 01566b364569..fbc7d5e5c143 100644 --- a/superset/mcp_service/app.py +++ b/superset/mcp_service/app.py @@ -654,6 +654,14 @@ def create_mcp_app( from superset.mcp_service.explore.tool import ( # noqa: F401, E402 generate_explore_link, ) +from superset.mcp_service.query.tool import ( # noqa: F401, E402 + get_query_info, + list_queries, +) +from superset.mcp_service.saved_query.tool import ( # noqa: F401, E402 + get_saved_query_info, + list_saved_queries, +) from superset.mcp_service.sql_lab.tool import ( # noqa: F401, E402 execute_sql, open_sql_lab_with_context, diff --git a/superset/mcp_service/query/__init__.py b/superset/mcp_service/query/__init__.py new file mode 100644 index 000000000000..13a83393a912 --- /dev/null +++ b/superset/mcp_service/query/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/superset/mcp_service/query/schemas.py b/superset/mcp_service/query/schemas.py new file mode 100644 index 000000000000..07c4bdcab38d --- /dev/null +++ b/superset/mcp_service/query/schemas.py @@ -0,0 +1,290 @@ +# 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. + +""" +Pydantic schemas for query history-related responses +""" + +from __future__ import annotations + +from datetime import datetime +from typing import Annotated, Any, Dict, List, Literal + +from pydantic import ( + BaseModel, + ConfigDict, + Field, + field_validator, + model_serializer, + model_validator, + PositiveInt, +) + +from superset.daos.base import ColumnOperator, ColumnOperatorEnum +from superset.mcp_service.constants import MAX_PAGE_SIZE +from superset.mcp_service.privacy import filter_user_directory_fields +from superset.mcp_service.system.schemas import PaginationInfo +from superset.mcp_service.utils.schema_utils import ( + parse_json_or_list, + parse_json_or_model_list, +) + +DEFAULT_QUERY_COLUMNS = ["id", "sql", "status", "start_time", "database_id", "schema"] +SORTABLE_QUERY_COLUMNS = [ + "id", + "start_time", + "end_time", + "status", + "database_id", +] +ALL_QUERY_COLUMNS = [ + "id", + "sql", + "status", + "start_time", + "end_time", + "rows", + "database_id", + "schema", + "tab_name", + "error_message", + "client_id", + "limit", + "progress", + "changed_on", +] + +DEFAULT_QUERY_PAGE_SIZE = 25 + + +class QueryFilter(ColumnOperator): + """ + Filter object for query history listing. + col: The column to filter on. Must be one of the allowed filter fields. + opr: The operator to use. Must be one of the supported operators. + value: The value to filter by (type depends on col and opr). + """ + + col: Literal["status", "database_id", "schema"] = Field( + ..., + description="Column to filter on.", + ) + opr: ColumnOperatorEnum = Field( + ..., + description="Operator to use.", + ) + value: str | int | float | bool | List[str | int | float | bool] = Field( + ..., description="Value to filter by (type depends on col and opr)" + ) + + +class QueryInfo(BaseModel): + id: int | None = Field(None, description="Query ID") + sql: str | None = Field(None, description="SQL query text") + status: str | None = Field(None, description="Query execution status") + start_time: float | None = Field( + None, description="Query start time (seconds since epoch)" + ) + end_time: float | None = Field( + None, description="Query end time (seconds since epoch)" + ) + rows: int | None = Field(None, description="Number of rows returned or affected") + database_id: int | None = Field(None, description="Database connection ID") + schema: str | None = Field(None, description="Database schema name") + tab_name: str | None = Field(None, description="SQL Lab tab name") + error_message: str | None = Field(None, description="Error message if query failed") + client_id: str | None = Field(None, description="Client-assigned query identifier") + limit: int | None = Field(None, description="Row limit applied to the query") + progress: int | None = Field(None, description="Query execution progress (0-100)") + changed_on: str | datetime | None = Field( + None, description="Last modification timestamp" + ) + model_config = ConfigDict( + from_attributes=True, + ser_json_timedelta="iso8601", + populate_by_name=True, + ) + + @model_serializer(mode="wrap") + def _filter_fields_by_context(self, serializer: Any, info: Any) -> Dict[str, Any]: + data = filter_user_directory_fields(serializer(self)) + + if info.context and isinstance(info.context, dict): + select_columns = info.context.get("select_columns") + if select_columns: + requested_fields = set(select_columns) + return {k: v for k, v in data.items() if k in requested_fields} + + return data + + +class QueryList(BaseModel): + queries: List[QueryInfo] + count: int + total_count: int + page: int + page_size: int + total_pages: int + has_previous: bool + has_next: bool + columns_requested: List[str] = Field( + default_factory=list, + description="Requested columns for the response", + ) + columns_loaded: List[str] = Field( + default_factory=list, + description="Columns that were actually loaded for each query", + ) + columns_available: List[str] = Field( + default_factory=list, + description="All columns available for selection via select_columns parameter", + ) + sortable_columns: List[str] = Field( + default_factory=list, + description="Columns that can be used with order_column parameter", + ) + filters_applied: List[QueryFilter] = Field( + default_factory=list, + description="List of advanced filter dicts applied to the query.", + ) + pagination: PaginationInfo | None = None + timestamp: datetime | None = None + model_config = ConfigDict(ser_json_timedelta="iso8601") + + +class ListQueriesRequest(BaseModel): + """Request schema for list_queries.""" + + filters: Annotated[ + List[QueryFilter], + Field( + default_factory=list, + description="List of filter objects (column, operator, value). Each " + "filter is an object with 'col', 'opr', and 'value' " + "properties. Cannot be used together with 'search'.", + ), + ] + select_columns: Annotated[ + List[str], + Field( + default_factory=list, + description="List of columns to select. Defaults to common columns if not " + "specified.", + ), + ] + search: Annotated[ + str | None, + Field( + default=None, + description="Text search string to match against query fields. " + "Cannot be used together with 'filters'.", + ), + ] + order_column: Annotated[ + str | None, + Field(default=None, description="Column to order results by"), + ] + order_direction: Annotated[ + Literal["asc", "desc"], + Field( + default="desc", + description="Direction to order results ('asc' or 'desc')", + ), + ] + page: Annotated[ + PositiveInt, + Field(default=1, description="Page number for pagination (1-based)"), + ] + page_size: Annotated[ + int, + Field( + default=DEFAULT_QUERY_PAGE_SIZE, + gt=0, + le=MAX_PAGE_SIZE, + description=f"Number of items per page (max {MAX_PAGE_SIZE})", + ), + ] + + @field_validator("filters", mode="before") + @classmethod + def parse_filters(cls, v: Any) -> List[QueryFilter]: + """Accept both JSON string and list of objects.""" + return parse_json_or_model_list(v, QueryFilter, "filters") + + @field_validator("select_columns", mode="before") + @classmethod + def parse_columns(cls, v: Any) -> List[str]: + """Accept JSON array, list, or comma-separated string.""" + return parse_json_or_list(v, "select_columns") + + @model_validator(mode="after") + def validate_search_and_filters(self) -> "ListQueriesRequest": + """Prevent using both search and filters simultaneously.""" + if self.search and self.filters: + raise ValueError( + "Cannot use both 'search' and 'filters' parameters simultaneously. " + "Use either 'search' for text-based searching across multiple fields, " + "or 'filters' for precise column-based filtering, but not both." + ) + return self + + +class QueryError(BaseModel): + error: str = Field(..., description="Error message") + error_type: str = Field(..., description="Type of error") + timestamp: str | datetime | None = Field(None, description="Error timestamp") + model_config = ConfigDict(ser_json_timedelta="iso8601") + + @classmethod + def create(cls, error: str, error_type: str) -> "QueryError": + """Create a standardized QueryError with timestamp.""" + from datetime import datetime, timezone + + return cls( + error=error, error_type=error_type, timestamp=datetime.now(timezone.utc) + ) + + +class GetQueryInfoRequest(BaseModel): + """Request schema for get_query_info with support for numeric ID only.""" + + identifier: Annotated[ + int, + Field(description="Query ID (numeric)"), + ] + + +def serialize_query_object(query: Any) -> QueryInfo | None: + if not query: + return None + + return QueryInfo( + id=getattr(query, "id", None), + sql=getattr(query, "sql", None), + status=getattr(query, "status", None), + start_time=getattr(query, "start_time", None), + end_time=getattr(query, "end_time", None), + rows=getattr(query, "rows", None), + database_id=getattr(query, "database_id", None), + schema=getattr(query, "schema", None), + tab_name=getattr(query, "tab_name", None), + error_message=getattr(query, "error_message", None), + client_id=getattr(query, "client_id", None), + limit=getattr(query, "limit", None), + progress=getattr(query, "progress", None), + changed_on=getattr(query, "changed_on", None), + ) diff --git a/superset/mcp_service/query/tool/__init__.py b/superset/mcp_service/query/tool/__init__.py new file mode 100644 index 000000000000..3e6edcbda474 --- /dev/null +++ b/superset/mcp_service/query/tool/__init__.py @@ -0,0 +1,24 @@ +# 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. + +from .get_query_info import get_query_info +from .list_queries import list_queries + +__all__ = [ + "list_queries", + "get_query_info", +] diff --git a/superset/mcp_service/query/tool/get_query_info.py b/superset/mcp_service/query/tool/get_query_info.py new file mode 100644 index 000000000000..dc94a947d6cd --- /dev/null +++ b/superset/mcp_service/query/tool/get_query_info.py @@ -0,0 +1,122 @@ +# 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. + +""" +Get query info FastMCP tool + +This module contains the FastMCP tool for getting detailed information +about a specific SQL query from the query history. +""" + +import logging +from datetime import datetime, timezone + +from fastmcp import Context +from superset_core.mcp.decorators import tool, ToolAnnotations + +from superset.extensions import event_logger +from superset.mcp_service.mcp_core import ModelGetInfoCore +from superset.mcp_service.query.schemas import ( + GetQueryInfoRequest, + QueryError, + QueryInfo, + serialize_query_object, +) + +logger = logging.getLogger(__name__) + + +@tool( + tags=["discovery"], + class_permission_name="Query", + annotations=ToolAnnotations( + title="Get query info", + readOnlyHint=True, + destructiveHint=False, + ), +) +async def get_query_info( + request: GetQueryInfoRequest, ctx: Context +) -> QueryInfo | QueryError: + """Get SQL query history details by ID. + + Returns query details including SQL text, execution status, timing, + row count, and any error messages. + + IMPORTANT FOR LLM CLIENTS: + - Use numeric ID (e.g., 123) + - To find a query ID, use the list_queries tool first + + Example usage: + ```json + { + "identifier": 123 + } + ``` + """ + await ctx.info( + "Retrieving query information: identifier=%s" % (request.identifier,) + ) + + try: + from superset.daos.query import QueryDAO + + with event_logger.log_context(action="mcp.get_query_info.lookup"): + get_tool = ModelGetInfoCore( + dao_class=QueryDAO, + output_schema=QueryInfo, + error_schema=QueryError, + serializer=serialize_query_object, + supports_slug=False, + logger=logger, + ) + + result = get_tool.run_tool(request.identifier) + + if isinstance(result, QueryInfo): + await ctx.info( + "Query information retrieved successfully: " + "query_id=%s, status=%s, database_id=%s" + % ( + result.id, + result.status, + result.database_id, + ) + ) + else: + await ctx.warning( + "Query retrieval failed: error_type=%s, error=%s" + % (result.error_type, result.error) + ) + + return result + + except Exception as e: + await ctx.error( + "Query information retrieval failed: identifier=%s, error=%s, " + "error_type=%s" + % ( + request.identifier, + str(e), + type(e).__name__, + ) + ) + return QueryError( + error=f"Failed to get query info: {str(e)}", + error_type="InternalError", + timestamp=datetime.now(timezone.utc), + ) diff --git a/superset/mcp_service/query/tool/list_queries.py b/superset/mcp_service/query/tool/list_queries.py new file mode 100644 index 000000000000..c39f1f6062dc --- /dev/null +++ b/superset/mcp_service/query/tool/list_queries.py @@ -0,0 +1,156 @@ +# 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. + +""" +List queries FastMCP tool + +This module contains the FastMCP tool for listing SQL query history +with filtering, search, and pagination. +""" + +import logging + +from fastmcp import Context +from superset_core.mcp.decorators import tool, ToolAnnotations + +from superset.extensions import event_logger +from superset.mcp_service.mcp_core import ModelListCore +from superset.mcp_service.query.schemas import ( + ALL_QUERY_COLUMNS, + DEFAULT_QUERY_COLUMNS, + ListQueriesRequest, + QueryError, + QueryFilter, + QueryInfo, + QueryList, + serialize_query_object, + SORTABLE_QUERY_COLUMNS, +) + +logger = logging.getLogger(__name__) + +_DEFAULT_LIST_QUERIES_REQUEST = ListQueriesRequest() + + +@tool( + tags=["core"], + class_permission_name="Query", + annotations=ToolAnnotations( + title="List queries", + readOnlyHint=True, + destructiveHint=False, + ), +) +async def list_queries( + request: ListQueriesRequest | None = None, + ctx: Context | None = None, +) -> QueryList | QueryError: + """List SQL query history with filtering and search. + + Returns recent queries executed by the current user (or all queries for + admins), including SQL text, status, timing, and database information. + Results are ordered by start_time descending (most recent first) by default. + + Sortable columns for order_column: id, start_time, end_time, status, + database_id + """ + if ctx is None: + raise RuntimeError("FastMCP context is required for list_queries") + + request = request or _DEFAULT_LIST_QUERIES_REQUEST.model_copy(deep=True) + + await ctx.info( + "Listing queries: page=%s, page_size=%s, search=%s" + % ( + request.page, + request.page_size, + request.search, + ) + ) + await ctx.debug( + "Query listing parameters: filters=%s, order_column=%s, " + "order_direction=%s, select_columns=%s" + % ( + request.filters, + request.order_column, + request.order_direction, + request.select_columns, + ) + ) + + try: + from superset.daos.query import QueryDAO + + def _serialize_query(obj: object, cols: list[str] | None) -> QueryInfo | None: + return serialize_query_object(obj) + + list_tool = ModelListCore( + dao_class=QueryDAO, + output_schema=QueryInfo, + item_serializer=_serialize_query, + filter_type=QueryFilter, + default_columns=DEFAULT_QUERY_COLUMNS, + search_columns=["tab_name"], + list_field_name="queries", + output_list_schema=QueryList, + all_columns=ALL_QUERY_COLUMNS, + sortable_columns=SORTABLE_QUERY_COLUMNS, + logger=logger, + ) + + with event_logger.log_context(action="mcp.list_queries.query"): + result = list_tool.run_tool( + filters=request.filters, + search=request.search, + select_columns=request.select_columns, + order_column=request.order_column or "start_time", + order_direction=request.order_direction, + page=max(request.page - 1, 0), + page_size=request.page_size, + ) + + await ctx.info( + "Queries listed successfully: count=%s, total_count=%s, total_pages=%s" + % ( + len(result.queries) if hasattr(result, "queries") else 0, + getattr(result, "total_count", None), + getattr(result, "total_pages", None), + ) + ) + + columns_to_filter = result.columns_requested + await ctx.debug( + "Applying field filtering via serialization context: columns=%s" + % (columns_to_filter,) + ) + with event_logger.log_context(action="mcp.list_queries.serialization"): + return result.model_dump( + mode="json", + context={"select_columns": columns_to_filter}, + ) + + except Exception as e: + await ctx.error( + "Query listing failed: page=%s, page_size=%s, error=%s, error_type=%s" + % ( + request.page, + request.page_size, + str(e), + type(e).__name__, + ) + ) + raise diff --git a/superset/mcp_service/saved_query/__init__.py b/superset/mcp_service/saved_query/__init__.py new file mode 100644 index 000000000000..13a83393a912 --- /dev/null +++ b/superset/mcp_service/saved_query/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/superset/mcp_service/saved_query/schemas.py b/superset/mcp_service/saved_query/schemas.py new file mode 100644 index 000000000000..c55298637e89 --- /dev/null +++ b/superset/mcp_service/saved_query/schemas.py @@ -0,0 +1,270 @@ +# 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. + +""" +Pydantic schemas for saved query-related responses +""" + +from __future__ import annotations + +from datetime import datetime +from typing import Annotated, Any, Dict, List, Literal + +from pydantic import ( + BaseModel, + ConfigDict, + Field, + field_validator, + model_serializer, + model_validator, + PositiveInt, +) + +from superset.daos.base import ColumnOperator, ColumnOperatorEnum +from superset.mcp_service.constants import DEFAULT_PAGE_SIZE, MAX_PAGE_SIZE +from superset.mcp_service.privacy import filter_user_directory_fields +from superset.mcp_service.system.schemas import PaginationInfo +from superset.mcp_service.utils.schema_utils import ( + parse_json_or_list, + parse_json_or_model_list, +) + +DEFAULT_SAVED_QUERY_COLUMNS = ["id", "label", "db_id", "schema", "uuid"] +SORTABLE_SAVED_QUERY_COLUMNS = [ + "id", + "label", + "db_id", + "schema", + "changed_on", + "created_on", +] +ALL_SAVED_QUERY_COLUMNS = [ + "id", + "label", + "db_id", + "schema", + "uuid", + "sql", + "description", + "changed_on", + "created_on", +] + + +class SavedQueryFilter(ColumnOperator): + """ + Filter object for saved query listing. + col: The column to filter on. Must be one of the allowed filter fields. + opr: The operator to use. Must be one of the supported operators. + value: The value to filter by (type depends on col and opr). + """ + + col: Literal["label", "db_id", "schema"] = Field( + ..., + description="Column to filter on.", + ) + opr: ColumnOperatorEnum = Field( + ..., + description="Operator to use.", + ) + value: str | int | float | bool | List[str | int | float | bool] = Field( + ..., description="Value to filter by (type depends on col and opr)" + ) + + +class SavedQueryInfo(BaseModel): + id: int | None = Field(None, description="Saved query ID") + uuid: str | None = Field(None, description="Saved query UUID") + label: str | None = Field(None, description="Saved query label/name") + sql: str | None = Field(None, description="SQL query text") + db_id: int | None = Field(None, description="Database connection ID") + schema: str | None = Field(None, description="Database schema name") + description: str | None = Field(None, description="User-provided description") + changed_on: str | datetime | None = Field( + None, description="Last modification timestamp" + ) + created_on: str | datetime | None = Field(None, description="Creation timestamp") + model_config = ConfigDict( + from_attributes=True, + ser_json_timedelta="iso8601", + populate_by_name=True, + ) + + @model_serializer(mode="wrap") + def _filter_fields_by_context(self, serializer: Any, info: Any) -> Dict[str, Any]: + data = filter_user_directory_fields(serializer(self)) + + if info.context and isinstance(info.context, dict): + select_columns = info.context.get("select_columns") + if select_columns: + requested_fields = set(select_columns) + return {k: v for k, v in data.items() if k in requested_fields} + + return data + + +class SavedQueryList(BaseModel): + saved_queries: List[SavedQueryInfo] + count: int + total_count: int + page: int + page_size: int + total_pages: int + has_previous: bool + has_next: bool + columns_requested: List[str] = Field( + default_factory=list, + description="Requested columns for the response", + ) + columns_loaded: List[str] = Field( + default_factory=list, + description="Columns that were actually loaded for each saved query", + ) + columns_available: List[str] = Field( + default_factory=list, + description="All columns available for selection via select_columns parameter", + ) + sortable_columns: List[str] = Field( + default_factory=list, + description="Columns that can be used with order_column parameter", + ) + filters_applied: List[SavedQueryFilter] = Field( + default_factory=list, + description="List of advanced filter dicts applied to the query.", + ) + pagination: PaginationInfo | None = None + timestamp: datetime | None = None + model_config = ConfigDict(ser_json_timedelta="iso8601") + + +class ListSavedQueriesRequest(BaseModel): + """Request schema for list_saved_queries.""" + + filters: Annotated[ + List[SavedQueryFilter], + Field( + default_factory=list, + description="List of filter objects (column, operator, value). Each " + "filter is an object with 'col', 'opr', and 'value' " + "properties. Cannot be used together with 'search'.", + ), + ] + select_columns: Annotated[ + List[str], + Field( + default_factory=list, + description="List of columns to select. Defaults to common columns if not " + "specified.", + ), + ] + search: Annotated[ + str | None, + Field( + default=None, + description="Text search string to match against saved query fields. " + "Cannot be used together with 'filters'.", + ), + ] + order_column: Annotated[ + str | None, Field(default=None, description="Column to order results by") + ] + order_direction: Annotated[ + Literal["asc", "desc"], + Field( + default="desc", description="Direction to order results ('asc' or 'desc')" + ), + ] + page: Annotated[ + PositiveInt, + Field(default=1, description="Page number for pagination (1-based)"), + ] + page_size: Annotated[ + int, + Field( + default=DEFAULT_PAGE_SIZE, + gt=0, + le=MAX_PAGE_SIZE, + description=f"Number of items per page (max {MAX_PAGE_SIZE})", + ), + ] + + @field_validator("filters", mode="before") + @classmethod + def parse_filters(cls, v: Any) -> List[SavedQueryFilter]: + """Accept both JSON string and list of objects.""" + return parse_json_or_model_list(v, SavedQueryFilter, "filters") + + @field_validator("select_columns", mode="before") + @classmethod + def parse_columns(cls, v: Any) -> List[str]: + """Accept JSON array, list, or comma-separated string.""" + return parse_json_or_list(v, "select_columns") + + @model_validator(mode="after") + def validate_search_and_filters(self) -> "ListSavedQueriesRequest": + """Prevent using both search and filters simultaneously.""" + if self.search and self.filters: + raise ValueError( + "Cannot use both 'search' and 'filters' parameters simultaneously. " + "Use either 'search' for text-based searching across multiple fields, " + "or 'filters' for precise column-based filtering, but not both." + ) + return self + + +class SavedQueryError(BaseModel): + error: str = Field(..., description="Error message") + error_type: str = Field(..., description="Type of error") + timestamp: str | datetime | None = Field(None, description="Error timestamp") + model_config = ConfigDict(ser_json_timedelta="iso8601") + + @classmethod + def create(cls, error: str, error_type: str) -> "SavedQueryError": + """Create a standardized SavedQueryError with timestamp.""" + from datetime import datetime, timezone + + return cls( + error=error, error_type=error_type, timestamp=datetime.now(timezone.utc) + ) + + +class GetSavedQueryInfoRequest(BaseModel): + """Request schema for get_saved_query_info with support for ID or UUID.""" + + identifier: Annotated[ + int | str, + Field(description="Saved query identifier - can be numeric ID or UUID string"), + ] + + +def serialize_saved_query_object(saved_query: Any) -> SavedQueryInfo | None: + if not saved_query: + return None + + return SavedQueryInfo( + id=getattr(saved_query, "id", None), + uuid=str(getattr(saved_query, "uuid", "")) + if getattr(saved_query, "uuid", None) + else None, + label=getattr(saved_query, "label", None), + sql=getattr(saved_query, "sql", None), + db_id=getattr(saved_query, "db_id", None), + schema=getattr(saved_query, "schema", None), + description=getattr(saved_query, "description", None), + changed_on=getattr(saved_query, "changed_on", None), + created_on=getattr(saved_query, "created_on", None), + ) diff --git a/superset/mcp_service/saved_query/tool/__init__.py b/superset/mcp_service/saved_query/tool/__init__.py new file mode 100644 index 000000000000..af366fd53122 --- /dev/null +++ b/superset/mcp_service/saved_query/tool/__init__.py @@ -0,0 +1,24 @@ +# 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. + +from .get_saved_query_info import get_saved_query_info +from .list_saved_queries import list_saved_queries + +__all__ = [ + "list_saved_queries", + "get_saved_query_info", +] diff --git a/superset/mcp_service/saved_query/tool/get_saved_query_info.py b/superset/mcp_service/saved_query/tool/get_saved_query_info.py new file mode 100644 index 000000000000..9b3a1be74b22 --- /dev/null +++ b/superset/mcp_service/saved_query/tool/get_saved_query_info.py @@ -0,0 +1,129 @@ +# 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. + +""" +Get saved query info FastMCP tool + +This module contains the FastMCP tool for getting detailed information +about a specific saved SQL query. +""" + +import logging +from datetime import datetime, timezone + +from fastmcp import Context +from superset_core.mcp.decorators import tool, ToolAnnotations + +from superset.extensions import event_logger +from superset.mcp_service.mcp_core import ModelGetInfoCore +from superset.mcp_service.saved_query.schemas import ( + GetSavedQueryInfoRequest, + SavedQueryError, + SavedQueryInfo, + serialize_saved_query_object, +) + +logger = logging.getLogger(__name__) + + +@tool( + tags=["discovery"], + class_permission_name="SavedQuery", + annotations=ToolAnnotations( + title="Get saved query info", + readOnlyHint=True, + destructiveHint=False, + ), +) +async def get_saved_query_info( + request: GetSavedQueryInfoRequest, ctx: Context +) -> SavedQueryInfo | SavedQueryError: + """Get saved query details by ID or UUID. + + Returns the full saved query including SQL text, label, database, + schema, and timestamps. + + IMPORTANT FOR LLM CLIENTS: + - Use numeric ID (e.g., 42) or UUID string (e.g., "a1b2c3d4-...") + - To find a saved query ID, use the list_saved_queries tool first + + Example usage: + ```json + { + "identifier": 42 + } + ``` + + Or with UUID: + ```json + { + "identifier": "a1b2c3d4-5678-90ab-cdef-1234567890ab" + } + ``` + """ + await ctx.info( + "Retrieving saved query information: identifier=%s" % (request.identifier,) + ) + + try: + from superset.daos.query import SavedQueryDAO + + with event_logger.log_context(action="mcp.get_saved_query_info.lookup"): + get_tool = ModelGetInfoCore( + dao_class=SavedQueryDAO, + output_schema=SavedQueryInfo, + error_schema=SavedQueryError, + serializer=serialize_saved_query_object, + supports_slug=False, + logger=logger, + ) + + result = get_tool.run_tool(request.identifier) + + if isinstance(result, SavedQueryInfo): + await ctx.info( + "Saved query information retrieved successfully: " + "saved_query_id=%s, label=%s, db_id=%s" + % ( + result.id, + result.label, + result.db_id, + ) + ) + else: + await ctx.warning( + "Saved query retrieval failed: error_type=%s, error=%s" + % (result.error_type, result.error) + ) + + return result + + except Exception as e: + await ctx.error( + "Saved query information retrieval failed: identifier=%s, error=%s, " + "error_type=%s" + % ( + request.identifier, + str(e), + type(e).__name__, + ) + ) + return SavedQueryError( + error=f"Failed to get saved query info: {str(e)}", + error_type="InternalError", + timestamp=datetime.now(timezone.utc), + ) diff --git a/superset/mcp_service/saved_query/tool/list_saved_queries.py b/superset/mcp_service/saved_query/tool/list_saved_queries.py new file mode 100644 index 000000000000..d1820517ca85 --- /dev/null +++ b/superset/mcp_service/saved_query/tool/list_saved_queries.py @@ -0,0 +1,159 @@ +# 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. + +""" +List saved queries FastMCP tool + +This module contains the FastMCP tool for listing saved SQL queries +with filtering, search, and pagination. +""" + +import logging + +from fastmcp import Context +from superset_core.mcp.decorators import tool, ToolAnnotations + +from superset.extensions import event_logger +from superset.mcp_service.mcp_core import ModelListCore +from superset.mcp_service.saved_query.schemas import ( + ALL_SAVED_QUERY_COLUMNS, + DEFAULT_SAVED_QUERY_COLUMNS, + ListSavedQueriesRequest, + SavedQueryError, + SavedQueryFilter, + SavedQueryInfo, + SavedQueryList, + serialize_saved_query_object, + SORTABLE_SAVED_QUERY_COLUMNS, +) + +logger = logging.getLogger(__name__) + +_DEFAULT_LIST_SAVED_QUERIES_REQUEST = ListSavedQueriesRequest() + + +@tool( + tags=["core"], + class_permission_name="SavedQuery", + annotations=ToolAnnotations( + title="List saved queries", + readOnlyHint=True, + destructiveHint=False, + ), +) +async def list_saved_queries( + request: ListSavedQueriesRequest | None = None, + ctx: Context | None = None, +) -> SavedQueryList | SavedQueryError: + """List saved SQL queries with filtering and search. + + Returns saved queries owned by the current user, including label, SQL, + database ID, and schema. + + Sortable columns for order_column: id, label, db_id, schema, + changed_on, created_on + """ + if ctx is None: + raise RuntimeError("FastMCP context is required for list_saved_queries") + + request = request or _DEFAULT_LIST_SAVED_QUERIES_REQUEST.model_copy(deep=True) + + await ctx.info( + "Listing saved queries: page=%s, page_size=%s, search=%s" + % ( + request.page, + request.page_size, + request.search, + ) + ) + await ctx.debug( + "Saved query listing parameters: filters=%s, order_column=%s, " + "order_direction=%s, select_columns=%s" + % ( + request.filters, + request.order_column, + request.order_direction, + request.select_columns, + ) + ) + + try: + from superset.daos.query import SavedQueryDAO + + def _serialize_saved_query( + obj: object, cols: list[str] | None + ) -> SavedQueryInfo | None: + return serialize_saved_query_object(obj) + + list_tool = ModelListCore( + dao_class=SavedQueryDAO, + output_schema=SavedQueryInfo, + item_serializer=_serialize_saved_query, + filter_type=SavedQueryFilter, + default_columns=DEFAULT_SAVED_QUERY_COLUMNS, + search_columns=["label", "description"], + list_field_name="saved_queries", + output_list_schema=SavedQueryList, + all_columns=ALL_SAVED_QUERY_COLUMNS, + sortable_columns=SORTABLE_SAVED_QUERY_COLUMNS, + logger=logger, + ) + + with event_logger.log_context(action="mcp.list_saved_queries.query"): + result = list_tool.run_tool( + filters=request.filters, + search=request.search, + select_columns=request.select_columns, + order_column=request.order_column, + order_direction=request.order_direction, + page=max(request.page - 1, 0), + page_size=request.page_size, + ) + + await ctx.info( + "Saved queries listed successfully: count=%s, total_count=%s, " + "total_pages=%s" + % ( + len(result.saved_queries) if hasattr(result, "saved_queries") else 0, + getattr(result, "total_count", None), + getattr(result, "total_pages", None), + ) + ) + + columns_to_filter = result.columns_requested + await ctx.debug( + "Applying field filtering via serialization context: columns=%s" + % (columns_to_filter,) + ) + with event_logger.log_context(action="mcp.list_saved_queries.serialization"): + return result.model_dump( + mode="json", + context={"select_columns": columns_to_filter}, + ) + + except Exception as e: + await ctx.error( + "Saved query listing failed: page=%s, page_size=%s, error=%s, " + "error_type=%s" + % ( + request.page, + request.page_size, + str(e), + type(e).__name__, + ) + ) + raise diff --git a/tests/unit_tests/mcp_service/query/__init__.py b/tests/unit_tests/mcp_service/query/__init__.py new file mode 100644 index 000000000000..13a83393a912 --- /dev/null +++ b/tests/unit_tests/mcp_service/query/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/tests/unit_tests/mcp_service/query/tool/__init__.py b/tests/unit_tests/mcp_service/query/tool/__init__.py new file mode 100644 index 000000000000..13a83393a912 --- /dev/null +++ b/tests/unit_tests/mcp_service/query/tool/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/tests/unit_tests/mcp_service/query/tool/test_query_tools.py b/tests/unit_tests/mcp_service/query/tool/test_query_tools.py new file mode 100644 index 000000000000..8e12d109a540 --- /dev/null +++ b/tests/unit_tests/mcp_service/query/tool/test_query_tools.py @@ -0,0 +1,271 @@ +# 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. + + +import logging +from unittest.mock import MagicMock, patch + +import pytest +from fastmcp import Client +from pydantic import ValidationError + +from superset.mcp_service.app import mcp +from superset.mcp_service.query.schemas import ( + ListQueriesRequest, + QueryFilter, +) +from superset.utils import json + +logging.basicConfig(level=logging.DEBUG) +logger = logging.getLogger(__name__) + + +class TestQueryFilterSchema: + """Tests for QueryFilter schema — filterable columns.""" + + def test_invalid_filter_column_rejected(self): + """Columns not in the Literal set must be rejected.""" + with pytest.raises(ValidationError): + QueryFilter(col="not_a_real_column", opr="eq", value="test") + + def test_user_id_is_rejected_as_filter_column(self): + """user_id is an internal field and should not be a filter column.""" + with pytest.raises(ValidationError): + QueryFilter(col="user_id", opr="eq", value=1) + + def test_valid_status_filter_accepted(self): + """status is a valid filter column.""" + f = QueryFilter(col="status", opr="eq", value="success") + assert f.col == "status" + + def test_valid_database_id_filter_accepted(self): + """database_id is a valid filter column.""" + f = QueryFilter(col="database_id", opr="eq", value=1) + assert f.col == "database_id" + + def test_valid_schema_filter_accepted(self): + """schema is a valid filter column.""" + f = QueryFilter(col="schema", opr="eq", value="public") + assert f.col == "schema" + + +def create_mock_query( + query_id: int = 1, + sql: str = "SELECT * FROM table", + status: str = "success", + start_time: float = 1700000000.0, + end_time: float = 1700000001.0, + rows: int = 100, + database_id: int = 1, + schema: str = "public", + tab_name: str = "SQL Lab 1", + error_message: str | None = None, + client_id: str = "abc123", +) -> MagicMock: + """Factory function to create mock query objects with sensible defaults.""" + query = MagicMock() + query.id = query_id + query.sql = sql + query.status = status + query.start_time = start_time + query.end_time = end_time + query.rows = rows + query.database_id = database_id + query.schema = schema + query.tab_name = tab_name + query.error_message = error_message + query.client_id = client_id + query.limit = 1000 + query.progress = 100 + query.changed_on = None + return query + + +@pytest.fixture +def mcp_server(): + return mcp + + +@pytest.fixture(autouse=True) +def mock_auth(): + """Mock authentication for all tests.""" + from unittest.mock import Mock, patch + + with patch("superset.mcp_service.auth.get_user_from_request") as mock_get_user: + mock_user = Mock() + mock_user.id = 1 + mock_user.username = "admin" + mock_get_user.return_value = mock_user + yield mock_get_user + + +@patch("superset.daos.query.QueryDAO.list") +@pytest.mark.asyncio +async def test_list_queries_basic(mock_list, mcp_server): + """Test basic query listing functionality.""" + query = create_mock_query() + query._mapping = { + "id": query.id, + "sql": query.sql, + "status": query.status, + "start_time": query.start_time, + "database_id": query.database_id, + "schema": query.schema, + } + mock_list.return_value = ([query], 1) + async with Client(mcp_server) as client: + request = ListQueriesRequest(page=1, page_size=10) + result = await client.call_tool( + "list_queries", {"request": request.model_dump()} + ) + assert result.content is not None + data = json.loads(result.content[0].text) + assert data["queries"] is not None + assert len(data["queries"]) == 1 + assert data["queries"][0]["id"] == 1 + assert data["queries"][0]["status"] == "success" + + +@patch("superset.daos.query.QueryDAO.list") +@pytest.mark.asyncio +async def test_list_queries_with_status_filter(mock_list, mcp_server): + """Test query listing with status filter.""" + query = create_mock_query(status="failed", error_message="Syntax error") + query._mapping = { + "id": query.id, + "sql": query.sql, + "status": query.status, + "error_message": query.error_message, + } + mock_list.return_value = ([query], 1) + async with Client(mcp_server) as client: + request = ListQueriesRequest( + page=1, + page_size=10, + filters=[ + {"col": "status", "opr": "eq", "value": "failed"}, + ], + ) + result = await client.call_tool( + "list_queries", {"request": request.model_dump()} + ) + assert result.content is not None + data = json.loads(result.content[0].text) + assert data["queries"] is not None + assert len(data["queries"]) == 1 + assert data["queries"][0]["status"] == "failed" + + +@patch("superset.daos.query.QueryDAO.list") +@pytest.mark.asyncio +async def test_list_queries_default_page_size(mock_list, mcp_server): + """Test that default page size is 25 for query history.""" + mock_list.return_value = ([], 0) + async with Client(mcp_server) as client: + result = await client.call_tool("list_queries", {}) + assert result.content is not None + data = json.loads(result.content[0].text) + assert data["page_size"] == 25 + + +def test_list_queries_request_rejects_both_search_and_filters(): + """Cannot use search and filters simultaneously.""" + with pytest.raises(ValidationError): + ListQueriesRequest( + search="test", + filters=[{"col": "status", "opr": "eq", "value": "success"}], + ) + + +@patch("superset.daos.query.QueryDAO.find_by_id") +@pytest.mark.asyncio +async def test_get_query_info_basic(mock_find, mcp_server): + """Test basic get query info functionality.""" + query = create_mock_query() + mock_find.return_value = query + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_query_info", {"request": {"identifier": 1}} + ) + assert result.content is not None + data = json.loads(result.content[0].text) + assert data["id"] == 1 + assert data["status"] == "success" + assert data["database_id"] == 1 + + +@patch("superset.daos.query.QueryDAO.find_by_id") +@pytest.mark.asyncio +async def test_get_query_info_not_found(mock_find, mcp_server): + """Test get query info when query does not exist.""" + mock_find.return_value = None + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_query_info", {"request": {"identifier": 999}} + ) + assert result.data["error_type"] == "not_found" + + +@patch("superset.daos.query.QueryDAO.list") +@pytest.mark.asyncio +async def test_list_queries_empty(mock_list, mcp_server): + """Test query listing returns empty list when no results.""" + mock_list.return_value = ([], 0) + async with Client(mcp_server) as client: + request = ListQueriesRequest(page=1, page_size=10) + result = await client.call_tool( + "list_queries", {"request": request.model_dump()} + ) + assert result.content is not None + data = json.loads(result.content[0].text) + assert data["queries"] == [] + assert data["count"] == 0 + assert data["total_count"] == 0 + + +@patch("superset.daos.query.QueryDAO.list") +@pytest.mark.asyncio +async def test_list_queries_pagination_info(mock_list, mcp_server): + """Test that pagination info is correctly returned.""" + queries = [create_mock_query(query_id=i) for i in range(1, 4)] + for q in queries: + q._mapping = {"id": q.id, "sql": q.sql, "status": q.status} + mock_list.return_value = (queries, 100) + async with Client(mcp_server) as client: + request = ListQueriesRequest(page=1, page_size=3) + result = await client.call_tool( + "list_queries", {"request": request.model_dump()} + ) + data = json.loads(result.content[0].text) + assert data["total_count"] == 100 + assert data["page_size"] == 3 + assert data["has_next"] is True + assert data["has_previous"] is False + + +@patch("superset.daos.query.QueryDAO.list") +@pytest.mark.asyncio +async def test_list_queries_default_order_is_start_time_desc(mock_list, mcp_server): + """Test that default ordering is start_time descending.""" + mock_list.return_value = ([], 0) + async with Client(mcp_server) as client: + result = await client.call_tool("list_queries", {}) + assert result.content is not None + mock_list.assert_called_once() + call_kwargs = mock_list.call_args + assert call_kwargs.kwargs.get("order_column") == "start_time" + assert call_kwargs.kwargs.get("order_direction") == "desc" diff --git a/tests/unit_tests/mcp_service/saved_query/__init__.py b/tests/unit_tests/mcp_service/saved_query/__init__.py new file mode 100644 index 000000000000..13a83393a912 --- /dev/null +++ b/tests/unit_tests/mcp_service/saved_query/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/tests/unit_tests/mcp_service/saved_query/tool/__init__.py b/tests/unit_tests/mcp_service/saved_query/tool/__init__.py new file mode 100644 index 000000000000..13a83393a912 --- /dev/null +++ b/tests/unit_tests/mcp_service/saved_query/tool/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py b/tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py new file mode 100644 index 000000000000..2a89f341cb70 --- /dev/null +++ b/tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py @@ -0,0 +1,256 @@ +# 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. + + +import logging +from unittest.mock import MagicMock, patch + +import pytest +from fastmcp import Client +from pydantic import ValidationError + +from superset.mcp_service.app import mcp +from superset.mcp_service.saved_query.schemas import ( + ListSavedQueriesRequest, + SavedQueryFilter, +) +from superset.utils import json + +logging.basicConfig(level=logging.DEBUG) +logger = logging.getLogger(__name__) + + +class TestSavedQueryFilterSchema: + """Tests for SavedQueryFilter schema — filterable columns.""" + + def test_invalid_filter_column_rejected(self): + """Columns not in the Literal set must be rejected.""" + with pytest.raises(ValidationError): + SavedQueryFilter(col="not_a_real_column", opr="eq", value="test") + + def test_user_id_is_rejected_as_filter_column(self): + """user_id is an internal field and should not be a filter column.""" + with pytest.raises(ValidationError): + SavedQueryFilter(col="user_id", opr="eq", value=1) + + def test_valid_label_filter_accepted(self): + """label is a valid filter column.""" + f = SavedQueryFilter(col="label", opr="eq", value="my query") + assert f.col == "label" + + def test_valid_db_id_filter_accepted(self): + """db_id is a valid filter column.""" + f = SavedQueryFilter(col="db_id", opr="eq", value=1) + assert f.col == "db_id" + + def test_valid_schema_filter_accepted(self): + """schema is a valid filter column.""" + f = SavedQueryFilter(col="schema", opr="eq", value="public") + assert f.col == "schema" + + +def create_mock_saved_query( + saved_query_id: int = 1, + label: str = "My Query", + sql: str = "SELECT 1", + db_id: int = 1, + schema: str = "public", + description: str = "Test query", + uuid: str = "test-uuid-1234", +) -> MagicMock: + """Factory function to create mock saved query objects with sensible defaults.""" + saved_query = MagicMock() + saved_query.id = saved_query_id + saved_query.label = label + saved_query.sql = sql + saved_query.db_id = db_id + saved_query.schema = schema + saved_query.description = description + saved_query.uuid = uuid + saved_query.changed_on = None + saved_query.created_on = None + return saved_query + + +@pytest.fixture +def mcp_server(): + return mcp + + +@pytest.fixture(autouse=True) +def mock_auth(): + """Mock authentication for all tests.""" + from unittest.mock import Mock, patch + + with patch("superset.mcp_service.auth.get_user_from_request") as mock_get_user: + mock_user = Mock() + mock_user.id = 1 + mock_user.username = "admin" + mock_get_user.return_value = mock_user + yield mock_get_user + + +@patch("superset.daos.query.SavedQueryDAO.list") +@pytest.mark.asyncio +async def test_list_saved_queries_basic(mock_list, mcp_server): + """Test basic saved query listing functionality.""" + saved_query = create_mock_saved_query() + saved_query._mapping = { + "id": saved_query.id, + "label": saved_query.label, + "db_id": saved_query.db_id, + "schema": saved_query.schema, + "uuid": saved_query.uuid, + } + mock_list.return_value = ([saved_query], 1) + async with Client(mcp_server) as client: + request = ListSavedQueriesRequest(page=1, page_size=10) + result = await client.call_tool( + "list_saved_queries", {"request": request.model_dump()} + ) + assert result.content is not None + data = json.loads(result.content[0].text) + assert data["saved_queries"] is not None + assert len(data["saved_queries"]) == 1 + assert data["saved_queries"][0]["id"] == 1 + assert data["saved_queries"][0]["label"] == "My Query" + + +@patch("superset.daos.query.SavedQueryDAO.list") +@pytest.mark.asyncio +async def test_list_saved_queries_with_search(mock_list, mcp_server): + """Test saved query listing with search functionality.""" + saved_query = create_mock_saved_query(label="Production Query") + saved_query._mapping = { + "id": saved_query.id, + "label": saved_query.label, + } + mock_list.return_value = ([saved_query], 1) + async with Client(mcp_server) as client: + request = ListSavedQueriesRequest(page=1, page_size=10, search="Production") + result = await client.call_tool( + "list_saved_queries", {"request": request.model_dump()} + ) + assert result.content is not None + data = json.loads(result.content[0].text) + assert data["saved_queries"] is not None + assert len(data["saved_queries"]) == 1 + assert data["saved_queries"][0]["label"] == "Production Query" + + +@patch("superset.daos.query.SavedQueryDAO.list") +@pytest.mark.asyncio +async def test_list_saved_queries_with_filters(mock_list, mcp_server): + """Test saved query listing with filters.""" + saved_query = create_mock_saved_query(db_id=2) + saved_query._mapping = { + "id": saved_query.id, + "label": saved_query.label, + "db_id": saved_query.db_id, + } + mock_list.return_value = ([saved_query], 1) + async with Client(mcp_server) as client: + request = ListSavedQueriesRequest( + page=1, + page_size=10, + filters=[ + {"col": "db_id", "opr": "eq", "value": 2}, + ], + ) + result = await client.call_tool( + "list_saved_queries", {"request": request.model_dump()} + ) + assert result.content is not None + data = json.loads(result.content[0].text) + assert data["saved_queries"] is not None + assert len(data["saved_queries"]) == 1 + + +def test_list_saved_queries_request_rejects_both_search_and_filters(): + """Cannot use search and filters simultaneously.""" + with pytest.raises(ValidationError): + ListSavedQueriesRequest( + search="test", + filters=[{"col": "label", "opr": "eq", "value": "test"}], + ) + + +@patch("superset.daos.query.SavedQueryDAO.find_by_id") +@pytest.mark.asyncio +async def test_get_saved_query_info_basic(mock_find, mcp_server): + """Test basic get saved query info functionality.""" + saved_query = create_mock_saved_query() + mock_find.return_value = saved_query + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_saved_query_info", {"request": {"identifier": 1}} + ) + assert result.content is not None + data = json.loads(result.content[0].text) + assert data["id"] == 1 + assert data["label"] == "My Query" + assert data["sql"] == "SELECT 1" + assert data["db_id"] == 1 + + +@patch("superset.daos.query.SavedQueryDAO.find_by_id") +@pytest.mark.asyncio +async def test_get_saved_query_info_not_found(mock_find, mcp_server): + """Test get saved query info when saved query does not exist.""" + mock_find.return_value = None + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_saved_query_info", {"request": {"identifier": 999}} + ) + assert result.data["error_type"] == "not_found" + + +@patch("superset.daos.query.SavedQueryDAO.list") +@pytest.mark.asyncio +async def test_list_saved_queries_empty(mock_list, mcp_server): + """Test saved query listing returns empty list when no results.""" + mock_list.return_value = ([], 0) + async with Client(mcp_server) as client: + request = ListSavedQueriesRequest(page=1, page_size=10) + result = await client.call_tool( + "list_saved_queries", {"request": request.model_dump()} + ) + assert result.content is not None + data = json.loads(result.content[0].text) + assert data["saved_queries"] == [] + assert data["count"] == 0 + assert data["total_count"] == 0 + + +@patch("superset.daos.query.SavedQueryDAO.list") +@pytest.mark.asyncio +async def test_list_saved_queries_pagination_info(mock_list, mcp_server): + """Test that pagination info is correctly returned.""" + saved_queries = [create_mock_saved_query(saved_query_id=i) for i in range(1, 4)] + for sq in saved_queries: + sq._mapping = {"id": sq.id, "label": sq.label} + mock_list.return_value = (saved_queries, 25) + async with Client(mcp_server) as client: + request = ListSavedQueriesRequest(page=1, page_size=3) + result = await client.call_tool( + "list_saved_queries", {"request": request.model_dump()} + ) + data = json.loads(result.content[0].text) + assert data["total_count"] == 25 + assert data["page_size"] == 3 + assert data["has_next"] is True + assert data["has_previous"] is False From d813dd24c28509f9cc5efde44fe9007649b86910 Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Thu, 21 May 2026 08:07:04 +0000 Subject: [PATCH 02/10] feat(mcp): add sql to search columns for query/saved-query list tools; add test --- superset/mcp_service/app.py | 4 ++++ superset/mcp_service/query/tool/list_queries.py | 2 +- .../saved_query/tool/list_saved_queries.py | 2 +- .../saved_query/tool/test_saved_query_tools.py | 17 +++++++++++++++++ 4 files changed, 23 insertions(+), 2 deletions(-) diff --git a/superset/mcp_service/app.py b/superset/mcp_service/app.py index fbc7d5e5c143..8ecaaefdd203 100644 --- a/superset/mcp_service/app.py +++ b/superset/mcp_service/app.py @@ -155,6 +155,10 @@ def get_default_instructions( - execute_sql: Execute SQL queries and get results (requires database_id and SQL access) - save_sql_query: Save a SQL query to Saved Queries list (requires write access) - open_sql_lab_with_context: Generate SQL Lab URL with pre-filled sql +- list_saved_queries: List saved SQL queries with filtering and search (1-based pagination) +- get_saved_query_info: Get saved query details by ID or UUID +- list_queries: List SQL query history with filtering and search (most recent first) +- get_query_info: Get SQL query history details by ID Schema Discovery: - get_schema: Get schema metadata for chart/dataset/dashboard (columns, filters) diff --git a/superset/mcp_service/query/tool/list_queries.py b/superset/mcp_service/query/tool/list_queries.py index c39f1f6062dc..ae621de2cb64 100644 --- a/superset/mcp_service/query/tool/list_queries.py +++ b/superset/mcp_service/query/tool/list_queries.py @@ -104,7 +104,7 @@ def _serialize_query(obj: object, cols: list[str] | None) -> QueryInfo | None: item_serializer=_serialize_query, filter_type=QueryFilter, default_columns=DEFAULT_QUERY_COLUMNS, - search_columns=["tab_name"], + search_columns=["tab_name", "sql"], list_field_name="queries", output_list_schema=QueryList, all_columns=ALL_QUERY_COLUMNS, diff --git a/superset/mcp_service/saved_query/tool/list_saved_queries.py b/superset/mcp_service/saved_query/tool/list_saved_queries.py index d1820517ca85..2e26bf2ce18f 100644 --- a/superset/mcp_service/saved_query/tool/list_saved_queries.py +++ b/superset/mcp_service/saved_query/tool/list_saved_queries.py @@ -105,7 +105,7 @@ def _serialize_saved_query( item_serializer=_serialize_saved_query, filter_type=SavedQueryFilter, default_columns=DEFAULT_SAVED_QUERY_COLUMNS, - search_columns=["label", "description"], + search_columns=["label", "description", "sql"], list_field_name="saved_queries", output_list_schema=SavedQueryList, all_columns=ALL_SAVED_QUERY_COLUMNS, diff --git a/tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py b/tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py index 2a89f341cb70..da35d8b7607a 100644 --- a/tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py +++ b/tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py @@ -236,6 +236,23 @@ async def test_list_saved_queries_empty(mock_list, mcp_server): assert data["total_count"] == 0 +@patch("superset.daos.query.SavedQueryDAO.find_by_id") +@pytest.mark.asyncio +async def test_get_saved_query_info_by_uuid(mock_find, mcp_server): + """Test get saved query info by UUID string.""" + saved_query = create_mock_saved_query(uuid="a1b2c3d4-5678-90ab-cdef-1234567890ab") + mock_find.return_value = saved_query + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_saved_query_info", + {"request": {"identifier": "a1b2c3d4-5678-90ab-cdef-1234567890ab"}}, + ) + assert result.content is not None + data = json.loads(result.content[0].text) + assert data["id"] == 1 + assert data["uuid"] == "a1b2c3d4-5678-90ab-cdef-1234567890ab" + + @patch("superset.daos.query.SavedQueryDAO.list") @pytest.mark.asyncio async def test_list_saved_queries_pagination_info(mock_list, mcp_server): From 192954752b655a17abfe8568b6e40d1a2080eb05 Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Fri, 22 May 2026 20:18:42 +0000 Subject: [PATCH 03/10] fix(mcp): address review feedback on list_queries and list_saved_queries tools MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Drop sql from DEFAULT_QUERY_COLUMNS (25 rows × full SQL bodies is too heavy for default LLM responses; callers use select_columns or get_query_info to access SQL) - Add changed_on to SORTABLE_QUERY_COLUMNS for queries (column is indexed on the model, same treatment as saved queries) - Remove QueryError.create() and SavedQueryError.create() dead code - Fix list_queries instruction text in app.py: use '(1-based pagination)' to match the wording used by all other list tools - Add tests: select_columns field projection and invalid order_column rejection --- superset/mcp_service/app.py | 2 +- superset/mcp_service/query/schemas.py | 12 ++----- .../mcp_service/query/tool/list_queries.py | 2 +- superset/mcp_service/saved_query/schemas.py | 9 ------ .../query/tool/test_query_tools.py | 31 +++++++++++++++++++ 5 files changed, 35 insertions(+), 21 deletions(-) diff --git a/superset/mcp_service/app.py b/superset/mcp_service/app.py index 8ecaaefdd203..f929ec89d598 100644 --- a/superset/mcp_service/app.py +++ b/superset/mcp_service/app.py @@ -157,7 +157,7 @@ def get_default_instructions( - open_sql_lab_with_context: Generate SQL Lab URL with pre-filled sql - list_saved_queries: List saved SQL queries with filtering and search (1-based pagination) - get_saved_query_info: Get saved query details by ID or UUID -- list_queries: List SQL query history with filtering and search (most recent first) +- list_queries: List SQL query history with filtering and search (1-based pagination) - get_query_info: Get SQL query history details by ID Schema Discovery: diff --git a/superset/mcp_service/query/schemas.py b/superset/mcp_service/query/schemas.py index 07c4bdcab38d..214ae02a8b28 100644 --- a/superset/mcp_service/query/schemas.py +++ b/superset/mcp_service/query/schemas.py @@ -43,13 +43,14 @@ parse_json_or_model_list, ) -DEFAULT_QUERY_COLUMNS = ["id", "sql", "status", "start_time", "database_id", "schema"] +DEFAULT_QUERY_COLUMNS = ["id", "status", "start_time", "database_id", "schema"] SORTABLE_QUERY_COLUMNS = [ "id", "start_time", "end_time", "status", "database_id", + "changed_on", ] ALL_QUERY_COLUMNS = [ "id", @@ -249,15 +250,6 @@ class QueryError(BaseModel): timestamp: str | datetime | None = Field(None, description="Error timestamp") model_config = ConfigDict(ser_json_timedelta="iso8601") - @classmethod - def create(cls, error: str, error_type: str) -> "QueryError": - """Create a standardized QueryError with timestamp.""" - from datetime import datetime, timezone - - return cls( - error=error, error_type=error_type, timestamp=datetime.now(timezone.utc) - ) - class GetQueryInfoRequest(BaseModel): """Request schema for get_query_info with support for numeric ID only.""" diff --git a/superset/mcp_service/query/tool/list_queries.py b/superset/mcp_service/query/tool/list_queries.py index ae621de2cb64..6e8de1b32756 100644 --- a/superset/mcp_service/query/tool/list_queries.py +++ b/superset/mcp_service/query/tool/list_queries.py @@ -66,7 +66,7 @@ async def list_queries( Results are ordered by start_time descending (most recent first) by default. Sortable columns for order_column: id, start_time, end_time, status, - database_id + database_id, changed_on """ if ctx is None: raise RuntimeError("FastMCP context is required for list_queries") diff --git a/superset/mcp_service/saved_query/schemas.py b/superset/mcp_service/saved_query/schemas.py index c55298637e89..96afdefc6a37 100644 --- a/superset/mcp_service/saved_query/schemas.py +++ b/superset/mcp_service/saved_query/schemas.py @@ -232,15 +232,6 @@ class SavedQueryError(BaseModel): timestamp: str | datetime | None = Field(None, description="Error timestamp") model_config = ConfigDict(ser_json_timedelta="iso8601") - @classmethod - def create(cls, error: str, error_type: str) -> "SavedQueryError": - """Create a standardized SavedQueryError with timestamp.""" - from datetime import datetime, timezone - - return cls( - error=error, error_type=error_type, timestamp=datetime.now(timezone.utc) - ) - class GetSavedQueryInfoRequest(BaseModel): """Request schema for get_saved_query_info with support for ID or UUID.""" diff --git a/tests/unit_tests/mcp_service/query/tool/test_query_tools.py b/tests/unit_tests/mcp_service/query/tool/test_query_tools.py index 8e12d109a540..e907f05e8e61 100644 --- a/tests/unit_tests/mcp_service/query/tool/test_query_tools.py +++ b/tests/unit_tests/mcp_service/query/tool/test_query_tools.py @@ -269,3 +269,34 @@ async def test_list_queries_default_order_is_start_time_desc(mock_list, mcp_serv call_kwargs = mock_list.call_args assert call_kwargs.kwargs.get("order_column") == "start_time" assert call_kwargs.kwargs.get("order_direction") == "desc" + + +@patch("superset.daos.query.QueryDAO.list") +@pytest.mark.asyncio +async def test_list_queries_select_columns_projects_fields(mock_list, mcp_server): + """select_columns limits which fields appear in each query result.""" + query = create_mock_query() + query._mapping = {"id": query.id, "status": query.status} + mock_list.return_value = ([query], 1) + async with Client(mcp_server) as client: + request = ListQueriesRequest( + page=1, page_size=10, select_columns=["id", "status"] + ) + result = await client.call_tool( + "list_queries", {"request": request.model_dump()} + ) + data = json.loads(result.content[0].text) + assert data["queries"] is not None + q = data["queries"][0] + assert set(q.keys()) == {"id", "status"} + assert q["id"] == 1 + assert q["status"] == "success" + + +@pytest.mark.asyncio +async def test_list_queries_invalid_order_column_raises(mcp_server): + """order_column not in SORTABLE_QUERY_COLUMNS must be rejected.""" + request = ListQueriesRequest(page=1, page_size=10, order_column="tab_name") + async with Client(mcp_server) as client: + with pytest.raises(Exception, match="Invalid order_column"): + await client.call_tool("list_queries", {"request": request.model_dump()}) From 8dc800a7e504e87678a1664e42d3d96deacbb162 Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Tue, 26 May 2026 20:18:16 +0000 Subject: [PATCH 04/10] fix(mcp): address review feedback on list_queries and list_saved_queries tools - Narrow exception type in order_column test from Exception to ValueError - Add count assertion to saved query pagination test --- tests/unit_tests/mcp_service/query/tool/test_query_tools.py | 2 +- .../mcp_service/saved_query/tool/test_saved_query_tools.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/mcp_service/query/tool/test_query_tools.py b/tests/unit_tests/mcp_service/query/tool/test_query_tools.py index e907f05e8e61..439a5444bfef 100644 --- a/tests/unit_tests/mcp_service/query/tool/test_query_tools.py +++ b/tests/unit_tests/mcp_service/query/tool/test_query_tools.py @@ -298,5 +298,5 @@ async def test_list_queries_invalid_order_column_raises(mcp_server): """order_column not in SORTABLE_QUERY_COLUMNS must be rejected.""" request = ListQueriesRequest(page=1, page_size=10, order_column="tab_name") async with Client(mcp_server) as client: - with pytest.raises(Exception, match="Invalid order_column"): + with pytest.raises(ValueError, match="Invalid order_column"): await client.call_tool("list_queries", {"request": request.model_dump()}) diff --git a/tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py b/tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py index da35d8b7607a..bfd6ccbd859c 100644 --- a/tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py +++ b/tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py @@ -268,6 +268,7 @@ async def test_list_saved_queries_pagination_info(mock_list, mcp_server): ) data = json.loads(result.content[0].text) assert data["total_count"] == 25 + assert data["count"] == 3 assert data["page_size"] == 3 assert data["has_next"] is True assert data["has_previous"] is False From 2c742bb4fa7601d6eb7f6d5f837ad45d5d0ff33c Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Wed, 27 May 2026 15:21:33 +0000 Subject: [PATCH 05/10] fix(mcp): avoid leaking raw exception text in get_saved_query_info error response The full error details are already logged via ctx.error() to the server log; return a generic message to the client to avoid exposing internal DB errors. --- superset/mcp_service/saved_query/tool/get_saved_query_info.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/superset/mcp_service/saved_query/tool/get_saved_query_info.py b/superset/mcp_service/saved_query/tool/get_saved_query_info.py index 9b3a1be74b22..f1a4f19b3f13 100644 --- a/superset/mcp_service/saved_query/tool/get_saved_query_info.py +++ b/superset/mcp_service/saved_query/tool/get_saved_query_info.py @@ -123,7 +123,7 @@ async def get_saved_query_info( ) ) return SavedQueryError( - error=f"Failed to get saved query info: {str(e)}", + error="Failed to get saved query info", error_type="InternalError", timestamp=datetime.now(timezone.utc), ) From fdff27f20a9676d6d6177c51893bd6d68549d0bf Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Wed, 27 May 2026 18:17:36 +0000 Subject: [PATCH 06/10] fix(mcp): change list_queries default sort from start_time to changed_on start_time is a nullable epoch float that is not set for queries inserted outside of normal SQL Lab execution. Sorting by changed_on (always populated) gives stable, deterministic ordering regardless of how the query record was created. --- superset/mcp_service/query/tool/list_queries.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/superset/mcp_service/query/tool/list_queries.py b/superset/mcp_service/query/tool/list_queries.py index 6e8de1b32756..53bdbcf888e3 100644 --- a/superset/mcp_service/query/tool/list_queries.py +++ b/superset/mcp_service/query/tool/list_queries.py @@ -63,7 +63,8 @@ async def list_queries( Returns recent queries executed by the current user (or all queries for admins), including SQL text, status, timing, and database information. - Results are ordered by start_time descending (most recent first) by default. + Results are ordered by changed_on descending by default (start_time is not + always populated for all query records). Sortable columns for order_column: id, start_time, end_time, status, database_id, changed_on @@ -117,7 +118,7 @@ def _serialize_query(obj: object, cols: list[str] | None) -> QueryInfo | None: filters=request.filters, search=request.search, select_columns=request.select_columns, - order_column=request.order_column or "start_time", + order_column=request.order_column or "changed_on", order_direction=request.order_direction, page=max(request.page - 1, 0), page_size=request.page_size, From 163012b52fc505342e37ee121bc0be70d43ea673 Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Wed, 27 May 2026 18:26:52 +0000 Subject: [PATCH 07/10] fix(mcp): address review feedback on list_queries and list_saved_queries tools MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Fix test_list_queries_default_order_is_start_time_desc → changed_on after default sort column change - Sanitize get_query_info error response: return generic message instead of str(e) to match get_saved_query_info (prevents leaking internal error text) - Add test_get_query_info_internal_error and test_get_saved_query_info_internal_error covering the InternalError exception path in both get-tools - Add test_list_saved_queries_select_columns_projects_fields and test_list_saved_queries_invalid_order_column_raises to match query tool coverage --- .../mcp_service/query/tool/get_query_info.py | 2 +- .../query/tool/test_query_tools.py | 20 ++++++-- .../tool/test_saved_query_tools.py | 47 +++++++++++++++++++ 3 files changed, 65 insertions(+), 4 deletions(-) diff --git a/superset/mcp_service/query/tool/get_query_info.py b/superset/mcp_service/query/tool/get_query_info.py index dc94a947d6cd..f26dac2349ad 100644 --- a/superset/mcp_service/query/tool/get_query_info.py +++ b/superset/mcp_service/query/tool/get_query_info.py @@ -116,7 +116,7 @@ async def get_query_info( ) ) return QueryError( - error=f"Failed to get query info: {str(e)}", + error="Failed to get query info", error_type="InternalError", timestamp=datetime.now(timezone.utc), ) diff --git a/tests/unit_tests/mcp_service/query/tool/test_query_tools.py b/tests/unit_tests/mcp_service/query/tool/test_query_tools.py index 439a5444bfef..da680453facc 100644 --- a/tests/unit_tests/mcp_service/query/tool/test_query_tools.py +++ b/tests/unit_tests/mcp_service/query/tool/test_query_tools.py @@ -259,15 +259,15 @@ async def test_list_queries_pagination_info(mock_list, mcp_server): @patch("superset.daos.query.QueryDAO.list") @pytest.mark.asyncio -async def test_list_queries_default_order_is_start_time_desc(mock_list, mcp_server): - """Test that default ordering is start_time descending.""" +async def test_list_queries_default_order_is_changed_on_desc(mock_list, mcp_server): + """Test that default ordering is changed_on descending.""" mock_list.return_value = ([], 0) async with Client(mcp_server) as client: result = await client.call_tool("list_queries", {}) assert result.content is not None mock_list.assert_called_once() call_kwargs = mock_list.call_args - assert call_kwargs.kwargs.get("order_column") == "start_time" + assert call_kwargs.kwargs.get("order_column") == "changed_on" assert call_kwargs.kwargs.get("order_direction") == "desc" @@ -300,3 +300,17 @@ async def test_list_queries_invalid_order_column_raises(mcp_server): async with Client(mcp_server) as client: with pytest.raises(ValueError, match="Invalid order_column"): await client.call_tool("list_queries", {"request": request.model_dump()}) + + +@patch("superset.daos.query.QueryDAO.find_by_id") +@pytest.mark.asyncio +async def test_get_query_info_internal_error(mock_find, mcp_server): + """When an unexpected exception is raised, get_query_info returns InternalError.""" + mock_find.side_effect = RuntimeError("unexpected db failure") + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_query_info", {"request": {"identifier": 1}} + ) + data = json.loads(result.content[0].text) + assert data["error_type"] == "InternalError" + assert data["error"] == "Failed to get query info" diff --git a/tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py b/tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py index bfd6ccbd859c..9e75963f791b 100644 --- a/tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py +++ b/tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py @@ -272,3 +272,50 @@ async def test_list_saved_queries_pagination_info(mock_list, mcp_server): assert data["page_size"] == 3 assert data["has_next"] is True assert data["has_previous"] is False + + +@patch("superset.daos.query.SavedQueryDAO.list") +@pytest.mark.asyncio +async def test_list_saved_queries_select_columns_projects_fields(mock_list, mcp_server): + """select_columns limits which fields appear in each saved query result.""" + saved_query = create_mock_saved_query() + saved_query._mapping = {"id": saved_query.id, "label": saved_query.label} + mock_list.return_value = ([saved_query], 1) + async with Client(mcp_server) as client: + request = ListSavedQueriesRequest( + page=1, page_size=10, select_columns=["id", "label"] + ) + result = await client.call_tool( + "list_saved_queries", {"request": request.model_dump()} + ) + data = json.loads(result.content[0].text) + assert data["saved_queries"] is not None + sq = data["saved_queries"][0] + assert set(sq.keys()) == {"id", "label"} + assert sq["id"] == 1 + assert sq["label"] == "My Query" + + +@pytest.mark.asyncio +async def test_list_saved_queries_invalid_order_column_raises(mcp_server): + """order_column not in SORTABLE_SAVED_QUERY_COLUMNS must be rejected.""" + request = ListSavedQueriesRequest(page=1, page_size=10, order_column="sql") + async with Client(mcp_server) as client: + with pytest.raises(ValueError, match="Invalid order_column"): + await client.call_tool( + "list_saved_queries", {"request": request.model_dump()} + ) + + +@patch("superset.daos.query.SavedQueryDAO.find_by_id") +@pytest.mark.asyncio +async def test_get_saved_query_info_internal_error(mock_find, mcp_server): + """Unexpected exception in get_saved_query_info returns InternalError.""" + mock_find.side_effect = RuntimeError("unexpected db failure") + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_saved_query_info", {"request": {"identifier": 1}} + ) + data = json.loads(result.content[0].text) + assert data["error_type"] == "InternalError" + assert data["error"] == "Failed to get saved query info" From ccd5c69e854cf9e7ae5d8d281f114f827dab57ba Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Wed, 27 May 2026 18:48:07 +0000 Subject: [PATCH 08/10] =?UTF-8?q?test(mcp):=20fix=20ValueError=20=E2=86=92?= =?UTF-8?q?=20ToolError=20in=20invalid=20order=5Fcolumn=20tests?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit fastmcp client wraps tool ValueError as ToolError at the protocol level, so pytest.raises(ValueError) never matched. Use ToolError instead. --- tests/unit_tests/mcp_service/query/tool/test_query_tools.py | 3 ++- .../mcp_service/saved_query/tool/test_saved_query_tools.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/mcp_service/query/tool/test_query_tools.py b/tests/unit_tests/mcp_service/query/tool/test_query_tools.py index da680453facc..828e8d917844 100644 --- a/tests/unit_tests/mcp_service/query/tool/test_query_tools.py +++ b/tests/unit_tests/mcp_service/query/tool/test_query_tools.py @@ -21,6 +21,7 @@ import pytest from fastmcp import Client +from fastmcp.exceptions import ToolError from pydantic import ValidationError from superset.mcp_service.app import mcp @@ -298,7 +299,7 @@ async def test_list_queries_invalid_order_column_raises(mcp_server): """order_column not in SORTABLE_QUERY_COLUMNS must be rejected.""" request = ListQueriesRequest(page=1, page_size=10, order_column="tab_name") async with Client(mcp_server) as client: - with pytest.raises(ValueError, match="Invalid order_column"): + with pytest.raises(ToolError, match="Invalid order_column"): await client.call_tool("list_queries", {"request": request.model_dump()}) diff --git a/tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py b/tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py index 9e75963f791b..025bc18a370f 100644 --- a/tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py +++ b/tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py @@ -21,6 +21,7 @@ import pytest from fastmcp import Client +from fastmcp.exceptions import ToolError from pydantic import ValidationError from superset.mcp_service.app import mcp @@ -301,7 +302,7 @@ async def test_list_saved_queries_invalid_order_column_raises(mcp_server): """order_column not in SORTABLE_SAVED_QUERY_COLUMNS must be rejected.""" request = ListSavedQueriesRequest(page=1, page_size=10, order_column="sql") async with Client(mcp_server) as client: - with pytest.raises(ValueError, match="Invalid order_column"): + with pytest.raises(ToolError, match="Invalid order_column"): await client.call_tool( "list_saved_queries", {"request": request.model_dump()} ) From 98b986b51376a7bbf03a66170641c97dae0bc335 Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Thu, 28 May 2026 16:16:04 +0000 Subject: [PATCH 09/10] feat(mcp): expand query and saved_query schemas per review feedback - QueryInfo: add executed_sql, catalog, user_id fields - QueryFilter.col: add user_id and start_time as filterable columns - SavedQueryInfo: add catalog and last_run fields - SavedQueryFilter.col: add catalog and created_by_fk as filterable columns - ALL_QUERY_COLUMNS / ALL_SAVED_QUERY_COLUMNS: include new fields - Tests: update filter schema tests to reflect new valid columns, update mock factories with new fields --- superset/mcp_service/query/schemas.py | 15 +++++++++++-- superset/mcp_service/saved_query/schemas.py | 10 ++++++++- .../query/tool/test_query_tools.py | 21 ++++++++++++++----- .../tool/test_saved_query_tools.py | 17 ++++++++++++++- 4 files changed, 54 insertions(+), 9 deletions(-) diff --git a/superset/mcp_service/query/schemas.py b/superset/mcp_service/query/schemas.py index 214ae02a8b28..6c8a9241bc2a 100644 --- a/superset/mcp_service/query/schemas.py +++ b/superset/mcp_service/query/schemas.py @@ -55,18 +55,21 @@ ALL_QUERY_COLUMNS = [ "id", "sql", + "executed_sql", "status", "start_time", "end_time", "rows", "database_id", "schema", + "catalog", "tab_name", "error_message", "client_id", "limit", "progress", "changed_on", + "user_id", ] DEFAULT_QUERY_PAGE_SIZE = 25 @@ -80,7 +83,7 @@ class QueryFilter(ColumnOperator): value: The value to filter by (type depends on col and opr). """ - col: Literal["status", "database_id", "schema"] = Field( + col: Literal["status", "database_id", "schema", "user_id", "start_time"] = Field( ..., description="Column to filter on.", ) @@ -95,7 +98,10 @@ class QueryFilter(ColumnOperator): class QueryInfo(BaseModel): id: int | None = Field(None, description="Query ID") - sql: str | None = Field(None, description="SQL query text") + sql: str | None = Field(None, description="SQL query text as submitted") + executed_sql: str | None = Field( + None, description="Actual SQL executed after templating/CTAS rewriting" + ) status: str | None = Field(None, description="Query execution status") start_time: float | None = Field( None, description="Query start time (seconds since epoch)" @@ -106,6 +112,7 @@ class QueryInfo(BaseModel): rows: int | None = Field(None, description="Number of rows returned or affected") database_id: int | None = Field(None, description="Database connection ID") schema: str | None = Field(None, description="Database schema name") + catalog: str | None = Field(None, description="Database catalog name") tab_name: str | None = Field(None, description="SQL Lab tab name") error_message: str | None = Field(None, description="Error message if query failed") client_id: str | None = Field(None, description="Client-assigned query identifier") @@ -114,6 +121,7 @@ class QueryInfo(BaseModel): changed_on: str | datetime | None = Field( None, description="Last modification timestamp" ) + user_id: int | None = Field(None, description="ID of the user who ran the query") model_config = ConfigDict( from_attributes=True, ser_json_timedelta="iso8601", @@ -267,16 +275,19 @@ def serialize_query_object(query: Any) -> QueryInfo | None: return QueryInfo( id=getattr(query, "id", None), sql=getattr(query, "sql", None), + executed_sql=getattr(query, "executed_sql", None), status=getattr(query, "status", None), start_time=getattr(query, "start_time", None), end_time=getattr(query, "end_time", None), rows=getattr(query, "rows", None), database_id=getattr(query, "database_id", None), schema=getattr(query, "schema", None), + catalog=getattr(query, "catalog", None), tab_name=getattr(query, "tab_name", None), error_message=getattr(query, "error_message", None), client_id=getattr(query, "client_id", None), limit=getattr(query, "limit", None), progress=getattr(query, "progress", None), changed_on=getattr(query, "changed_on", None), + user_id=getattr(query, "user_id", None), ) diff --git a/superset/mcp_service/saved_query/schemas.py b/superset/mcp_service/saved_query/schemas.py index 96afdefc6a37..b38c6215076c 100644 --- a/superset/mcp_service/saved_query/schemas.py +++ b/superset/mcp_service/saved_query/schemas.py @@ -57,11 +57,13 @@ "label", "db_id", "schema", + "catalog", "uuid", "sql", "description", "changed_on", "created_on", + "last_run", ] @@ -73,7 +75,7 @@ class SavedQueryFilter(ColumnOperator): value: The value to filter by (type depends on col and opr). """ - col: Literal["label", "db_id", "schema"] = Field( + col: Literal["label", "db_id", "schema", "catalog", "created_by_fk"] = Field( ..., description="Column to filter on.", ) @@ -93,11 +95,15 @@ class SavedQueryInfo(BaseModel): sql: str | None = Field(None, description="SQL query text") db_id: int | None = Field(None, description="Database connection ID") schema: str | None = Field(None, description="Database schema name") + catalog: str | None = Field(None, description="Database catalog name") description: str | None = Field(None, description="User-provided description") changed_on: str | datetime | None = Field( None, description="Last modification timestamp" ) created_on: str | datetime | None = Field(None, description="Creation timestamp") + last_run: str | datetime | None = Field( + None, description="Timestamp of last execution" + ) model_config = ConfigDict( from_attributes=True, ser_json_timedelta="iso8601", @@ -255,7 +261,9 @@ def serialize_saved_query_object(saved_query: Any) -> SavedQueryInfo | None: sql=getattr(saved_query, "sql", None), db_id=getattr(saved_query, "db_id", None), schema=getattr(saved_query, "schema", None), + catalog=getattr(saved_query, "catalog", None), description=getattr(saved_query, "description", None), changed_on=getattr(saved_query, "changed_on", None), created_on=getattr(saved_query, "created_on", None), + last_run=getattr(saved_query, "last_run", None), ) diff --git a/tests/unit_tests/mcp_service/query/tool/test_query_tools.py b/tests/unit_tests/mcp_service/query/tool/test_query_tools.py index 828e8d917844..3bf370a26278 100644 --- a/tests/unit_tests/mcp_service/query/tool/test_query_tools.py +++ b/tests/unit_tests/mcp_service/query/tool/test_query_tools.py @@ -43,11 +43,6 @@ def test_invalid_filter_column_rejected(self): with pytest.raises(ValidationError): QueryFilter(col="not_a_real_column", opr="eq", value="test") - def test_user_id_is_rejected_as_filter_column(self): - """user_id is an internal field and should not be a filter column.""" - with pytest.raises(ValidationError): - QueryFilter(col="user_id", opr="eq", value=1) - def test_valid_status_filter_accepted(self): """status is a valid filter column.""" f = QueryFilter(col="status", opr="eq", value="success") @@ -63,36 +58,52 @@ def test_valid_schema_filter_accepted(self): f = QueryFilter(col="schema", opr="eq", value="public") assert f.col == "schema" + def test_valid_user_id_filter_accepted(self): + """user_id filter enables admin-level filtering by user.""" + f = QueryFilter(col="user_id", opr="eq", value=42) + assert f.col == "user_id" + + def test_valid_start_time_filter_accepted(self): + """start_time filter enables time-range queries.""" + f = QueryFilter(col="start_time", opr="gt", value=1700000000.0) + assert f.col == "start_time" + def create_mock_query( query_id: int = 1, sql: str = "SELECT * FROM table", + executed_sql: str | None = None, status: str = "success", start_time: float = 1700000000.0, end_time: float = 1700000001.0, rows: int = 100, database_id: int = 1, schema: str = "public", + catalog: str | None = None, tab_name: str = "SQL Lab 1", error_message: str | None = None, client_id: str = "abc123", + user_id: int | None = 1, ) -> MagicMock: """Factory function to create mock query objects with sensible defaults.""" query = MagicMock() query.id = query_id query.sql = sql + query.executed_sql = executed_sql query.status = status query.start_time = start_time query.end_time = end_time query.rows = rows query.database_id = database_id query.schema = schema + query.catalog = catalog query.tab_name = tab_name query.error_message = error_message query.client_id = client_id query.limit = 1000 query.progress = 100 query.changed_on = None + query.user_id = user_id return query diff --git a/tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py b/tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py index 025bc18a370f..7deb3afedf5b 100644 --- a/tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py +++ b/tests/unit_tests/mcp_service/saved_query/tool/test_saved_query_tools.py @@ -17,6 +17,7 @@ import logging +from datetime import datetime from unittest.mock import MagicMock, patch import pytest @@ -44,7 +45,7 @@ def test_invalid_filter_column_rejected(self): SavedQueryFilter(col="not_a_real_column", opr="eq", value="test") def test_user_id_is_rejected_as_filter_column(self): - """user_id is an internal field and should not be a filter column.""" + """user_id is not a filter column; use created_by_fk instead.""" with pytest.raises(ValidationError): SavedQueryFilter(col="user_id", opr="eq", value=1) @@ -63,6 +64,16 @@ def test_valid_schema_filter_accepted(self): f = SavedQueryFilter(col="schema", opr="eq", value="public") assert f.col == "schema" + def test_valid_catalog_filter_accepted(self): + """catalog is a valid filter column.""" + f = SavedQueryFilter(col="catalog", opr="eq", value="my_catalog") + assert f.col == "catalog" + + def test_valid_created_by_fk_filter_accepted(self): + """created_by_fk enables filtering by the owner user ID.""" + f = SavedQueryFilter(col="created_by_fk", opr="eq", value=42) + assert f.col == "created_by_fk" + def create_mock_saved_query( saved_query_id: int = 1, @@ -70,8 +81,10 @@ def create_mock_saved_query( sql: str = "SELECT 1", db_id: int = 1, schema: str = "public", + catalog: str | None = None, description: str = "Test query", uuid: str = "test-uuid-1234", + last_run: datetime | None = None, ) -> MagicMock: """Factory function to create mock saved query objects with sensible defaults.""" saved_query = MagicMock() @@ -80,10 +93,12 @@ def create_mock_saved_query( saved_query.sql = sql saved_query.db_id = db_id saved_query.schema = schema + saved_query.catalog = catalog saved_query.description = description saved_query.uuid = uuid saved_query.changed_on = None saved_query.created_on = None + saved_query.last_run = last_run return saved_query From 2232717d70a086b5599bc28f24fe125f704d0655 Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Thu, 28 May 2026 16:23:09 +0000 Subject: [PATCH 10/10] refactor(mcp): derive columns_available from schema model_fields Use QueryInfo.model_fields.keys() / SavedQueryInfo.model_fields.keys() as the columns_available source in list_queries and list_saved_queries, rather than manually maintained ALL_QUERY_COLUMNS / ALL_SAVED_QUERY_COLUMNS constants. This ensures the advertised columns always exactly match what the response schema can serialize, preventing future drift between the constant and the schema definition. Pattern mirrors the fix applied to list_reports in #40348 per reviewer feedback from richardfogaca. --- superset/mcp_service/query/tool/list_queries.py | 3 +-- superset/mcp_service/saved_query/tool/list_saved_queries.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/superset/mcp_service/query/tool/list_queries.py b/superset/mcp_service/query/tool/list_queries.py index 53bdbcf888e3..64a97e0cb451 100644 --- a/superset/mcp_service/query/tool/list_queries.py +++ b/superset/mcp_service/query/tool/list_queries.py @@ -30,7 +30,6 @@ from superset.extensions import event_logger from superset.mcp_service.mcp_core import ModelListCore from superset.mcp_service.query.schemas import ( - ALL_QUERY_COLUMNS, DEFAULT_QUERY_COLUMNS, ListQueriesRequest, QueryError, @@ -108,7 +107,7 @@ def _serialize_query(obj: object, cols: list[str] | None) -> QueryInfo | None: search_columns=["tab_name", "sql"], list_field_name="queries", output_list_schema=QueryList, - all_columns=ALL_QUERY_COLUMNS, + all_columns=list(QueryInfo.model_fields.keys()), sortable_columns=SORTABLE_QUERY_COLUMNS, logger=logger, ) diff --git a/superset/mcp_service/saved_query/tool/list_saved_queries.py b/superset/mcp_service/saved_query/tool/list_saved_queries.py index 2e26bf2ce18f..d4f09bb60fce 100644 --- a/superset/mcp_service/saved_query/tool/list_saved_queries.py +++ b/superset/mcp_service/saved_query/tool/list_saved_queries.py @@ -30,7 +30,6 @@ from superset.extensions import event_logger from superset.mcp_service.mcp_core import ModelListCore from superset.mcp_service.saved_query.schemas import ( - ALL_SAVED_QUERY_COLUMNS, DEFAULT_SAVED_QUERY_COLUMNS, ListSavedQueriesRequest, SavedQueryError, @@ -108,7 +107,7 @@ def _serialize_saved_query( search_columns=["label", "description", "sql"], list_field_name="saved_queries", output_list_schema=SavedQueryList, - all_columns=ALL_SAVED_QUERY_COLUMNS, + all_columns=list(SavedQueryInfo.model_fields.keys()), sortable_columns=SORTABLE_SAVED_QUERY_COLUMNS, logger=logger, )