This is an automated email from the ASF dual-hosted git repository. aminghadersohi pushed a commit to branch oss-40340 in repository https://gitbox.apache.org/repos/asf/superset.git
commit 178fe56c9c72f104e42916de074c21eef7e0940e Author: Amin Ghadersohi <[email protected]> AuthorDate: Fri May 22 01:44:41 2026 +0000 fix(mcp): fix create_dataset CI failures - schemas.py: restore full apache/master version and add CreateDatasetRequest (previous cherry-pick used an older shorter version missing helper functions _sanitize_dataset_info_for_llm_context, _humanize_timestamp, etc.) - create_dataset.py: remove parse_request decorator (not in apache/master yet) --- superset/mcp_service/dataset/schemas.py | 405 +++++++++++++++++++-- .../mcp_service/dataset/tool/create_dataset.py | 2 - 2 files changed, 384 insertions(+), 23 deletions(-) diff --git a/superset/mcp_service/dataset/schemas.py b/superset/mcp_service/dataset/schemas.py index d14b2fb4a11..0bbc4061f8c 100644 --- a/superset/mcp_service/dataset/schemas.py +++ b/superset/mcp_service/dataset/schemas.py @@ -24,21 +24,35 @@ from __future__ import annotations from datetime import datetime from typing import Annotated, Any, Dict, List, Literal +import humanize from pydantic import ( BaseModel, ConfigDict, Field, + field_validator, model_serializer, model_validator, PositiveInt, ) from superset.daos.base import ColumnOperator, ColumnOperatorEnum -from superset.mcp_service.common.cache_schemas import MetadataCacheControl +from superset.mcp_service.chart.schemas import DataColumn, PerformanceMetadata +from superset.mcp_service.common.cache_schemas import ( + CacheStatus, + CreatedByMeMixin, + MetadataCacheControl, + OwnedByMeMixin, + QueryCacheControl, +) +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, TagInfo, - UserInfo, +) +from superset.mcp_service.utils import ( + escape_llm_context_delimiters, + sanitize_for_llm_context, ) from superset.utils import json @@ -85,7 +99,11 @@ class TableColumnInfo(BaseModel): class SqlMetricInfo(BaseModel): - metric_name: str = Field(..., description="Metric name") + metric_name: str = Field( + ..., + description="Saved metric name. In chart configs, reference as " + '{"name": "<metric_name>", "saved_metric": true}.', + ) verbose_name: str | None = Field(None, description="Verbose name") expression: str | None = Field(None, description="SQL expression") description: str | None = Field(None, description="Metric description") @@ -98,22 +116,23 @@ class DatasetInfo(BaseModel): schema_name: str | None = Field(None, description="Schema name", alias="schema") database_name: str | None = Field(None, description="Database name") description: str | None = Field(None, description="Dataset description") - changed_by: str | None = Field(None, description="Last modifier (username)") + certified_by: str | None = Field( + None, description="Name of the person or team who certified this dataset" + ) + certification_details: str | None = Field( + None, description="Certification details or reason" + ) changed_on: str | datetime | None = Field( None, description="Last modification timestamp" ) changed_on_humanized: str | None = Field( None, description="Humanized modification time" ) - created_by: str | None = Field(None, description="Dataset creator (username)") created_on: str | datetime | None = Field(None, description="Creation timestamp") created_on_humanized: str | None = Field( None, description="Humanized creation time" ) tags: List[TagInfo] = Field(default_factory=list, description="Dataset tags") - owners: List[UserInfo] = Field( - default_factory=list, description="DatasetInfo owners" - ) is_virtual: bool | None = Field( None, description="Whether the dataset is virtual (uses SQL)" ) @@ -134,7 +153,9 @@ class DatasetInfo(BaseModel): default_factory=list, description="Columns in the dataset" ) metrics: List[SqlMetricInfo] = Field( - default_factory=list, description="Metrics in the dataset" + default_factory=list, + description="Saved metrics (pre-defined aggregations). " + "NOT columns — use saved_metric=true in chart configs.", ) is_favorite: bool | None = Field( None, description="Whether this dataset is favorited by the current user" @@ -145,7 +166,7 @@ class DatasetInfo(BaseModel): populate_by_name=True, # Allow both 'schema' (alias) and 'schema_name' (field) ) - @model_serializer(mode="wrap", when_used="json") + @model_serializer(mode="wrap") def _filter_fields_by_context(self, serializer: Any, info: Any) -> Dict[str, Any]: """Filter fields based on serialization context. @@ -153,16 +174,18 @@ class DatasetInfo(BaseModel): Otherwise, include all fields (default behavior). """ # Get full serialization - data = serializer(self) + data = filter_user_directory_fields(serializer(self)) + + # Normalize alias: Pydantic serializes as 'schema_name' (field name) + # but the DAO column and API convention is 'schema' + if "schema_name" in data: + data["schema"] = data.pop("schema_name") # Check if we have a context with select_columns if info.context and isinstance(info.context, dict): select_columns = info.context.get("select_columns") if select_columns: - # Handle alias: 'schema' -> 'schema_name' requested_fields = set(select_columns) - if "schema" in requested_fields: - requested_fields.add("schema_name") # Filter to only requested fields return {k: v for k, v in data.items() if k in requested_fields} @@ -205,7 +228,7 @@ class DatasetList(BaseModel): model_config = ConfigDict(ser_json_timedelta="iso8601") -class ListDatasetsRequest(MetadataCacheControl): +class ListDatasetsRequest(OwnedByMeMixin, CreatedByMeMixin, MetadataCacheControl): """Request schema for list_datasets with clear, unambiguous types.""" filters: Annotated[ @@ -247,13 +270,18 @@ class ListDatasetsRequest(MetadataCacheControl): Field(default=1, description="Page number for pagination (1-based)"), ] page_size: Annotated[ - PositiveInt, Field(default=10, description="Number of items per page") + int, + Field( + default=DEFAULT_PAGE_SIZE, + gt=0, + le=MAX_PAGE_SIZE, + description=f"Number of items per page (max {MAX_PAGE_SIZE})", + ), ] @model_validator(mode="after") def validate_search_and_filters(self) -> "ListDatasetsRequest": - """Prevent using both search and filters simultaneously to avoid query - conflicts.""" + """Prevent using both search and filters simultaneously.""" if self.search and self.filters: raise ValueError( "Cannot use both 'search' and 'filters' parameters simultaneously. " @@ -269,12 +297,22 @@ class DatasetError(BaseModel): timestamp: str | datetime | None = Field(None, description="Error timestamp") model_config = ConfigDict(ser_json_timedelta="iso8601") + @field_validator("error") + @classmethod + def sanitize_error_for_llm_context(cls, value: str) -> str: + """Wrap error text before it is exposed to LLM context.""" + return sanitize_for_llm_context(value, field_path=("error",)) + @classmethod def create(cls, error: str, error_type: str) -> "DatasetError": """Create a standardized DatasetError with timestamp.""" - from datetime import datetime + from datetime import datetime, timezone - return cls(error=error, error_type=error_type, timestamp=datetime.now()) + return cls( + error=error, + error_type=error_type, + timestamp=datetime.now(timezone.utc), + ) class GetDatasetInfoRequest(MetadataCacheControl): @@ -307,14 +345,339 @@ class CreateDatasetRequest(BaseModel): List[int] | None, Field( default=None, - description="Optional list of owner user IDs. Defaults to the calling user.", + description="Optional list of owner user IDs. " + "Defaults to the calling user.", + ), + ] + + +class CreateVirtualDatasetRequest(BaseModel): + """Request schema for create_virtual_dataset.""" + + model_config = ConfigDict(populate_by_name=True) + + database_id: int = Field( + ..., + description="ID of the database connection to use. " + "Use list_databases to find valid IDs.", + ) + sql: str = Field( + ..., + description="SQL query to save as a virtual dataset. " + "Can be a JOIN, CTE, aggregation, or any valid SELECT.", + ) + dataset_name: str = Field( + ..., + min_length=1, + max_length=250, + description="Name for the new virtual dataset.", + ) + schema_name: str | None = Field( + None, + alias="schema", + description="Schema to associate with the dataset (optional).", + ) + catalog: str | None = Field( + None, + description="Catalog to associate with the dataset (optional).", + ) + description: str | None = Field( + None, + description="Human-readable description of the dataset (optional).", + ) + + @field_validator("sql") + @classmethod + def sql_must_not_be_empty(cls, v: str) -> str: + if not v.strip(): + raise ValueError("sql must not be empty") + return v.strip() + + @field_validator("dataset_name") + @classmethod + def dataset_name_must_not_be_empty(cls, v: str) -> str: + if not v.strip(): + raise ValueError("dataset_name must not be empty") + return v.strip() + + +class CreateVirtualDatasetResponse(BaseModel): + """Response schema for create_virtual_dataset.""" + + id: int | None = Field( + None, + description="Dataset ID. Pass this as dataset_id to generate_chart " + "or generate_explore_link. None if creation failed.", + ) + dataset_name: str = Field(..., description="Name of the created dataset.") + sql: str = Field(..., description="SQL query stored in the dataset.") + database_id: int = Field(..., description="Database ID used.") + columns: List[str] = Field( + default_factory=list, + description="Column names available for charting. " + "Use these when building chart configs.", + ) + url: str | None = Field( + None, + description="URL to view/edit the dataset in Superset. None if failed.", + ) + error: str | None = Field( + None, + description="Error message if creation failed, otherwise null.", + ) + + +VALID_FILTER_OPS = Literal[ + "==", + "!=", + ">", + "<", + ">=", + "<=", + "LIKE", + "NOT LIKE", + "ILIKE", + "NOT ILIKE", + "IN", + "NOT IN", + "IS NULL", + "IS NOT NULL", + "IS TRUE", + "IS FALSE", + "TEMPORAL_RANGE", +] + + +class QueryDatasetFilter(BaseModel): + """A single filter condition for dataset queries.""" + + col: str = Field(..., description="Column name to filter on") + op: VALID_FILTER_OPS = Field( + ..., + description=( + 'Filter operator. Use "==" for equals, "!=" for not equals, ' + '"IN" / "NOT IN" for membership, "IS NULL" / "IS NOT NULL", ' + '"LIKE" for pattern matching, "TEMPORAL_RANGE" for time filters.' + ), + ) + val: Any = Field( + default=None, + description="Filter value (omit for IS NULL/IS NOT NULL)", + ) + + +class QueryDatasetRequest(QueryCacheControl): + """Request schema for query_dataset tool.""" + + dataset_id: int | str = Field( + ..., + description="Dataset identifier — numeric ID or UUID string.", + ) + metrics: List[str] = Field( + default_factory=list, + description=( + "Saved metric names to compute (e.g. ['count', 'avg_revenue']). " + "Use get_dataset_info to discover available metrics." + ), + ) + columns: List[str] = Field( + default_factory=list, + description=( + "Column/dimension names for GROUP BY or SELECT " + "(e.g. ['category', 'region']). " + "Use get_dataset_info to discover available columns." + ), + ) + filters: List[QueryDatasetFilter] = Field( + default_factory=list, + description=( + 'Filter conditions (e.g. [{"col": "status", "op": "==", "val": "active"}]).' + ), + ) + time_range: str | None = Field( + default=None, + description=( + "Time range filter (e.g. 'Last 7 days', 'Last month', " + "'2024-01-01 : 2024-12-31'). Requires a temporal column " + "on the dataset." + ), + ) + time_column: str | None = Field( + default=None, + description=( + "Temporal column to apply time_range to. " + "Defaults to the dataset's main datetime column." ), + ) + order_by: List[str] | None = Field( + default=None, + description="Column or metric names to sort results by.", + ) + order_desc: bool = Field( + default=True, + description="Sort descending (True) or ascending (False).", + ) + row_limit: int = Field( + default=1000, + ge=1, + le=50000, + description="Maximum number of rows to return (default 1000, max 50000).", + ) + + @model_validator(mode="after") + def validate_metrics_or_columns(self) -> "QueryDatasetRequest": + """At least one of metrics or columns must be provided.""" + if not self.metrics and not self.columns: + raise ValueError( + "At least one of 'metrics' or 'columns' must be provided. " + "Use get_dataset_info to discover available metrics and columns." + ) + return self + + +class QueryDatasetResponse(BaseModel): + """Response schema for query_dataset tool.""" + + model_config = ConfigDict(ser_json_timedelta="iso8601") + + dataset_id: int = Field(..., description="Dataset ID") + dataset_name: str = Field(..., description="Dataset name") + columns: List[DataColumn] = Field( + default_factory=list, description="Column metadata for returned data" + ) + data: List[Dict[str, Any]] = Field( + default_factory=list, description="Query result rows" + ) + row_count: int = Field(0, description="Number of rows returned") + total_rows: int | None = Field( + None, description="Total row count from the query engine" + ) + summary: str = Field("", description="Human-readable summary of the results") + performance: PerformanceMetadata | None = Field( + None, description="Query performance metadata" + ) + cache_status: CacheStatus | None = Field( + None, description="Cache hit/miss information" + ) + applied_filters: List[QueryDatasetFilter] = Field( + default_factory=list, description="Filters that were applied to the query" + ) + warnings: List[str] = Field( + default_factory=list, description="Any warnings encountered during execution" + ) + + +def _parse_json_field(obj: Any, field_name: str) -> Dict[str, Any] | None: + """Parse a field that may be stored as a JSON string into a dict.""" + value = getattr(obj, field_name, None) + if isinstance(value, str): + try: + parsed = json.loads(value) + if isinstance(parsed, dict): + return parsed + except (ValueError, TypeError): + pass + return None + return value + + +def _humanize_timestamp(dt: datetime | None) -> str | None: + """Convert a datetime to a humanized string like '2 hours ago'.""" + if dt is None: + return None + return humanize.naturaltime(datetime.now() - dt) + + +def _sanitize_dataset_info_for_llm_context(dataset_info: DatasetInfo) -> DatasetInfo: + """Wrap dataset read-path descriptive fields before LLM exposure.""" + payload = dataset_info.model_dump(mode="python") + + for field_name in ("description", "certified_by", "certification_details", "sql"): + payload[field_name] = sanitize_for_llm_context( + payload.get(field_name), + field_path=(field_name,), + ) + + for field_name in ("table_name", "schema_name", "database_name", "schema_perm"): + payload[field_name] = escape_llm_context_delimiters(payload.get(field_name)) + + payload["extra"] = sanitize_for_llm_context( + payload.get("extra"), + field_path=("extra",), + excluded_field_names=frozenset(), + ) + + for field_name in ("params", "template_params"): + payload[field_name] = sanitize_for_llm_context( + payload.get(field_name), + field_path=(field_name,), + excluded_field_names=frozenset(), + ) + + payload["columns"] = [ + { + **column, + "column_name": escape_llm_context_delimiters( + column.get("column_name"), + ), + "description": sanitize_for_llm_context( + column.get("description"), + field_path=("columns", str(index), "description"), + ), + "verbose_name": sanitize_for_llm_context( + column.get("verbose_name"), + field_path=("columns", str(index), "verbose_name"), + ), + } + for index, column in enumerate(payload.get("columns", [])) + ] + + payload["metrics"] = [ + { + **metric, + "metric_name": escape_llm_context_delimiters( + metric.get("metric_name"), + ), + "expression": sanitize_for_llm_context( + metric.get("expression"), + field_path=("metrics", str(index), "expression"), + ), + "description": sanitize_for_llm_context( + metric.get("description"), + field_path=("metrics", str(index), "description"), + ), + "verbose_name": sanitize_for_llm_context( + metric.get("verbose_name"), + field_path=("metrics", str(index), "verbose_name"), + ), + } + for index, metric in enumerate(payload.get("metrics", [])) ] + payload["tags"] = [ + { + **tag, + "name": sanitize_for_llm_context( + tag.get("name"), + field_path=("tags", str(index), "name"), + ), + "description": sanitize_for_llm_context( + tag.get("description"), + field_path=("tags", str(index), "description"), + ), + } + for index, tag in enumerate(payload.get("tags", [])) + ] + + return DatasetInfo.model_validate(payload) + def serialize_dataset_object(dataset: Any) -> DatasetInfo | None: if not dataset: return None + + from superset.mcp_service.utils.url_utils import get_superset_base_url + params = getattr(dataset, "params", None) if isinstance(params, str): try: diff --git a/superset/mcp_service/dataset/tool/create_dataset.py b/superset/mcp_service/dataset/tool/create_dataset.py index 0dfb8443240..6484f8ab36d 100644 --- a/superset/mcp_service/dataset/tool/create_dataset.py +++ b/superset/mcp_service/dataset/tool/create_dataset.py @@ -37,14 +37,12 @@ from superset.mcp_service.dataset.schemas import ( DatasetInfo, serialize_dataset_object, ) -from superset.mcp_service.utils.schema_utils import parse_request logger = logging.getLogger(__name__) @mcp.tool(tags=["mutate"]) @mcp_auth_hook -@parse_request(CreateDatasetRequest) def create_dataset( request: CreateDatasetRequest, ctx: Context ) -> DatasetInfo | DatasetError:
