You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by mi...@apache.org on 2026/04/13 19:05:29 UTC
(superset) 24/45: fix(mcp): resolve null fields in list_datasets, list_databases, and save_sql_query (#39206)
This is an automated email from the ASF dual-hosted git repository.
michaelsmolina pushed a commit to branch 6.1
in repository https://gitbox.apache.org/repos/asf/superset.git
commit 25fb1188835037d09630efe1df00ef4360196088
Author: Amin Ghadersohi <am...@gmail.com>
AuthorDate: Wed Apr 8 18:39:56 2026 -0400
fix(mcp): resolve null fields in list_datasets, list_databases, and save_sql_query (#39206)
(cherry picked from commit 1bde6f3bfda0b99de425e2c6b4da786b29833908)
---
superset/daos/base.py | 16 +++++++---
superset/mcp_service/dataset/tool/list_datasets.py | 5 +++
superset/mcp_service/sql_lab/schemas.py | 36 +++++++++++++++++++---
.../mcp_service/dataset/tool/test_dataset_tools.py | 8 ++++-
4 files changed, 55 insertions(+), 10 deletions(-)
diff --git a/superset/daos/base.py b/superset/daos/base.py
index 2e5290776e3..fe72f1c63d0 100644
--- a/superset/daos/base.py
+++ b/superset/daos/base.py
@@ -623,6 +623,7 @@ class BaseDAO(CoreBaseDAO[T], Generic[T]):
column_attrs = []
relationship_loads = []
+ needs_full_model = False
if columns is None:
columns = []
for name in columns:
@@ -634,11 +635,16 @@ class BaseDAO(CoreBaseDAO[T], Generic[T]):
column_attrs.append(attr)
elif isinstance(prop, RelationshipProperty):
relationship_loads.append(joinedload(attr))
- # Ignore properties and other non-queryable attributes
-
- if relationship_loads:
- # If any relationships are requested, query the full model
- # but don't add the joins yet - we'll add them after counting
+ else:
+ # Python @property or other descriptor — requires a full
+ # model instance (Row objects don't support descriptors)
+ needs_full_model = True
+
+ if relationship_loads or needs_full_model:
+ # Need full model for relationships or Python @property access.
+ # Do NOT apply load_only() here — @property descriptors and
+ # serializers may access columns beyond the explicitly requested
+ # set (e.g., Slice.datasource_type accessed during serialization).
query = data_model.session.query(cls.model_cls)
elif column_attrs:
# Only columns requested
diff --git a/superset/mcp_service/dataset/tool/list_datasets.py b/superset/mcp_service/dataset/tool/list_datasets.py
index 4804e453757..e4d2d3bd427 100644
--- a/superset/mcp_service/dataset/tool/list_datasets.py
+++ b/superset/mcp_service/dataset/tool/list_datasets.py
@@ -44,10 +44,15 @@ from superset.mcp_service.mcp_core import ModelListCore
logger = logging.getLogger(__name__)
# Minimal defaults for reduced token usage - users can request more via select_columns
+# NOTE: "database" (relationship) is included so the DAO eagerly loads it
+# via joinedload, which avoids N+1 lazy-load queries when the serializer
+# accesses dataset.database.name (via the database_name @property).
DEFAULT_DATASET_COLUMNS = [
"id",
"table_name",
"schema",
+ "database_name",
+ "database",
"description",
"certified_by",
"certification_details",
diff --git a/superset/mcp_service/sql_lab/schemas.py b/superset/mcp_service/sql_lab/schemas.py
index b20bdc419e4..f328a752fc7 100644
--- a/superset/mcp_service/sql_lab/schemas.py
+++ b/superset/mcp_service/sql_lab/schemas.py
@@ -17,9 +17,33 @@
"""Schemas for SQL Lab MCP tools."""
-from typing import Any
+from typing import Any, Dict
-from pydantic import AliasChoices, BaseModel, ConfigDict, Field, field_validator
+from pydantic import (
+ AliasChoices,
+ BaseModel,
+ ConfigDict,
+ Field,
+ field_validator,
+ model_serializer,
+)
+
+
+class _SchemaFieldNormalizer(BaseModel):
+ """Mixin that renames schema_name → schema in JSON output.
+
+ Pydantic serializes using field names by default, but the MCP/API
+ convention uses the alias ``schema``. Adding this as a base class
+ to any response model that carries a ``schema_name`` field with
+ ``alias="schema"`` keeps the JSON key consistent.
+ """
+
+ @model_serializer(mode="wrap", when_used="json")
+ def _normalize_schema_field(self, serializer: Any, _info: Any) -> Dict[str, Any]:
+ data = serializer(self)
+ if "schema_name" in data:
+ data["schema"] = data.pop("schema_name")
+ return data
class ExecuteSqlRequest(BaseModel):
@@ -151,6 +175,8 @@ class ExecuteSqlResponse(BaseModel):
class SaveSqlQueryRequest(BaseModel):
"""Request schema for saving a SQL query."""
+ model_config = ConfigDict(populate_by_name=True)
+
database_id: int = Field(
..., description="Database connection ID the query runs against"
)
@@ -189,9 +215,11 @@ class SaveSqlQueryRequest(BaseModel):
return v.strip()
-class SaveSqlQueryResponse(BaseModel):
+class SaveSqlQueryResponse(_SchemaFieldNormalizer):
"""Response schema for a saved SQL query."""
+ model_config = ConfigDict(populate_by_name=True)
+
id: int = Field(..., description="Saved query ID")
label: str = Field(..., description="Query name")
sql: str = Field(..., description="SQL query text")
@@ -231,7 +259,7 @@ class OpenSqlLabRequest(BaseModel):
title: str | None = Field(None, description="Title for the SQL Lab tab/query")
-class SqlLabResponse(BaseModel):
+class SqlLabResponse(_SchemaFieldNormalizer):
"""Response schema for SQL Lab URL generation."""
model_config = ConfigDict(populate_by_name=True)
diff --git a/tests/unit_tests/mcp_service/dataset/tool/test_dataset_tools.py b/tests/unit_tests/mcp_service/dataset/tool/test_dataset_tools.py
index 04735064914..3be98887b1c 100644
--- a/tests/unit_tests/mcp_service/dataset/tool/test_dataset_tools.py
+++ b/tests/unit_tests/mcp_service/dataset/tool/test_dataset_tools.py
@@ -1293,6 +1293,8 @@ class TestDatasetDefaultColumnFiltering:
"id",
"table_name",
"schema",
+ "database_name",
+ "database",
"description",
"certified_by",
"certification_details",
@@ -1380,10 +1382,15 @@ class TestDatasetDefaultColumnFiltering:
dataset_item = data["datasets"][0]
# Verify ONLY default columns are present in the response item
+ # Note: "database" is in DEFAULT_DATASET_COLUMNS to trigger eager
+ # loading of the relationship, but it's not a field on DatasetInfo
+ # so it won't appear in the serialized output. "database_name" IS
+ # a field and will appear.
expected_keys = {
"id",
"table_name",
"schema",
+ "database_name",
"description",
"certified_by",
"certification_details",
@@ -1402,7 +1409,6 @@ class TestDatasetDefaultColumnFiltering:
# Verify non-default columns are NOT present (not even with null values)
non_default_columns = [
- "database_name",
"changed_by",
"columns",
"metrics",