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/06/16 03:25:37 UTC
(superset) branch master updated: feat(mcp): add create_dataset tool to register physical tables as datasets (#40340)
This is an automated email from the ASF dual-hosted git repository.
aminghadersohi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new 7cb4990403f feat(mcp): add create_dataset tool to register physical tables as datasets (#40340)
7cb4990403f is described below
commit 7cb4990403f570ce8aa02973f6c3dd65eba5a835
Author: Amin Ghadersohi <am...@gmail.com>
AuthorDate: Mon Jun 15 23:25:29 2026 -0400
feat(mcp): add create_dataset tool to register physical tables as datasets (#40340)
Co-authored-by: Claude Sonnet 4.6 <no...@anthropic.com>
---
superset/mcp_service/app.py | 16 +-
superset/mcp_service/dataset/schemas.py | 76 ++-
superset/mcp_service/dataset/tool/__init__.py | 4 +-
.../mcp_service/dataset/tool/create_dataset.py | 171 ++++++
.../dataset/tool/test_create_dataset.py | 585 +++++++++++++++++++++
5 files changed, 841 insertions(+), 11 deletions(-)
diff --git a/superset/mcp_service/app.py b/superset/mcp_service/app.py
index 41568df0429..a45e1578d7e 100644
--- a/superset/mcp_service/app.py
+++ b/superset/mcp_service/app.py
@@ -162,6 +162,7 @@ Alerts & Reports:
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_dataset: Register a physical table as a dataset against an existing DB connection (requires write access)
- create_virtual_dataset: Save a SQL query as a virtual dataset for charting (requires write access)
- query_dataset: Query a dataset using its semantic layer (saved metrics, dimensions, filters) without needing a saved chart
@@ -422,7 +423,7 @@ Input format:
{_feature_availability}Permission Awareness:
{_instance_info_role_bullet}- ALWAYS check the user's roles BEFORE suggesting write operations (creating datasets,
charts, or dashboards). SQL execution is a separate permission — see execute_sql below.
-- Write tools (generate_chart, generate_dashboard, update_chart, create_virtual_dataset,
+- Write tools (generate_chart, generate_dashboard, update_chart, create_dataset, create_virtual_dataset,
save_sql_query, add_chart_to_existing_dashboard, update_chart_preview) require write
permissions. These tools are only listed for users who have the necessary access.
If a write tool does not appear in the tool list, the current user lacks write access.
@@ -631,9 +632,9 @@ def create_mcp_app(
# Create default MCP instance for backward compatibility
mcp = create_mcp_app()
-# Initialize MCP dependency injection BEFORE importing tools/prompts
-# This replaces the abstract @tool and @prompt decorators in superset_core.api.mcp
-# with concrete implementations that can register with the mcp instance
+# Initialize MCP dependency injection BEFORE importing tools/prompts.
+# Replaces the stub @tool/@prompt decorators in superset_core.mcp.decorators
+# with concrete implementations bound to this mcp instance.
from superset.core.mcp.core_mcp_injection import ( # noqa: E402
initialize_core_mcp_dependencies,
)
@@ -658,6 +659,7 @@ warnings.filterwarnings(
module=r"google\..*",
)
+
# Import all MCP tools to register them with the mcp instance
# NOTE: Always add new tool imports here when creating new MCP tools.
# Tools use the @tool decorator from `superset-core` and register automatically
@@ -698,6 +700,7 @@ from superset.mcp_service.database.tool import ( # noqa: F401, E402
list_databases,
)
from superset.mcp_service.dataset.tool import ( # noqa: F401, E402
+ create_dataset,
create_virtual_dataset,
get_dataset_info,
list_datasets,
@@ -835,8 +838,9 @@ def init_fastmcp_server(
Returns:
The global FastMCP instance configured with the provided settings
"""
- # Read branding from Flask config's APP_NAME
- from superset.mcp_service.flask_singleton import app as flask_app
+ # circular import: flask_singleton imports from superset.extensions which
+ # re-enters mcp_service during startup; must stay lazy inside the function.
+ from superset.mcp_service.flask_singleton import app as flask_app # noqa: PLC0415
# Derive branding from Superset's APP_NAME config (defaults to "Superset")
app_name = flask_app.config.get("APP_NAME", "Superset")
diff --git a/superset/mcp_service/dataset/schemas.py b/superset/mcp_service/dataset/schemas.py
index 17c25cd0a67..d9c66b764b2 100644
--- a/superset/mcp_service/dataset/schemas.py
+++ b/superset/mcp_service/dataset/schemas.py
@@ -21,7 +21,7 @@ Pydantic schemas for dataset-related responses
from __future__ import annotations
-from datetime import datetime
+from datetime import datetime, timezone
from typing import Annotated, Any, Dict, List, Literal
from pydantic import (
@@ -325,8 +325,6 @@ class DatasetError(BaseModel):
@classmethod
def create(cls, error: str, error_type: str) -> "DatasetError":
"""Create a standardized DatasetError with timestamp."""
- from datetime import datetime, timezone
-
return cls(
error=error,
error_type=error_type,
@@ -410,6 +408,76 @@ class GetDatasetInfoRequest(MetadataCacheControl):
return parsed
+class CreateDatasetRequest(BaseModel):
+ """Request schema for create_dataset to register a physical table as a dataset."""
+
+ model_config = ConfigDict(populate_by_name=True)
+
+ database_id: Annotated[
+ int,
+ Field(
+ description="ID of the database connection to register the table against"
+ ),
+ ]
+ schema_: Annotated[
+ str | None,
+ Field(
+ default=None,
+ alias="schema",
+ serialization_alias="schema",
+ max_length=250,
+ description="Schema (namespace) where the table lives, e.g. 'public'. "
+ "Omit or pass None for databases without schema namespaces (e.g. SQLite).",
+ ),
+ ]
+ catalog: Annotated[
+ str | None,
+ Field(
+ default=None,
+ max_length=250,
+ description="Catalog where the table lives. Omit for databases without "
+ "catalog support.",
+ ),
+ ]
+ table_name: Annotated[
+ str,
+ Field(
+ min_length=1,
+ max_length=250,
+ description="Name of the physical table to register as a dataset",
+ ),
+ ]
+ owners: Annotated[
+ List[int] | None,
+ Field(
+ default=None,
+ description="Optional list of owner user IDs. "
+ "Defaults to the calling user.",
+ ),
+ ]
+
+ @field_validator("schema_", "catalog", mode="before")
+ @classmethod
+ def _normalize_optional_str(cls, v: object) -> object:
+ """Strip whitespace and convert blank strings to None.
+
+ Non-string values pass through unchanged so Pydantic's type validation
+ rejects them, rather than silently treating a malformed value (e.g. an
+ int or dict) as an omitted namespace.
+ """
+ if isinstance(v, str):
+ return v.strip() or None
+ return v
+
+ @field_validator("table_name", mode="before")
+ @classmethod
+ def _strip_table_name(cls, v: object) -> object:
+ """Strip leading/trailing whitespace from table_name."""
+ if isinstance(v, str):
+ return v.strip()
+ return v
+
+
class CreateVirtualDatasetRequest(BaseModel):
"""Request schema for create_virtual_dataset."""
@@ -734,7 +802,7 @@ def serialize_dataset_object(dataset: Any) -> DatasetInfo | None:
if isinstance(params, str):
try:
params = json.loads(params)
- except Exception:
+ except (json.JSONDecodeError, TypeError):
params = None
columns = [
TableColumnInfo(
diff --git a/superset/mcp_service/dataset/tool/__init__.py b/superset/mcp_service/dataset/tool/__init__.py
index cad8d4ed569..6fd3c12133c 100644
--- a/superset/mcp_service/dataset/tool/__init__.py
+++ b/superset/mcp_service/dataset/tool/__init__.py
@@ -15,14 +15,16 @@
# specific language governing permissions and limitations
# under the License.
+from .create_dataset import create_dataset
from .create_virtual_dataset import create_virtual_dataset
from .get_dataset_info import get_dataset_info
from .list_datasets import list_datasets
from .query_dataset import query_dataset
__all__ = [
+ "create_dataset",
"create_virtual_dataset",
- "list_datasets",
"get_dataset_info",
+ "list_datasets",
"query_dataset",
]
diff --git a/superset/mcp_service/dataset/tool/create_dataset.py b/superset/mcp_service/dataset/tool/create_dataset.py
new file mode 100644
index 00000000000..5767a450551
--- /dev/null
+++ b/superset/mcp_service/dataset/tool/create_dataset.py
@@ -0,0 +1,171 @@
+# 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.
+
+"""
+Create dataset FastMCP tool
+
+Registers a physical table as a Superset dataset against an existing
+database connection — the programmatic equivalent of Data → Datasets → +Dataset.
+Returns the same DatasetInfo shape as get_dataset_info so the caller can feed
+the resulting dataset_id directly into generate_chart.
+"""
+
+import logging
+
+from fastmcp import Context
+from superset_core.mcp.decorators import tool, ToolAnnotations
+
+from superset.commands.dataset.create import CreateDatasetCommand
+from superset.commands.dataset.exceptions import (
+ DatasetCreateFailedError,
+ DatasetInvalidError,
+)
+from superset.extensions import event_logger
+from superset.mcp_service.dataset.schemas import (
+ CreateDatasetRequest,
+ DatasetError,
+ DatasetInfo,
+ serialize_dataset_object,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _classify_invalid_error(exc: DatasetInvalidError) -> DatasetError:
+ """Map DatasetInvalidError sub-exceptions to typed DatasetError responses."""
+ classnames = exc.get_list_classnames()
+ messages = exc.normalized_messages()
+ if "DatabaseNotFoundValidationError" in classnames:
+ return DatasetError.create(
+ error="Database not found", error_type="DatabaseNotFoundError"
+ )
+ if "DatasetDataAccessIsNotAllowed" in classnames:
+ return DatasetError.create(
+ error="Access denied", error_type="AccessDeniedError"
+ )
+ if "DatasetExistsValidationError" in classnames:
+ return DatasetError.create(error=str(messages), error_type="DatasetExistsError")
+ if "TableNotFoundValidationError" in classnames:
+ return DatasetError.create(error=str(messages), error_type="TableNotFoundError")
+ # Other DatasetInvalidError sub-types are returned as generic ValidationError.
+ # Add explicit branches here when callers need to distinguish them.
+ return DatasetError.create(error=str(messages), error_type="ValidationError")
+
+
+@tool(
+ tags=["mutate"],
+ class_permission_name="Dataset",
+ method_permission_name="write",
+ annotations=ToolAnnotations(
+ title="Register physical table as dataset",
+ readOnlyHint=False,
+ destructiveHint=False,
+ ),
+)
+async def create_dataset(
+ request: CreateDatasetRequest, ctx: Context
+) -> DatasetInfo | DatasetError:
+ """Register a physical table as a Superset dataset.
+
+ Wraps POST /api/v1/dataset/ — the same endpoint the UI uses when you click
+ Data → Datasets → +Dataset. Returns full dataset metadata (same shape as
+ get_dataset_info) so you can pass the resulting dataset_id straight into
+ generate_chart.
+
+ Required fields:
+ - database_id: ID of the existing database connection
+ - table_name: Exact name of the physical table to register
+
+ Optional fields:
+ - schema: Schema/namespace where the table lives (e.g. "public"). Omit for
+ databases without schema namespaces (e.g. SQLite).
+ - catalog: Catalog where the table lives. Omit for databases without catalog
+ support.
+ - owners: List of user IDs to set as owners (defaults to calling user)
+
+ Example:
+ ```json
+ {
+ "database_id": 1,
+ "schema": "public",
+ "table_name": "orders"
+ }
+ ```
+
+ Returns DatasetInfo on success or DatasetError on failure.
+ Use list_databases to find the correct database_id.
+ """
+ schema = request.schema_
+ table_name = request.table_name
+ catalog = request.catalog
+
+ await ctx.info(
+ "Registering physical table as dataset: database_id=%s, table=%s"
+ % (request.database_id, f"{schema}.{table_name}" if schema else table_name)
+ )
+
+ try:
+ dataset_properties: dict[str, object] = {
+ "database": request.database_id,
+ "table_name": table_name,
+ }
+ if schema is not None:
+ dataset_properties["schema"] = schema
+ if catalog is not None:
+ dataset_properties["catalog"] = catalog
+ if request.owners is not None:
+ dataset_properties["owners"] = request.owners
+
+ with event_logger.log_context(action="mcp.create_dataset.create"):
+ dataset = CreateDatasetCommand(dataset_properties).run()
+
+ result = serialize_dataset_object(dataset)
+ if result is None:
+ return DatasetError.create(
+ error="Dataset was created but could not be serialized",
+ error_type="SerializationError",
+ )
+
+ await ctx.info(
+ "Dataset registered: id=%s, table=%s"
+ % (
+ dataset.id,
+ f"{schema}.{table_name}" if schema else table_name,
+ )
+ )
+ return result
+
+ except DatasetInvalidError as exc:
+ # CreateDatasetCommand.validate() collects validation errors into
+ # DatasetInvalidError.exceptions, never raising them directly.
+ # Inspect the wrapped class names for a typed response.
+ error_response = _classify_invalid_error(exc)
+ await ctx.warning(
+ "Dataset validation failed (%s): %s"
+ % (error_response.error_type, error_response.error)
+ )
+ return error_response
+ except DatasetCreateFailedError:
+ logger.exception("Dataset creation failed")
+ await ctx.error("Dataset creation failed")
+ return DatasetError.create(
+ error="Dataset creation failed", error_type="CreateFailedError"
+ )
+ except Exception as exc:
+ logger.exception("Unexpected error in create_dataset")
+ await ctx.error("Unexpected error: %s" % (type(exc).__name__,))
+ raise
diff --git a/tests/unit_tests/mcp_service/dataset/tool/test_create_dataset.py b/tests/unit_tests/mcp_service/dataset/tool/test_create_dataset.py
new file mode 100644
index 00000000000..ec949e5f3d2
--- /dev/null
+++ b/tests/unit_tests/mcp_service/dataset/tool/test_create_dataset.py
@@ -0,0 +1,585 @@
+# 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.
+
+"""Unit tests for create_dataset MCP tool."""
+
+import importlib
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+from fastmcp import Client
+from fastmcp.exceptions import ToolError
+
+from superset.commands.dataset.exceptions import (
+ DatabaseNotFoundValidationError,
+ DatasetCreateFailedError,
+ DatasetDataAccessIsNotAllowed,
+ DatasetExistsValidationError,
+ DatasetInvalidError,
+ TableNotFoundValidationError,
+)
+from superset.mcp_service.app import mcp
+from superset.sql.parse import Table
+from superset.utils import json
+
+# Use importlib to get the module object directly, bypassing the __init__.py
+# attribute binding that shadows the module name with the exported function.
+# This ensures patch.object resolves to the module in all Python versions.
+create_dataset_module = importlib.import_module(
+ "superset.mcp_service.dataset.tool.create_dataset"
+)
+
+
+def _make_mock_dataset(
+ dataset_id: int = 42,
+ table_name: str = "orders",
+ schema: str = "public",
+ database_name: str = "main_db",
+) -> MagicMock:
+ dataset = MagicMock()
+ dataset.id = dataset_id
+ dataset.table_name = table_name
+ dataset.schema = schema
+ dataset.description = None
+ dataset.certified_by = None
+ dataset.certification_details = None
+ dataset.changed_by_name = "admin"
+ dataset.changed_on = None
+ dataset.changed_on_humanized = None
+ dataset.created_by_name = "admin"
+ dataset.created_on = None
+ dataset.created_on_humanized = None
+ dataset.tags = []
+ dataset.owners = []
+ dataset.is_virtual = False
+ dataset.is_favorite = None
+ dataset.database_id = 1
+ dataset.schema_perm = f"[{database_name}].[{schema}]"
+ dataset.url = f"/tablemodelview/edit/{dataset_id}"
+ dataset.database = MagicMock()
+ dataset.database.database_name = database_name
+ dataset.sql = None
+ dataset.main_dttm_col = None
+ dataset.offset = 0
+ dataset.cache_timeout = 0
+ dataset.params = {}
+ dataset.template_params = {}
+ dataset.extra = {}
+ dataset.uuid = f"dataset-uuid-{dataset_id}"
+ dataset.columns = []
+ dataset.metrics = []
+ return dataset
+
+
+@pytest.fixture
+def mcp_server():
+ return mcp
+
+
+@pytest.fixture(autouse=True)
+def mock_auth():
+ 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
+
+
+class TestCreateDataset:
+ """Tests for the create_dataset MCP tool."""
+
+ @patch.object(create_dataset_module, "CreateDatasetCommand")
+ @pytest.mark.asyncio
+ async def test_create_dataset_success(self, mock_command_class, mcp_server) -> None:
+ """Happy path: tool creates dataset and returns DatasetInfo."""
+ mock_dataset = _make_mock_dataset()
+ mock_command = MagicMock()
+ mock_command.run.return_value = mock_dataset
+ mock_command_class.return_value = mock_command
+
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "create_dataset",
+ {
+ "request": {
+ "database_id": 1,
+ "schema": "public",
+ "table_name": "orders",
+ }
+ },
+ )
+
+ assert result.content is not None
+ data = json.loads(result.content[0].text)
+ assert data["id"] == 42
+ assert data["table_name"] == "orders"
+ assert data["schema"] == "public"
+
+ # Verify the command was called with the right properties
+ call_kwargs = mock_command_class.call_args[0][0]
+ assert call_kwargs["database"] == 1
+ assert call_kwargs["schema"] == "public"
+ assert call_kwargs["table_name"] == "orders"
+ assert "owners" not in call_kwargs
+
+ @patch.object(create_dataset_module, "CreateDatasetCommand")
+ @pytest.mark.asyncio
+ async def test_create_dataset_with_owners(
+ self, mock_command_class, mcp_server
+ ) -> None:
+ """Owners list is forwarded to the command when supplied."""
+ mock_dataset = _make_mock_dataset()
+ mock_command = MagicMock()
+ mock_command.run.return_value = mock_dataset
+ mock_command_class.return_value = mock_command
+
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "create_dataset",
+ {
+ "request": {
+ "database_id": 2,
+ "schema": "sales",
+ "table_name": "transactions",
+ "owners": [5, 10],
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["id"] == 42
+
+ call_kwargs = mock_command_class.call_args[0][0]
+ assert call_kwargs["owners"] == [5, 10]
+
+ @patch.object(create_dataset_module, "CreateDatasetCommand")
+ @pytest.mark.asyncio
+ async def test_create_dataset_already_exists(
+ self, mock_command_class, mcp_server
+ ) -> None:
+ """Returns DatasetExistsError when the table is already registered.
+
+ CreateDatasetCommand.validate() wraps DatasetExistsValidationError inside
+ DatasetInvalidError. The tool must inspect get_list_classnames() to surface
+ the typed error response.
+ """
+ mock_command = MagicMock()
+ mock_command.run.side_effect = DatasetInvalidError(
+ exceptions=[DatasetExistsValidationError(Table("orders", "public", None))]
+ )
+ mock_command_class.return_value = mock_command
+
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "create_dataset",
+ {
+ "request": {
+ "database_id": 1,
+ "schema": "public",
+ "table_name": "orders",
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["error_type"] == "DatasetExistsError"
+ assert "error" in data
+
+ @patch.object(create_dataset_module, "CreateDatasetCommand")
+ @pytest.mark.asyncio
+ async def test_create_dataset_table_not_found(
+ self, mock_command_class, mcp_server
+ ) -> None:
+ """Returns TableNotFoundError when the physical table does not exist in the DB.
+
+ CreateDatasetCommand.validate() wraps TableNotFoundValidationError inside
+ DatasetInvalidError. The tool must inspect get_list_classnames() to surface
+ the typed error response.
+ """
+ mock_command = MagicMock()
+ mock_command.run.side_effect = DatasetInvalidError(
+ exceptions=[
+ TableNotFoundValidationError(Table("missing_table", "public", None))
+ ]
+ )
+ mock_command_class.return_value = mock_command
+
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "create_dataset",
+ {
+ "request": {
+ "database_id": 1,
+ "schema": "public",
+ "table_name": "missing_table",
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["error_type"] == "TableNotFoundError"
+
+ @patch.object(create_dataset_module, "CreateDatasetCommand")
+ @pytest.mark.asyncio
+ async def test_create_dataset_with_catalog(
+ self, mock_command_class, mcp_server
+ ) -> None:
+ """Catalog field is normalized and forwarded to the command when supplied."""
+ mock_dataset = _make_mock_dataset()
+ mock_command = MagicMock()
+ mock_command.run.return_value = mock_dataset
+ mock_command_class.return_value = mock_command
+
+ async with Client(mcp_server) as client:
+ await client.call_tool(
+ "create_dataset",
+ {
+ "request": {
+ "database_id": 1,
+ "catalog": " hive ",
+ "schema": "default",
+ "table_name": "events",
+ }
+ },
+ )
+
+ call_kwargs = mock_command_class.call_args[0][0]
+ assert call_kwargs["catalog"] == "hive"
+
+ @patch.object(create_dataset_module, "CreateDatasetCommand")
+ @pytest.mark.asyncio
+ async def test_create_dataset_invalid_error(
+ self, mock_command_class, mcp_server
+ ) -> None:
+ """DatasetInvalidError is returned as ValidationError type."""
+ mock_command = MagicMock()
+ mock_command.run.side_effect = DatasetInvalidError()
+ mock_command_class.return_value = mock_command
+
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "create_dataset",
+ {
+ "request": {
+ "database_id": 1,
+ "schema": "public",
+ "table_name": "orders",
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["error_type"] == "ValidationError"
+ assert "error" in data
+
+ @patch.object(create_dataset_module, "CreateDatasetCommand")
+ @pytest.mark.asyncio
+ async def test_create_dataset_create_failed_error(
+ self, mock_command_class, mcp_server
+ ) -> None:
+ """DatasetCreateFailedError is returned as CreateFailedError type."""
+ mock_command = MagicMock()
+ mock_command.run.side_effect = DatasetCreateFailedError()
+ mock_command_class.return_value = mock_command
+
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "create_dataset",
+ {
+ "request": {
+ "database_id": 1,
+ "schema": "public",
+ "table_name": "orders",
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["error_type"] == "CreateFailedError"
+ assert "Dataset creation failed" in data["error"]
+
+ @patch.object(create_dataset_module, "CreateDatasetCommand")
+ @pytest.mark.asyncio
+ async def test_create_dataset_database_not_found(
+ self, mock_command_class, mcp_server
+ ) -> None:
+ """Returns DatabaseNotFoundError when CreateDatasetCommand raises it."""
+ mock_command = MagicMock()
+ mock_command.run.side_effect = DatasetInvalidError(
+ exceptions=[DatabaseNotFoundValidationError()]
+ )
+ mock_command_class.return_value = mock_command
+
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "create_dataset",
+ {
+ "request": {
+ "database_id": 999,
+ "table_name": "orders",
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["error_type"] == "DatabaseNotFoundError"
+ assert "Database not found" in data["error"]
+
+ @patch.object(create_dataset_module, "CreateDatasetCommand")
+ @pytest.mark.asyncio
+ async def test_create_dataset_access_denied(
+ self, mock_command_class, mcp_server
+ ) -> None:
+ """Returns AccessDeniedError when CreateDatasetCommand raises it."""
+ mock_command = MagicMock()
+ mock_command.run.side_effect = DatasetInvalidError(
+ exceptions=[DatasetDataAccessIsNotAllowed("Access is Denied")]
+ )
+ mock_command_class.return_value = mock_command
+
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "create_dataset",
+ {
+ "request": {
+ "database_id": 1,
+ "schema": "secret",
+ "table_name": "restricted_table",
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["error_type"] == "AccessDeniedError"
+ assert "Access denied" in data["error"]
+
+ @patch.object(create_dataset_module, "CreateDatasetCommand")
+ @pytest.mark.asyncio
+ async def test_create_dataset_unexpected_error(
+ self, mock_command_class, mcp_server
+ ) -> None:
+ """Unexpected exceptions are re-raised as ToolError (handled by middleware)."""
+ mock_command = MagicMock()
+ mock_command.run.side_effect = RuntimeError("DB connection lost")
+ mock_command_class.return_value = mock_command
+
+ async with Client(mcp_server) as client:
+ with pytest.raises(ToolError):
+ await client.call_tool(
+ "create_dataset",
+ {
+ "request": {
+ "database_id": 1,
+ "schema": "public",
+ "table_name": "orders",
+ }
+ },
+ )
+
+ @pytest.mark.asyncio
+ async def test_create_dataset_missing_required_fields(self, mcp_server) -> None:
+ """Missing required fields raise a validation error before the tool runs."""
+ async with Client(mcp_server) as client:
+ with pytest.raises(ToolError):
+ await client.call_tool(
+ "create_dataset",
+ {
+ "request": {
+ # database_id and table_name are omitted intentionally
+ "schema": "public",
+ }
+ },
+ )
+
+ @patch.object(create_dataset_module, "CreateDatasetCommand")
+ @pytest.mark.asyncio
+ async def test_create_dataset_returns_full_dataset_info(
+ self, mock_command_class, mcp_server
+ ) -> None:
+ """The returned DatasetInfo includes columns, metrics, and all core fields."""
+ mock_dataset = _make_mock_dataset(
+ dataset_id=99, table_name="sales", schema="dw"
+ )
+
+ col = MagicMock()
+ col.column_name = "amount"
+ col.verbose_name = "Amount"
+ col.type = "NUMERIC"
+ col.is_dttm = False
+ col.groupby = True
+ col.filterable = True
+ col.description = "Sale amount"
+ mock_dataset.columns = [col]
+
+ metric = MagicMock()
+ metric.metric_name = "total_sales"
+ metric.verbose_name = "Total Sales"
+ metric.expression = "SUM(amount)"
+ metric.description = "Sum of amounts"
+ metric.d3format = None
+ mock_dataset.metrics = [metric]
+
+ mock_command = MagicMock()
+ mock_command.run.return_value = mock_dataset
+ mock_command_class.return_value = mock_command
+
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "create_dataset",
+ {
+ "request": {
+ "database_id": 1,
+ "schema": "dw",
+ "table_name": "sales",
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["id"] == 99
+ assert data["table_name"] == "sales"
+ assert data["schema"] == "dw"
+ assert data["is_virtual"] is False
+ assert len(data["columns"]) == 1
+ assert data["columns"][0]["column_name"] == "amount"
+ assert len(data["metrics"]) == 1
+ assert data["metrics"][0]["metric_name"] == "total_sales"
+
+ @patch.object(create_dataset_module, "CreateDatasetCommand")
+ @pytest.mark.asyncio
+ async def test_create_dataset_table_name_whitespace_normalized(
+ self, mock_command_class, mcp_server
+ ) -> None:
+ """Whitespace in table_name is stripped before forwarding to the command."""
+ mock_dataset = _make_mock_dataset()
+ mock_command = MagicMock()
+ mock_command.run.return_value = mock_dataset
+ mock_command_class.return_value = mock_command
+
+ async with Client(mcp_server) as client:
+ await client.call_tool(
+ "create_dataset",
+ {
+ "request": {
+ "database_id": 1,
+ "schema": "public",
+ "table_name": " orders ",
+ }
+ },
+ )
+
+ call_kwargs = mock_command_class.call_args[0][0]
+ assert call_kwargs["table_name"] == "orders"
+
+ @patch.object(create_dataset_module, "CreateDatasetCommand")
+ @pytest.mark.asyncio
+ async def test_create_dataset_blank_schema_normalized_to_none(
+ self, mock_command_class, mcp_server
+ ) -> None:
+ """Blank schema string is treated as absent: not forwarded to the command."""
+ mock_dataset = _make_mock_dataset(schema="")
+ mock_command = MagicMock()
+ mock_command.run.return_value = mock_dataset
+ mock_command_class.return_value = mock_command
+
+ async with Client(mcp_server) as client:
+ await client.call_tool(
+ "create_dataset",
+ {
+ "request": {
+ "database_id": 1,
+ "schema": "",
+ "table_name": "orders",
+ }
+ },
+ )
+
+ call_kwargs = mock_command_class.call_args[0][0]
+ assert "schema" not in call_kwargs
+
+ @patch.object(create_dataset_module, "CreateDatasetCommand")
+ @pytest.mark.asyncio
+ async def test_create_dataset_blank_catalog_normalized_to_none(
+ self, mock_command_class, mcp_server
+ ) -> None:
+ """Blank catalog string is treated as absent: not forwarded to the command."""
+ mock_dataset = _make_mock_dataset()
+ mock_command = MagicMock()
+ mock_command.run.return_value = mock_dataset
+ mock_command_class.return_value = mock_command
+
+ async with Client(mcp_server) as client:
+ await client.call_tool(
+ "create_dataset",
+ {
+ "request": {
+ "database_id": 1,
+ "catalog": "",
+ "table_name": "orders",
+ }
+ },
+ )
+
+ call_kwargs = mock_command_class.call_args[0][0]
+ assert "catalog" not in call_kwargs
+
+ @pytest.mark.asyncio
+ async def test_create_dataset_non_string_namespace_rejected(
+ self, mcp_server
+ ) -> None:
+ """Non-string schema/catalog values fail validation, not silently dropped."""
+ async with Client(mcp_server) as client:
+ for field, value in (("schema", 123), ("catalog", {"name": "hive"})):
+ with pytest.raises(ToolError):
+ await client.call_tool(
+ "create_dataset",
+ {
+ "request": {
+ "database_id": 1,
+ "table_name": "orders",
+ field: value,
+ }
+ },
+ )
+
+ @patch.object(create_dataset_module, "CreateDatasetCommand")
+ @patch.object(create_dataset_module, "serialize_dataset_object", return_value=None)
+ @pytest.mark.asyncio
+ async def test_create_dataset_when_serialize_returns_none(
+ self, mock_serialize, mock_command_class, mcp_server
+ ) -> None:
+ """Returns SerializationError when serialize_dataset_object returns None."""
+ mock_dataset = _make_mock_dataset()
+ mock_command = MagicMock()
+ mock_command.run.return_value = mock_dataset
+ mock_command_class.return_value = mock_command
+
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "create_dataset",
+ {
+ "request": {
+ "database_id": 1,
+ "schema": "public",
+ "table_name": "orders",
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["error_type"] == "SerializationError"