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/15 19:42:01 UTC
(superset) 04/10: feat(mcp): add create_virtual_dataset tool to save SQL queries as datasets (#39279)
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 94dc92df02f4d11f9b16fdb586ac2166a535a096
Author: Gabriel Torres Ruiz <ga...@gmail.com>
AuthorDate: Wed Apr 15 13:04:32 2026 -0300
feat(mcp): add create_virtual_dataset tool to save SQL queries as datasets (#39279)
(cherry picked from commit 18d6feb49943bf0f23a968171d1af0b569b49cc6)
---
superset/mcp_service/app.py | 9 +
superset/mcp_service/dataset/schemas.py | 77 ++++++
superset/mcp_service/dataset/tool/__init__.py | 2 +
.../dataset/tool/create_virtual_dataset.py | 137 ++++++++++
.../mcp_service/dataset/tool/test_dataset_tools.py | 293 ++++++++++++++++++++-
5 files changed, 517 insertions(+), 1 deletion(-)
diff --git a/superset/mcp_service/app.py b/superset/mcp_service/app.py
index c38f88e19ee..b29cd121477 100644
--- a/superset/mcp_service/app.py
+++ b/superset/mcp_service/app.py
@@ -61,6 +61,7 @@ Database Connections:
Dataset Management:
- list_datasets: List datasets with advanced filters (1-based pagination)
- get_dataset_info: Get detailed dataset information by ID (includes columns/metrics)
+- create_virtual_dataset: Save a SQL query as a virtual dataset for charting
Chart Management:
- list_charts: List charts with advanced filters (1-based pagination)
@@ -115,6 +116,13 @@ To explore data with SQL:
3. save_sql_query(database_id, label, sql) -> save query for later reuse
4. open_sql_lab_with_context(database_id) -> open SQL Lab UI
+To chart from a SQL query (JOIN, CTE, aggregation):
+1. execute_sql(database_id, sql) -> verify the query returns expected data
+2. Ask the user if they want to save it as a dataset
+3. create_virtual_dataset(database_id, sql, dataset_name) -> save as chartable dataset
+4. generate_explore_link(dataset_id, config) or generate_chart(dataset_id, config)
+ (use the column names returned by create_virtual_dataset)
+
generate_explore_link vs generate_chart:
- Use generate_explore_link for exploration (no permanent chart created)
- Use generate_chart with save_chart=True only when user wants to save permanently
@@ -428,6 +436,7 @@ from superset.mcp_service.database.tool import ( # noqa: F401, E402
list_databases,
)
from superset.mcp_service.dataset.tool import ( # noqa: F401, E402
+ create_virtual_dataset,
get_dataset_info,
list_datasets,
)
diff --git a/superset/mcp_service/dataset/schemas.py b/superset/mcp_service/dataset/schemas.py
index ad5f33fbba6..ba34b55cdcb 100644
--- a/superset/mcp_service/dataset/schemas.py
+++ b/superset/mcp_service/dataset/schemas.py
@@ -29,6 +29,7 @@ from pydantic import (
BaseModel,
ConfigDict,
Field,
+ field_validator,
model_serializer,
model_validator,
PositiveInt,
@@ -300,6 +301,82 @@ class GetDatasetInfoRequest(MetadataCacheControl):
]
+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.",
+ )
+
+
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)
diff --git a/superset/mcp_service/dataset/tool/__init__.py b/superset/mcp_service/dataset/tool/__init__.py
index 4446d2173ff..d396ce9e2d1 100644
--- a/superset/mcp_service/dataset/tool/__init__.py
+++ b/superset/mcp_service/dataset/tool/__init__.py
@@ -15,10 +15,12 @@
# specific language governing permissions and limitations
# under the License.
+from .create_virtual_dataset import create_virtual_dataset
from .get_dataset_info import get_dataset_info
from .list_datasets import list_datasets
__all__ = [
+ "create_virtual_dataset",
"list_datasets",
"get_dataset_info",
]
diff --git a/superset/mcp_service/dataset/tool/create_virtual_dataset.py b/superset/mcp_service/dataset/tool/create_virtual_dataset.py
new file mode 100644
index 00000000000..272509b1572
--- /dev/null
+++ b/superset/mcp_service/dataset/tool/create_virtual_dataset.py
@@ -0,0 +1,137 @@
+# 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 typing import Any
+
+from fastmcp import Context
+from superset_core.mcp.decorators import tool, ToolAnnotations
+
+from superset.extensions import event_logger
+from superset.mcp_service.dataset.schemas import (
+ CreateVirtualDatasetRequest,
+ CreateVirtualDatasetResponse,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@tool(
+ tags=["mutate"],
+ class_permission_name="Dataset",
+ method_permission_name="write",
+ annotations=ToolAnnotations(
+ title="Create virtual dataset from SQL",
+ readOnlyHint=False,
+ destructiveHint=False,
+ ),
+)
+async def create_virtual_dataset(
+ request: CreateVirtualDatasetRequest, ctx: Context
+) -> CreateVirtualDatasetResponse:
+ """Save a SQL query as a virtual dataset so it can be charted.
+
+ Use this tool when a user wants to visualize data from a SQL query
+ (e.g., a JOIN or complex aggregation) that doesn't map to a single
+ physical table.
+
+ Workflow:
+ 1. Call this tool with the SQL query and a dataset name
+ 2. Use the returned ``id`` as the ``dataset_id`` in generate_chart or
+ generate_explore_link
+ 3. Use the returned ``columns`` list to pick columns for the chart config
+ """
+ await ctx.info(
+ "Creating virtual dataset: database_id=%s, dataset_name=%r"
+ % (request.database_id, request.dataset_name)
+ )
+
+ try:
+ from superset.commands.dataset.create import CreateDatasetCommand
+ from superset.commands.dataset.exceptions import (
+ DatasetCreateFailedError,
+ DatasetInvalidError,
+ )
+ from superset.mcp_service.utils.url_utils import get_superset_base_url
+
+ # Create the virtual dataset — CreateDatasetCommand enforces access control
+ # and validates that the database exists (raises DatasetInvalidError otherwise)
+ with event_logger.log_context(action="mcp.create_virtual_dataset.create"):
+ properties: dict[str, Any] = {
+ "database": request.database_id,
+ "table_name": request.dataset_name,
+ "sql": request.sql,
+ }
+ if request.schema_name is not None:
+ properties["schema"] = request.schema_name
+ if request.catalog is not None:
+ properties["catalog"] = request.catalog
+ if request.description is not None:
+ properties["description"] = request.description
+
+ dataset = CreateDatasetCommand(properties).run()
+
+ # Build response
+ columns = [col.column_name for col in dataset.columns]
+ dataset_url = (
+ f"{get_superset_base_url()}"
+ f"/explore/?datasource_type=table&datasource_id={dataset.id}"
+ )
+
+ await ctx.info(
+ "Virtual dataset created: id=%s, dataset_name=%r, columns=%s"
+ % (dataset.id, dataset.table_name, columns)
+ )
+
+ return CreateVirtualDatasetResponse(
+ id=dataset.id,
+ dataset_name=dataset.table_name,
+ sql=request.sql,
+ database_id=request.database_id,
+ columns=columns,
+ url=dataset_url,
+ )
+
+ except DatasetInvalidError as exc:
+ messages = exc.normalized_messages()
+ await ctx.error("Virtual dataset validation failed: %s" % (messages,))
+ return CreateVirtualDatasetResponse(
+ id=None,
+ dataset_name=request.dataset_name,
+ sql=request.sql,
+ database_id=request.database_id,
+ columns=[],
+ url=None,
+ error=str(messages),
+ )
+ except DatasetCreateFailedError as exc:
+ await ctx.error("Virtual dataset creation failed: %s" % (str(exc),))
+ return CreateVirtualDatasetResponse(
+ id=None,
+ dataset_name=request.dataset_name,
+ sql=request.sql,
+ database_id=request.database_id,
+ columns=[],
+ url=None,
+ error=f"Failed to create dataset: {exc}",
+ )
+ except Exception as exc:
+ await ctx.error(
+ "Unexpected error creating virtual dataset: %s: %s"
+ % (type(exc).__name__, str(exc))
+ )
+ raise
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 3be98887b1c..c64801df22f 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
@@ -25,7 +25,10 @@ from fastmcp import Client
from fastmcp.exceptions import ToolError
from superset.mcp_service.app import mcp
-from superset.mcp_service.dataset.schemas import ListDatasetsRequest
+from superset.mcp_service.dataset.schemas import (
+ CreateVirtualDatasetRequest,
+ ListDatasetsRequest,
+)
from superset.utils import json
logging.basicConfig(level=logging.DEBUG)
@@ -1506,3 +1509,291 @@ class TestDatasetSortableColumns:
data = json.loads(result.content[0].text)
assert data["count"] == 0
assert data["datasets"] == []
+
+
+# ---------------------------------------------------------------------------
+# create_virtual_dataset tests
+# ---------------------------------------------------------------------------
+
+
+def _make_mock_db(database_name: str = "examples") -> MagicMock:
+ """Create a mock database object with a configurable database name."""
+ db = MagicMock()
+ db.database_name = database_name
+ return db
+
+
+def _make_mock_virtual_dataset(
+ id: int = 21,
+ table_name: str = "Customer Revenue",
+ column_names: list[str] | None = None,
+) -> MagicMock:
+ """Create a mock virtual dataset object with configurable columns."""
+ if column_names is None:
+ column_names = ["name", "revenue"]
+ dataset = MagicMock()
+ dataset.id = id
+ dataset.table_name = table_name
+ dataset.columns = [MagicMock(column_name=c) for c in column_names]
+ return dataset
+
+
+# --- Schema tests ---
+
+
+def test_create_virtual_dataset_request_valid() -> None:
+ req = CreateVirtualDatasetRequest(
+ database_id=1,
+ sql="SELECT a.name, COUNT(*) FROM a JOIN b ON a.id = b.a_id GROUP BY a.name",
+ dataset_name="Customer Revenue",
+ )
+ assert req.database_id == 1
+ assert req.dataset_name == "Customer Revenue"
+ assert req.schema_name is None
+
+
+def test_create_virtual_dataset_request_empty_sql_fails() -> None:
+ from pydantic import ValidationError
+
+ with pytest.raises(ValidationError, match="sql must not be empty"):
+ CreateVirtualDatasetRequest(database_id=1, sql=" ", dataset_name="Test")
+
+
+def test_create_virtual_dataset_request_empty_name_fails() -> None:
+ from pydantic import ValidationError
+
+ with pytest.raises(ValidationError, match="dataset_name must not be empty"):
+ CreateVirtualDatasetRequest(database_id=1, sql="SELECT 1", dataset_name=" ")
+
+
+def test_create_virtual_dataset_request_name_too_long() -> None:
+ from pydantic import ValidationError
+
+ with pytest.raises(ValidationError):
+ CreateVirtualDatasetRequest(
+ database_id=1, sql="SELECT 1", dataset_name="a" * 251
+ )
+
+
+def test_create_virtual_dataset_request_optional_fields() -> None:
+ req = CreateVirtualDatasetRequest(
+ database_id=1,
+ sql="SELECT 1",
+ dataset_name="Test",
+ schema_name="public",
+ catalog="main",
+ description="A virtual dataset",
+ )
+ assert req.schema_name == "public"
+ assert req.catalog == "main"
+ assert req.description == "A virtual dataset"
+
+
+# --- Tool logic tests ---
+
+
+@pytest.mark.asyncio
+async def test_create_virtual_dataset_success(mcp_server: object) -> None:
+ """Happy path: dataset created, columns and URL returned."""
+ mock_dataset = _make_mock_virtual_dataset(
+ id=21, table_name="Customer Revenue", column_names=["name", "revenue"]
+ )
+ mock_command = MagicMock()
+ mock_command.run.return_value = mock_dataset
+
+ with (
+ patch(
+ "superset.commands.dataset.create.CreateDatasetCommand",
+ return_value=mock_command,
+ ),
+ patch(
+ "superset.mcp_service.utils.url_utils.get_superset_base_url",
+ return_value="http://localhost:8088",
+ ),
+ ):
+ async with Client(mcp_server) as client:
+ sql = (
+ "SELECT a.name, SUM(b.revenue) FROM a"
+ " JOIN b ON a.id = b.a_id GROUP BY a.name"
+ )
+ request = CreateVirtualDatasetRequest(
+ database_id=1,
+ sql=sql,
+ dataset_name="Customer Revenue",
+ )
+ result = await client.call_tool(
+ "create_virtual_dataset", {"request": request.model_dump()}
+ )
+ data = json.loads(result.content[0].text)
+
+ assert data["id"] == 21
+ assert data["dataset_name"] == "Customer Revenue"
+ assert data["columns"] == ["name", "revenue"]
+ assert "datasource_type=table" in data["url"]
+ assert "datasource_id=21" in data["url"]
+ assert data["error"] is None
+
+
+@pytest.mark.asyncio
+async def test_create_virtual_dataset_db_not_found(mcp_server: object) -> None:
+ """When the database ID does not exist, CreateDatasetCommand raises
+ DatasetInvalidError containing DatabaseNotFoundValidationError."""
+ from superset.commands.dataset.exceptions import (
+ DatabaseNotFoundValidationError,
+ DatasetInvalidError,
+ )
+
+ invalid_exc = DatasetInvalidError()
+ invalid_exc.append(DatabaseNotFoundValidationError())
+ mock_command = MagicMock()
+ mock_command.run.side_effect = invalid_exc
+
+ with patch(
+ "superset.commands.dataset.create.CreateDatasetCommand",
+ return_value=mock_command,
+ ):
+ async with Client(mcp_server) as client:
+ request = CreateVirtualDatasetRequest(
+ database_id=999, sql="SELECT 1", dataset_name="Test"
+ )
+ result = await client.call_tool(
+ "create_virtual_dataset", {"request": request.model_dump()}
+ )
+ data = json.loads(result.content[0].text)
+
+ assert data["id"] is None
+ assert data["columns"] == []
+ assert data["error"] is not None
+
+
+@pytest.mark.asyncio
+async def test_create_virtual_dataset_invalid_error(mcp_server: object) -> None:
+ """DatasetInvalidError is caught and returned as an error response."""
+ from marshmallow.exceptions import ValidationError as MarshmallowValidationError
+
+ from superset.commands.dataset.exceptions import DatasetInvalidError
+
+ invalid_exc = DatasetInvalidError()
+ invalid_exc.append(
+ MarshmallowValidationError(
+ {"table_name": ["Dataset with this name already exists"]}
+ )
+ )
+ mock_command = MagicMock()
+ mock_command.run.side_effect = invalid_exc
+
+ with patch(
+ "superset.commands.dataset.create.CreateDatasetCommand",
+ return_value=mock_command,
+ ):
+ async with Client(mcp_server) as client:
+ request = CreateVirtualDatasetRequest(
+ database_id=1, sql="SELECT 1", dataset_name="Test"
+ )
+ result = await client.call_tool(
+ "create_virtual_dataset", {"request": request.model_dump()}
+ )
+ data = json.loads(result.content[0].text)
+
+ assert data["id"] is None
+ assert data["columns"] == []
+ assert data["error"] is not None
+
+
+@pytest.mark.asyncio
+async def test_create_virtual_dataset_create_failed(mcp_server: object) -> None:
+ """DatasetCreateFailedError is caught and returned as an error response."""
+ from superset.commands.dataset.exceptions import DatasetCreateFailedError
+
+ mock_command = MagicMock()
+ mock_command.run.side_effect = DatasetCreateFailedError()
+
+ with patch(
+ "superset.commands.dataset.create.CreateDatasetCommand",
+ return_value=mock_command,
+ ):
+ async with Client(mcp_server) as client:
+ request = CreateVirtualDatasetRequest(
+ database_id=1, sql="SELECT 1", dataset_name="Test"
+ )
+ result = await client.call_tool(
+ "create_virtual_dataset", {"request": request.model_dump()}
+ )
+ data = json.loads(result.content[0].text)
+
+ assert data["id"] is None
+ assert data["columns"] == []
+ assert data["error"] is not None
+ assert "Failed to create dataset" in data["error"]
+
+
+@pytest.mark.asyncio
+async def test_create_virtual_dataset_permission_denied(mcp_server: object) -> None:
+ """SQL access denied surfaces as DatasetInvalidError with id=None."""
+ from superset.commands.dataset.exceptions import (
+ DatasetDataAccessIsNotAllowed,
+ DatasetInvalidError,
+ )
+
+ invalid_exc = DatasetInvalidError()
+ invalid_exc.append(DatasetDataAccessIsNotAllowed("Access denied to schema public"))
+ mock_command = MagicMock()
+ mock_command.run.side_effect = invalid_exc
+
+ with patch(
+ "superset.commands.dataset.create.CreateDatasetCommand",
+ return_value=mock_command,
+ ):
+ async with Client(mcp_server) as client:
+ request = CreateVirtualDatasetRequest(
+ database_id=1,
+ sql="SELECT * FROM sensitive_table",
+ dataset_name="Restricted",
+ )
+ result = await client.call_tool(
+ "create_virtual_dataset", {"request": request.model_dump()}
+ )
+ data = json.loads(result.content[0].text)
+
+ assert data["id"] is None
+ assert data["columns"] == []
+ assert data["error"] is not None
+
+
+@pytest.mark.asyncio
+async def test_create_virtual_dataset_optional_fields_forwarded(
+ mcp_server: object,
+) -> None:
+ """schema_name, catalog, and description are forwarded to CreateDatasetCommand."""
+ mock_dataset = _make_mock_virtual_dataset(column_names=["col1"])
+ mock_command_instance = MagicMock()
+ mock_command_instance.run.return_value = mock_dataset
+ mock_command_cls = MagicMock(return_value=mock_command_instance)
+
+ with (
+ patch(
+ "superset.commands.dataset.create.CreateDatasetCommand",
+ mock_command_cls,
+ ),
+ patch(
+ "superset.mcp_service.utils.url_utils.get_superset_base_url",
+ return_value="http://localhost:8088",
+ ),
+ ):
+ async with Client(mcp_server) as client:
+ request = CreateVirtualDatasetRequest(
+ database_id=1,
+ sql="SELECT col1 FROM t",
+ dataset_name="My Dataset",
+ schema="public",
+ catalog="main",
+ description="A test dataset",
+ )
+ await client.call_tool(
+ "create_virtual_dataset", {"request": request.model_dump()}
+ )
+
+ props = mock_command_cls.call_args[0][0]
+ assert props["schema"] == "public"
+ assert props["catalog"] == "main"
+ assert props["description"] == "A test dataset"