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 9b6ff262fddb926c376d6de140070cd056e8d5b1 Author: Amin Ghadersohi <[email protected]> AuthorDate: Thu May 28 04:31:36 2026 +0000 fix(mcp): restore missing tool registrations and fix create_dataset tests - Restore create_virtual_dataset, query_dataset, get_database_info, list_databases, get_chart_sql, get_chart_type_schema, save_sql_query imports in app.py (accidentally dropped in the feat commit) - Restore create_virtual_dataset and query_dataset exports in dataset tool __init__.py - Make CreateDatasetRequest.schema optional (str | None, default None) - Refactor create_dataset.py to use @tool decorator pattern - Fix test_create_dataset.py: convert @patch class decorators to with patch() context managers (avoids pytest-asyncio arg injection issues), add get_superset_base_url mock for success paths, and set certified_by/certification_details=None on mock dataset --- superset/mcp_service/app.py | 9 + superset/mcp_service/dataset/schemas.py | 8 +- superset/mcp_service/dataset/tool/__init__.py | 8 +- .../mcp_service/dataset/tool/create_dataset.py | 68 +++---- .../dataset/tool/test_create_dataset.py | 206 ++++++++++++--------- 5 files changed, 172 insertions(+), 127 deletions(-) diff --git a/superset/mcp_service/app.py b/superset/mcp_service/app.py index 0198b6252f3..71d831a2ef9 100644 --- a/superset/mcp_service/app.py +++ b/superset/mcp_service/app.py @@ -555,6 +555,8 @@ from superset.mcp_service.chart.tool import ( # noqa: F401, E402 get_chart_data, get_chart_info, get_chart_preview, + get_chart_sql, + get_chart_type_schema, list_charts, update_chart, update_chart_preview, @@ -565,10 +567,16 @@ from superset.mcp_service.dashboard.tool import ( # noqa: F401, E402 get_dashboard_info, list_dashboards, ) +from superset.mcp_service.database.tool import ( # noqa: F401, E402 + get_database_info, + list_databases, +) from superset.mcp_service.dataset.tool import ( # noqa: F401, E402 create_dataset, + create_virtual_dataset, get_dataset_info, list_datasets, + query_dataset, ) from superset.mcp_service.explore.tool import ( # noqa: F401, E402 generate_explore_link, @@ -576,6 +584,7 @@ from superset.mcp_service.explore.tool import ( # noqa: F401, E402 from superset.mcp_service.sql_lab.tool import ( # noqa: F401, E402 execute_sql, open_sql_lab_with_context, + save_sql_query, ) from superset.mcp_service.system import ( # noqa: F401, E402 prompts as system_prompts, diff --git a/superset/mcp_service/dataset/schemas.py b/superset/mcp_service/dataset/schemas.py index 0bbc4061f8c..cf51d0afd00 100644 --- a/superset/mcp_service/dataset/schemas.py +++ b/superset/mcp_service/dataset/schemas.py @@ -334,8 +334,12 @@ class CreateDatasetRequest(BaseModel): ), ] schema: Annotated[ - str, - Field(description="Schema (namespace) where the table lives, e.g. 'public'"), + str | None, + Field( + default=None, + description="Schema (namespace) where the table lives, e.g. 'public'. " + "Optional: omit to use the database default schema.", + ), ] table_name: Annotated[ str, diff --git a/superset/mcp_service/dataset/tool/__init__.py b/superset/mcp_service/dataset/tool/__init__.py index 025b4ae1b9a..6fd3c12133c 100644 --- a/superset/mcp_service/dataset/tool/__init__.py +++ b/superset/mcp_service/dataset/tool/__init__.py @@ -16,11 +16,15 @@ # 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__ = [ - "list_datasets", - "get_dataset_info", "create_dataset", + "create_virtual_dataset", + "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 index 6484f8ab36d..7c2cc3c0d8c 100644 --- a/superset/mcp_service/dataset/tool/create_dataset.py +++ b/superset/mcp_service/dataset/tool/create_dataset.py @@ -25,12 +25,12 @@ the resulting dataset_id directly into generate_chart. """ import logging -from datetime import datetime, timezone +from typing import Any from fastmcp import Context +from superset_core.mcp.decorators import tool, ToolAnnotations -from superset.mcp_service.app import mcp -from superset.mcp_service.auth import mcp_auth_hook +from superset.extensions import event_logger from superset.mcp_service.dataset.schemas import ( CreateDatasetRequest, DatasetError, @@ -41,9 +41,17 @@ from superset.mcp_service.dataset.schemas import ( logger = logging.getLogger(__name__) [email protected](tags=["mutate"]) -@mcp_auth_hook -def create_dataset( +@tool( + tags=["mutate"], + class_permission_name="Dataset", + method_permission_name="write", + annotations=ToolAnnotations( + title="Create dataset", + readOnlyHint=False, + destructiveHint=False, + ), +) +async def create_dataset( request: CreateDatasetRequest, ctx: Context ) -> DatasetInfo | DatasetError: """Register a physical table as a Superset dataset. @@ -55,10 +63,10 @@ def create_dataset( Required fields: - database_id: ID of the existing database connection - - schema: Schema/namespace where the table lives (e.g. "public") - table_name: Exact name of the physical table to register Optional fields: + - schema: Schema/namespace where the table lives (e.g. "public") - owners: List of user IDs to set as owners (defaults to calling user) Example: @@ -73,6 +81,10 @@ def create_dataset( Returns DatasetInfo on success or DatasetError on failure. Use list_databases to find the correct database_id. """ + await ctx.info( + "Creating dataset: database_id=%s, schema=%r, table_name=%r" + % (request.database_id, request.schema, request.table_name) + ) try: from superset.commands.dataset.create import CreateDatasetCommand from superset.commands.dataset.exceptions import ( @@ -82,7 +94,7 @@ def create_dataset( TableNotFoundValidationError, ) - dataset_properties = { + dataset_properties: dict[str, Any] = { "database": request.database_id, "schema": request.schema, "table_name": request.table_name, @@ -90,15 +102,15 @@ def create_dataset( if request.owners is not None: dataset_properties["owners"] = request.owners - command = CreateDatasetCommand(dataset_properties) - dataset = command.run() + with event_logger.log_context(action="mcp.create_dataset"): + command = CreateDatasetCommand(dataset_properties) + dataset = command.run() result = serialize_dataset_object(dataset) if result is None: - return DatasetError( + return DatasetError.create( error="Dataset was created but could not be serialized", error_type="SerializationError", - timestamp=datetime.now(timezone.utc), ) logger.info( @@ -110,33 +122,21 @@ def create_dataset( return result except DatasetExistsValidationError as e: - return DatasetError( - error=str(e), - error_type="DatasetExistsError", - timestamp=datetime.now(timezone.utc), - ) + await ctx.error("Dataset already exists: %s" % (str(e),)) + return DatasetError.create(error=str(e), error_type="DatasetExistsError") except TableNotFoundValidationError as e: - return DatasetError( - error=str(e), - error_type="TableNotFoundError", - timestamp=datetime.now(timezone.utc), - ) + await ctx.error("Table not found: %s" % (str(e),)) + return DatasetError.create(error=str(e), error_type="TableNotFoundError") except DatasetInvalidError as e: - return DatasetError( - error=str(e), - error_type="ValidationError", - timestamp=datetime.now(timezone.utc), - ) + await ctx.error("Dataset validation failed: %s" % (str(e),)) + return DatasetError.create(error=str(e), error_type="ValidationError") except DatasetCreateFailedError as e: - return DatasetError( - error=str(e), - error_type="CreateFailedError", - timestamp=datetime.now(timezone.utc), - ) + await ctx.error("Dataset creation failed: %s" % (str(e),)) + return DatasetError.create(error=str(e), error_type="CreateFailedError") except Exception as e: logger.error("Failed to create dataset: %s", e, exc_info=True) - return DatasetError( + await ctx.error("Unexpected error: %s: %s" % (type(e).__name__, str(e))) + return DatasetError.create( error=f"Failed to create dataset: {str(e)}", error_type="InternalError", - timestamp=datetime.now(timezone.utc), ) 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 index ba2898eb660..01b0a0813c8 100644 --- a/tests/unit_tests/mcp_service/dataset/tool/test_create_dataset.py +++ b/tests/unit_tests/mcp_service/dataset/tool/test_create_dataset.py @@ -52,6 +52,8 @@ def _make_mock_dataset( dataset.owners = [] dataset.is_virtual = False dataset.database_id = 1 + dataset.certified_by = None + dataset.certification_details = None dataset.schema_perm = f"[{database_name}].[{schema}]" dataset.url = f"/tablemodelview/edit/{dataset_id}" dataset.database = MagicMock() @@ -87,26 +89,34 @@ def mock_auth(): class TestCreateDataset: """Tests for the create_dataset MCP tool.""" - @patch("superset.commands.dataset.create.CreateDatasetCommand") @pytest.mark.asyncio - async def test_create_dataset_success(self, mock_command_class, mcp_server): + async def test_create_dataset_success(self, mcp_server): """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", - } - }, - ) + with ( + patch( + "superset.commands.dataset.create.CreateDatasetCommand", + return_value=mock_command, + ) as mock_command_class, + patch( + "superset.mcp_service.utils.url_utils.get_superset_base_url", + return_value="http://localhost:8088";, + ), + ): + 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) @@ -121,27 +131,35 @@ class TestCreateDataset: assert call_kwargs["table_name"] == "orders" assert "owners" not in call_kwargs - @patch("superset.commands.dataset.create.CreateDatasetCommand") @pytest.mark.asyncio - async def test_create_dataset_with_owners(self, mock_command_class, mcp_server): + async def test_create_dataset_with_owners(self, mcp_server): """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], - } - }, - ) + with ( + patch( + "superset.commands.dataset.create.CreateDatasetCommand", + return_value=mock_command, + ) as mock_command_class, + patch( + "superset.mcp_service.utils.url_utils.get_superset_base_url", + return_value="http://localhost:8088";, + ), + ): + 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 @@ -149,9 +167,8 @@ class TestCreateDataset: call_kwargs = mock_command_class.call_args[0][0] assert call_kwargs["owners"] == [5, 10] - @patch("superset.commands.dataset.create.CreateDatasetCommand") @pytest.mark.asyncio - async def test_create_dataset_already_exists(self, mock_command_class, mcp_server): + async def test_create_dataset_already_exists(self, mcp_server): """Returns DatasetError when a dataset for the table already exists.""" from superset.commands.dataset.exceptions import DatasetExistsValidationError from superset.sql.parse import Table @@ -160,27 +177,29 @@ class TestCreateDataset: mock_command.run.side_effect = 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", - } - }, - ) + with patch( + "superset.commands.dataset.create.CreateDatasetCommand", + 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("superset.commands.dataset.create.CreateDatasetCommand") @pytest.mark.asyncio - async def test_create_dataset_table_not_found(self, mock_command_class, mcp_server): + async def test_create_dataset_table_not_found(self, mcp_server): """Returns DatasetError when the physical table does not exist in the DB.""" from superset.commands.dataset.exceptions import TableNotFoundValidationError from superset.sql.parse import Table @@ -189,44 +208,47 @@ class TestCreateDataset: mock_command.run.side_effect = 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", - } - }, - ) + with patch( + "superset.commands.dataset.create.CreateDatasetCommand", + 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("superset.commands.dataset.create.CreateDatasetCommand") @pytest.mark.asyncio - async def test_create_dataset_unexpected_error( - self, mock_command_class, mcp_server - ): + async def test_create_dataset_unexpected_error(self, mcp_server): """Unexpected exceptions are caught and returned as InternalError.""" 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: - result = await client.call_tool( - "create_dataset", - { - "request": { - "database_id": 1, - "schema": "public", - "table_name": "orders", - } - }, - ) + with patch( + "superset.commands.dataset.create.CreateDatasetCommand", + 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"] == "InternalError" @@ -247,11 +269,8 @@ class TestCreateDataset: }, ) - @patch("superset.commands.dataset.create.CreateDatasetCommand") @pytest.mark.asyncio - async def test_create_dataset_returns_full_dataset_info( - self, mock_command_class, mcp_server - ): + async def test_create_dataset_returns_full_dataset_info(self, mcp_server): """The returned DatasetInfo includes columns, metrics, and all core fields.""" mock_dataset = _make_mock_dataset( dataset_id=99, table_name="sales", schema="dw" @@ -277,19 +296,28 @@ class TestCreateDataset: 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", - } - }, - ) + 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: + 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
