This is an automated email from the ASF dual-hosted git repository.
msyavuz 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 2fac66d1a30 feat(mcp): add update_dataset_metric tool for editing
saved dataset metrics (#40975)
2fac66d1a30 is described below
commit 2fac66d1a308036c60203593afb913d5a0593420
Author: Mehmet Salih Yavuz <[email protected]>
AuthorDate: Mon Jul 13 15:44:57 2026 +0300
feat(mcp): add update_dataset_metric tool for editing saved dataset metrics
(#40975)
---
superset/mcp_service/app.py | 6 +-
superset/mcp_service/dataset/dataset_utils.py | 50 ++
superset/mcp_service/dataset/schemas.py | 151 ++++++
superset/mcp_service/dataset/tool/__init__.py | 2 +
superset/mcp_service/dataset/tool/query_dataset.py | 30 +-
.../dataset/tool/update_dataset_metric.py | 294 +++++++++++
.../mcp_service/dataset/tool/test_query_dataset.py | 32 +-
.../dataset/tool/test_update_dataset_metric.py | 541 +++++++++++++++++++++
8 files changed, 1060 insertions(+), 46 deletions(-)
diff --git a/superset/mcp_service/app.py b/superset/mcp_service/app.py
index 9fb597f66a5..017909d3d12 100644
--- a/superset/mcp_service/app.py
+++ b/superset/mcp_service/app.py
@@ -175,6 +175,7 @@ Dataset Management:
- 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)
+- update_dataset_metric: Update a saved metric on a dataset — expression,
name, verbose_name, format (requires dataset ownership)
- query_dataset: Query a dataset using its semantic layer (saved metrics,
dimensions, filters) without needing a saved chart
Chart Management:
@@ -463,8 +464,8 @@ Input format:
{_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,
duplicate_dashboard,
- create_dataset, create_virtual_dataset, save_sql_query,
add_chart_to_existing_dashboard,
- manage_native_filters, remove_chart_from_dashboard,
+ create_dataset, create_virtual_dataset, update_dataset_metric,
save_sql_query,
+ add_chart_to_existing_dashboard, manage_native_filters,
remove_chart_from_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.
@@ -754,6 +755,7 @@ from superset.mcp_service.dataset.tool import ( # noqa:
F401, E402
get_dataset_info,
list_datasets,
query_dataset,
+ update_dataset_metric,
)
from superset.mcp_service.explore.tool import ( # noqa: F401, E402
generate_explore_link,
diff --git a/superset/mcp_service/dataset/dataset_utils.py
b/superset/mcp_service/dataset/dataset_utils.py
new file mode 100644
index 00000000000..9a1a14fbb1b
--- /dev/null
+++ b/superset/mcp_service/dataset/dataset_utils.py
@@ -0,0 +1,50 @@
+# 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.
+
+"""Shared helpers for dataset MCP tools."""
+
+from typing import Any
+
+from superset.mcp_service.utils import _is_uuid
+
+
+def resolve_dataset(
+ identifier: int | str, eager_options: list[Any] | None = None
+) -> Any | None:
+ """Resolve a dataset by int ID or UUID string.
+
+ Replicates the identifier resolution logic from
ModelGetInfoCore._find_object().
+ """
+ from superset.daos.dataset import DatasetDAO
+
+ opts = eager_options or None
+
+ if isinstance(identifier, int):
+ return DatasetDAO.find_by_id(identifier, query_options=opts)
+
+ # Try parsing as int
+ try:
+ id_val = int(identifier)
+ return DatasetDAO.find_by_id(id_val, query_options=opts)
+ except (ValueError, TypeError):
+ pass
+
+ # Try UUID
+ if _is_uuid(str(identifier)):
+ return DatasetDAO.find_by_id(identifier, id_column="uuid",
query_options=opts)
+
+ return None
diff --git a/superset/mcp_service/dataset/schemas.py
b/superset/mcp_service/dataset/schemas.py
index d469a753a72..adbf10cd210 100644
--- a/superset/mcp_service/dataset/schemas.py
+++ b/superset/mcp_service/dataset/schemas.py
@@ -607,6 +607,157 @@ class CreateVirtualDatasetResponse(BaseModel):
)
+UPDATABLE_METRIC_FIELDS: frozenset[str] = frozenset(
+ {
+ "metric_name",
+ "expression",
+ "verbose_name",
+ "description",
+ "d3format",
+ "metric_type",
+ "currency",
+ "warning_text",
+ "extra",
+ }
+)
+
+
+class MetricCurrency(BaseModel):
+ """Currency formatting configuration for a metric."""
+
+ symbol: str | None = Field(
+ None, description="ISO 4217 currency code (e.g. 'USD', 'EUR')."
+ )
+ symbolPosition: str | None = Field( # noqa: N815
+ None,
+ description="Where to render the symbol relative to the value "
+ "(typically 'prefix' or 'suffix').",
+ )
+
+
+class UpdateDatasetMetricRequest(BaseModel):
+ """Request schema for update_dataset_metric."""
+
+ model_config = ConfigDict(populate_by_name=True)
+
+ dataset_id: int | str = Field(
+ ...,
+ description="Dataset identifier — numeric ID or UUID string. "
+ "Use list_datasets to find valid IDs.",
+ )
+ metric: int | str = Field(
+ ...,
+ description="Metric to update — numeric metric ID, metric UUID, or "
+ "metric_name (e.g. 'sum_revenue'). Numeric strings are treated as IDs.
"
+ "Use get_dataset_info to discover a dataset's saved metrics.",
+ )
+ metric_name: str | None = Field(
+ None,
+ max_length=255,
+ description="New metric name, unique within the dataset. "
+ "Only pass this to rename the metric.",
+ )
+ expression: str | None = Field(
+ None,
+ description="New SQL aggregation expression (e.g. 'SUM(revenue)').",
+ )
+ verbose_name: str | None = Field(
+ None, max_length=1024, description="Human-friendly display label."
+ )
+ description: str | None = Field(None, description="Metric description.")
+ d3format: str | None = Field(
+ None,
+ min_length=1,
+ max_length=128,
+ description="D3 number format string (e.g. ',.2f', '.1%').",
+ )
+ metric_type: str | None = Field(
+ None,
+ min_length=1,
+ max_length=32,
+ description="Metric type (e.g. 'count', 'sum').",
+ )
+ currency: MetricCurrency | None = Field(
+ None, description="Currency formatting configuration."
+ )
+ warning_text: str | None = Field(
+ None, description="Warning shown to users of this metric."
+ )
+ extra: str | None = Field(
+ None, description="JSON-encoded string with extra metric metadata."
+ )
+
+ def updates(self) -> Dict[str, Any]:
+ """Return only the metric properties explicitly provided by the caller.
+
+ ``exclude_unset`` distinguishes "not provided" (leave the stored value
+ alone) from an explicit ``null`` (clear the stored value).
+ """
+ return self.model_dump(
+ exclude_unset=True,
+ include=set(UPDATABLE_METRIC_FIELDS),
+ )
+
+ @model_validator(mode="after")
+ def validate_updates(self) -> "UpdateDatasetMetricRequest":
+ """Require at least one updatable property and reject empty/invalid
values.
+
+ Guards against no-op requests, empty ``metric_name``/``expression``,
and
+ malformed ``extra`` JSON before the update reaches the command layer.
+ """
+ provided = self.model_fields_set & UPDATABLE_METRIC_FIELDS
+ if not provided:
+ raise ValueError(
+ "At least one metric property must be provided to update. "
+ f"Updatable properties: {sorted(UPDATABLE_METRIC_FIELDS)}."
+ )
+ if "metric_name" in provided and not (self.metric_name or "").strip():
+ raise ValueError("metric_name cannot be empty or null")
+ if "expression" in provided and not (self.expression or "").strip():
+ raise ValueError("expression cannot be empty or null")
+ if "extra" in provided and self.extra is not None:
+ try:
+ json.loads(self.extra)
+ except (ValueError, TypeError) as ex:
+ raise ValueError("extra must be a valid JSON-encoded string")
from ex
+ return self
+
+
+class DatasetMetricDetail(SqlMetricInfo):
+ """Full saved-metric details, including identifiers."""
+
+ id: int | None = Field(None, description="Metric ID")
+ uuid: str | None = Field(None, description="Metric UUID")
+ metric_type: str | None = Field(None, description="Metric type")
+ currency: MetricCurrency | None = Field(
+ None, description="Currency formatting configuration"
+ )
+ warning_text: str | None = Field(None, description="Warning text")
+ extra: str | None = Field(
+ None, description="JSON-encoded string with extra metric metadata"
+ )
+
+
+class UpdateDatasetMetricResponse(BaseModel):
+ """Response schema for update_dataset_metric."""
+
+ dataset_id: int | None = Field(None, description="Dataset ID")
+ dataset_name: str | None = Field(None, description="Dataset name")
+ metric: DatasetMetricDetail | None = Field(
+ None, description="The metric after the update. None if the update
failed."
+ )
+ updated_properties: List[str] = Field(
+ default_factory=list,
+ description="Names of the metric properties that were updated.",
+ )
+ url: str | None = Field(
+ None, description="Explore URL for the dataset. None if the update
failed."
+ )
+ error: str | None = Field(
+ None, description="Error message if the update failed, otherwise null."
+ )
+
+
VALID_FILTER_OPS = Literal[
"==",
"!=",
diff --git a/superset/mcp_service/dataset/tool/__init__.py
b/superset/mcp_service/dataset/tool/__init__.py
index 6fd3c12133c..191ffd4c4d8 100644
--- a/superset/mcp_service/dataset/tool/__init__.py
+++ b/superset/mcp_service/dataset/tool/__init__.py
@@ -20,6 +20,7 @@ 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
+from .update_dataset_metric import update_dataset_metric
__all__ = [
"create_dataset",
@@ -27,4 +28,5 @@ __all__ = [
"get_dataset_info",
"list_datasets",
"query_dataset",
+ "update_dataset_metric",
]
diff --git a/superset/mcp_service/dataset/tool/query_dataset.py
b/superset/mcp_service/dataset/tool/query_dataset.py
index 6850d930873..8b92966e161 100644
--- a/superset/mcp_service/dataset/tool/query_dataset.py
+++ b/superset/mcp_service/dataset/tool/query_dataset.py
@@ -35,6 +35,7 @@ from superset.commands.exceptions import CommandException
from superset.exceptions import OAuth2Error, OAuth2RedirectError,
SupersetException
from superset.extensions import event_logger
from superset.mcp_service.chart.schemas import DataColumn, PerformanceMetadata
+from superset.mcp_service.dataset.dataset_utils import resolve_dataset
from superset.mcp_service.dataset.schemas import (
DatasetError,
QueryDatasetFilter,
@@ -46,7 +47,6 @@ from superset.mcp_service.privacy import (
requires_data_model_metadata_access,
user_can_view_data_model_metadata,
)
-from superset.mcp_service.utils import _is_uuid
from superset.mcp_service.utils.cache_utils import get_cache_status_from_result
from superset.mcp_service.utils.oauth2_utils import
build_oauth2_redirect_message
from superset.mcp_service.utils.query_utils import validate_names
@@ -55,32 +55,6 @@ from superset.mcp_service.utils.response_utils import
format_data_columns
logger = logging.getLogger(__name__)
-def _resolve_dataset(identifier: int | str, eager_options: list[Any]) -> Any |
None:
- """Resolve a dataset by int ID or UUID string.
-
- Replicates the identifier resolution logic from
ModelGetInfoCore._find_object().
- """
- from superset.daos.dataset import DatasetDAO
-
- opts = eager_options or None
-
- if isinstance(identifier, int):
- return DatasetDAO.find_by_id(identifier, query_options=opts)
-
- # Try parsing as int
- try:
- id_val = int(identifier)
- return DatasetDAO.find_by_id(id_val, query_options=opts)
- except (ValueError, TypeError):
- pass
-
- # Try UUID
- if _is_uuid(str(identifier)):
- return DatasetDAO.find_by_id(identifier, id_column="uuid",
query_options=opts)
-
- return None
-
-
_NO_SAVED_METRICS_HINT = (
"This dataset has no saved metrics, and query_dataset accepts saved "
"metric names only (not ad-hoc expressions). Use execute_sql for "
@@ -170,7 +144,7 @@ async def query_dataset( # noqa: C901
]
with event_logger.log_context(action="mcp.query_dataset.lookup"):
- dataset = _resolve_dataset(request.dataset_id, eager_options)
+ dataset = resolve_dataset(request.dataset_id, eager_options)
if dataset is None:
await ctx.error("Dataset not found: identifier=%s" %
(request.dataset_id,))
diff --git a/superset/mcp_service/dataset/tool/update_dataset_metric.py
b/superset/mcp_service/dataset/tool/update_dataset_metric.py
new file mode 100644
index 00000000000..f7f29b0a44e
--- /dev/null
+++ b/superset/mcp_service/dataset/tool/update_dataset_metric.py
@@ -0,0 +1,294 @@
+# 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.
+
+"""Update a saved metric on a dataset (FastMCP tool)."""
+
+import difflib
+import logging
+from typing import Any
+
+from fastmcp import Context
+from superset_core.mcp.decorators import tool, ToolAnnotations
+
+from superset import security_manager
+from superset.extensions import event_logger
+from superset.mcp_service.dataset.schemas import (
+ DatasetMetricDetail,
+ MetricCurrency,
+ UpdateDatasetMetricRequest,
+ UpdateDatasetMetricResponse,
+)
+from superset.mcp_service.utils import (
+ escape_llm_context_delimiters,
+ sanitize_for_llm_context,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _find_metric(metrics: list[Any], identifier: int | str) -> Any | None:
+ """Find a saved metric by ID, UUID, or metric_name."""
+ if isinstance(identifier, str):
+ try:
+ identifier = int(identifier)
+ except ValueError:
+ pass
+
+ if isinstance(identifier, int):
+ return next((m for m in metrics if m.id == identifier), None)
+
+ lowered = identifier.lower()
+ by_uuid = next(
+ (m for m in metrics if m.uuid and str(m.uuid).lower() == lowered), None
+ )
+ if by_uuid is not None:
+ return by_uuid
+ return next((m for m in metrics if m.metric_name == identifier), None)
+
+
+def _metric_not_found_message(metrics: list[Any], identifier: int | str) ->
str:
+ """Build a "metric not found" error, escaping caller- and stored-supplied
+ text so it can't break out of the LLM context delimiters (same treatment as
+ the success path in ``_serialize_metric``)."""
+ names = [m.metric_name for m in metrics]
+ safe_identifier = escape_llm_context_delimiters(str(identifier))
+ msg = f"Metric '{safe_identifier}' not found on this dataset."
+ if not names:
+ return f"{msg} This dataset has no saved metrics."
+ suggestions = difflib.get_close_matches(str(identifier), names, n=3,
cutoff=0.6)
+ if suggestions:
+ safe_suggestions = [escape_llm_context_delimiters(n) for n in
suggestions]
+ return f"{msg} Did you mean: {', '.join(safe_suggestions)}?"
+ safe_names = [escape_llm_context_delimiters(n) for n in sorted(names)]
+ return f"{msg} Available metrics: {', '.join(safe_names)}."
+
+
+def _serialize_metric(metric: Any) -> DatasetMetricDetail:
+ """Build a ``DatasetMetricDetail`` from a ``SqlMetric`` model.
+
+ Returns the metric's identifiers (id, uuid) and all updatable properties,
+ wrapping free-text fields in LLM-context sanitization the same way the
+ dataset read path does.
+ """
+ currency = getattr(metric, "currency", None)
+ return DatasetMetricDetail(
+ id=getattr(metric, "id", None),
+ uuid=str(metric.uuid) if getattr(metric, "uuid", None) else None,
+ metric_name=escape_llm_context_delimiters(metric.metric_name) or "",
+ verbose_name=sanitize_for_llm_context(
+ getattr(metric, "verbose_name", None),
+ field_path=("metric", "verbose_name"),
+ ),
+ expression=sanitize_for_llm_context(
+ getattr(metric, "expression", None),
+ field_path=("metric", "expression"),
+ ),
+ description=sanitize_for_llm_context(
+ getattr(metric, "description", None),
+ field_path=("metric", "description"),
+ ),
+ d3format=getattr(metric, "d3format", None),
+ metric_type=getattr(metric, "metric_type", None),
+ currency=MetricCurrency.model_validate(currency)
+ if isinstance(currency, dict)
+ else None,
+ warning_text=sanitize_for_llm_context(
+ getattr(metric, "warning_text", None),
+ field_path=("metric", "warning_text"),
+ ),
+ extra=sanitize_for_llm_context(
+ getattr(metric, "extra", None),
+ field_path=("metric", "extra"),
+ ),
+ )
+
+
+@tool(
+ tags=["mutate"],
+ class_permission_name="Dataset",
+ method_permission_name="write",
+ annotations=ToolAnnotations(
+ title="Update dataset metric",
+ readOnlyHint=False,
+ # Overwrites an existing metric's definition and affects every chart
+ # that uses it — non-additive, like update_chart.
+ destructiveHint=True,
+ ),
+)
+async def update_dataset_metric( # noqa: C901
+ request: UpdateDatasetMetricRequest, ctx: Context
+) -> UpdateDatasetMetricResponse:
+ """Update a saved metric on a dataset.
+
+ Modifies properties of an existing saved metric (the named, reusable
+ aggregations stored on a dataset, e.g. SUM(revenue)) — such as its SQL
+ expression, name, display label, or number format. Changes apply to
+ every chart that uses the metric.
+
+ Only the properties you pass are changed; everything else on the metric,
+ and all other metrics on the dataset, is left untouched. This tool cannot
+ create or delete metrics. Requires ownership of the dataset (or Admin).
+
+ Workflow:
+ 1. Call get_dataset_info to inspect the dataset's saved metrics
+ 2. Call this tool with the dataset ID, the metric identifier, and only
+ the properties to change
+
+ Example usage:
+ ```json
+ {
+ "dataset_id": 123,
+ "metric": "sum_revenue",
+ "expression": "SUM(net_revenue)",
+ "d3format": "$,.2f"
+ }
+ ```
+ """
+ updates = request.updates()
+ await ctx.info(
+ "Updating dataset metric: dataset_id=%s, metric=%r, properties=%s"
+ % (request.dataset_id, request.metric, sorted(updates))
+ )
+
+ try:
+ from sqlalchemy.orm import joinedload, subqueryload
+
+ from superset.commands.dataset.exceptions import (
+ DatasetForbiddenError,
+ DatasetInvalidError,
+ DatasetNotFoundError,
+ DatasetUpdateFailedError,
+ )
+ from superset.commands.dataset.update import UpdateDatasetCommand
+ from superset.connectors.sqla.models import SqlaTable
+ from superset.exceptions import SupersetSecurityException
+ from superset.mcp_service.dataset.dataset_utils import resolve_dataset
+ from superset.mcp_service.utils.url_utils import get_superset_base_url
+
+ eager_options = [
+ subqueryload(SqlaTable.metrics),
+ joinedload(SqlaTable.database),
+ ]
+
+ with
event_logger.log_context(action="mcp.update_dataset_metric.lookup"):
+ dataset = resolve_dataset(request.dataset_id, eager_options)
+
+ if dataset is None:
+ await ctx.warning("Dataset not found: %s" % (request.dataset_id,))
+ return UpdateDatasetMetricResponse(
+ error=(
+ f"No dataset found with identifier: {request.dataset_id}."
+ " Use list_datasets to get valid dataset IDs."
+ ),
+ )
+
+ # Enforce editorship before touching the metric list. Otherwise a
caller
+ # with read access but not edit rights could distinguish "metric
exists"
+ # from "metric not found" and enumerate metric names via the not-found
+ # suggestions. UpdateDatasetCommand re-checks this, but only after the
+ # lookup below would have already leaked that information.
+ try:
+ security_manager.raise_for_editorship(dataset)
+ except SupersetSecurityException:
+ await ctx.warning(
+ "Dataset metric update forbidden: dataset_id=%s" %
(dataset.id,)
+ )
+ return UpdateDatasetMetricResponse(
+ error="You must be an owner of this dataset (or an Admin) "
+ "to update its metrics.",
+ )
+
+ metrics = list(dataset.metrics)
+ target = _find_metric(metrics, request.metric)
+ if target is None:
+ message = _metric_not_found_message(metrics, request.metric)
+ await ctx.warning("Metric not found: %s" % (request.metric,))
+ return UpdateDatasetMetricResponse(
+ dataset_id=dataset.id,
+ dataset_name=dataset.table_name,
+ error=message,
+ )
+
+ # DatasetDAO.update_metrics deletes any metric missing from the list,
+ # so every metric is sent back: untouched ones as bare
+ # {id, metric_name} stubs (the bulk update merges them over the stored
+ # row), and the target with the requested changes applied on top.
+ metrics_payload: list[dict[str, Any]] = []
+ for metric in metrics:
+ entry: dict[str, Any] = {
+ "id": metric.id,
+ "metric_name": metric.metric_name,
+ }
+ if metric.id == target.id:
+ entry.update(updates)
+ metrics_payload.append(entry)
+
+ # UpdateDatasetCommand enforces dataset ownership and validates
+ # metric name uniqueness and the SQL expression
+ with
event_logger.log_context(action="mcp.update_dataset_metric.update"):
+ updated_dataset = UpdateDatasetCommand(
+ dataset.id, {"metrics": metrics_payload}
+ ).run()
+
+ updated_metric = next(
+ (m for m in updated_dataset.metrics if m.id == target.id), None
+ )
+
+ await ctx.info(
+ "Dataset metric updated: dataset_id=%s, metric_id=%s,
properties=%s"
+ % (updated_dataset.id, target.id, sorted(updates))
+ )
+
+ return UpdateDatasetMetricResponse(
+ dataset_id=updated_dataset.id,
+ dataset_name=updated_dataset.table_name,
+ metric=_serialize_metric(updated_metric) if updated_metric else
None,
+ updated_properties=sorted(updates),
+ url=(
+ f"{get_superset_base_url()}/explore/"
+ f"?datasource_type=table&datasource_id={updated_dataset.id}"
+ ),
+ )
+
+ except DatasetNotFoundError:
+ await ctx.warning("Dataset not found: %s" % (request.dataset_id,))
+ return UpdateDatasetMetricResponse(
+ error=f"No dataset found with identifier: {request.dataset_id}.",
+ )
+ except DatasetForbiddenError:
+ await ctx.warning(
+ "Dataset metric update forbidden: dataset_id=%s" %
(request.dataset_id,)
+ )
+ return UpdateDatasetMetricResponse(
+ error="You must be an owner of this dataset (or an Admin) "
+ "to update its metrics.",
+ )
+ except DatasetInvalidError as exc:
+ messages = exc.normalized_messages()
+ await ctx.warning("Dataset metric validation failed: %s" % (messages,))
+ return UpdateDatasetMetricResponse(error=str(messages))
+ except DatasetUpdateFailedError as exc:
+ await ctx.error("Dataset metric update failed: %s" % (str(exc),))
+ return UpdateDatasetMetricResponse(
+ error=f"Failed to update metric: {exc}",
+ )
+ except Exception as exc:
+ await ctx.error(
+ "Unexpected error updating dataset metric: %s: %s"
+ % (type(exc).__name__, str(exc))
+ )
+ raise
diff --git a/tests/unit_tests/mcp_service/dataset/tool/test_query_dataset.py
b/tests/unit_tests/mcp_service/dataset/tool/test_query_dataset.py
index eb9e241d283..08e17b5fbe0 100644
--- a/tests/unit_tests/mcp_service/dataset/tool/test_query_dataset.py
+++ b/tests/unit_tests/mcp_service/dataset/tool/test_query_dataset.py
@@ -143,7 +143,7 @@ async def test_query_dataset_success(mcp_server: FastMCP)
-> None:
with (
patch.object(
query_dataset_module,
- "_resolve_dataset",
+ "resolve_dataset",
return_value=dataset,
),
patch(
@@ -183,7 +183,7 @@ async def test_query_dataset_not_found(mcp_server: FastMCP)
-> None:
"""Dataset ID that doesn't exist returns error."""
with patch.object(
query_dataset_module,
- "_resolve_dataset",
+ "resolve_dataset",
return_value=None,
):
async with Client(mcp_server) as client:
@@ -209,7 +209,7 @@ async def test_query_dataset_invalid_metric(mcp_server:
FastMCP) -> None:
with patch.object(
query_dataset_module,
- "_resolve_dataset",
+ "resolve_dataset",
return_value=dataset,
):
async with Client(mcp_server) as client:
@@ -237,7 +237,7 @@ async def test_query_dataset_invalid_column(mcp_server:
FastMCP) -> None:
with patch.object(
query_dataset_module,
- "_resolve_dataset",
+ "resolve_dataset",
return_value=dataset,
):
async with Client(mcp_server) as client:
@@ -290,7 +290,7 @@ async def test_query_dataset_with_time_range(mcp_server:
FastMCP) -> None:
with (
patch.object(
query_dataset_module,
- "_resolve_dataset",
+ "resolve_dataset",
return_value=dataset,
),
patch(
@@ -342,7 +342,7 @@ async def
test_query_dataset_time_range_no_temporal_column(mcp_server: FastMCP)
with patch.object(
query_dataset_module,
- "_resolve_dataset",
+ "resolve_dataset",
return_value=dataset,
):
async with Client(mcp_server) as client:
@@ -376,7 +376,7 @@ async def test_query_dataset_with_filters(mcp_server:
FastMCP) -> None:
with (
patch.object(
query_dataset_module,
- "_resolve_dataset",
+ "resolve_dataset",
return_value=dataset,
),
patch(
@@ -433,7 +433,7 @@ async def test_query_dataset_empty_results(mcp_server:
FastMCP) -> None:
with (
patch.object(
query_dataset_module,
- "_resolve_dataset",
+ "resolve_dataset",
return_value=dataset,
),
patch(
@@ -474,7 +474,7 @@ async def test_query_dataset_by_uuid(mcp_server: FastMCP)
-> None:
with (
patch.object(
query_dataset_module,
- "_resolve_dataset",
+ "resolve_dataset",
return_value=dataset,
) as mock_resolve,
patch(
@@ -520,7 +520,7 @@ async def test_query_dataset_permission_denied(mcp_server:
FastMCP) -> None:
with (
patch.object(
query_dataset_module,
- "_resolve_dataset",
+ "resolve_dataset",
return_value=dataset,
),
patch(
@@ -567,7 +567,7 @@ async def test_query_dataset_order_by_valid(mcp_server:
FastMCP) -> None:
with (
patch.object(
query_dataset_module,
- "_resolve_dataset",
+ "resolve_dataset",
return_value=dataset,
),
patch(
@@ -611,7 +611,7 @@ async def test_query_dataset_order_by_invalid(mcp_server:
FastMCP) -> None:
with patch.object(
query_dataset_module,
- "_resolve_dataset",
+ "resolve_dataset",
return_value=dataset,
):
async with Client(mcp_server) as client:
@@ -645,7 +645,7 @@ async def
test_query_dataset_time_column_override(mcp_server: FastMCP) -> None:
with (
patch.object(
query_dataset_module,
- "_resolve_dataset",
+ "resolve_dataset",
return_value=dataset,
),
patch(
@@ -689,7 +689,7 @@ async def
test_query_dataset_non_dttm_time_column_warns(mcp_server: FastMCP) ->
with (
patch.object(
query_dataset_module,
- "_resolve_dataset",
+ "resolve_dataset",
return_value=dataset,
),
patch(
@@ -729,7 +729,7 @@ async def
test_query_dataset_invalid_filter_column(mcp_server: FastMCP) -> None:
with patch.object(
query_dataset_module,
- "_resolve_dataset",
+ "resolve_dataset",
return_value=dataset,
):
async with Client(mcp_server) as client:
@@ -769,7 +769,7 @@ async def
test_query_dataset_metadata_access_denied_no_suggestions(
with (
patch.object(
query_dataset_module,
- "_resolve_dataset",
+ "resolve_dataset",
return_value=dataset,
),
patch.object(
diff --git
a/tests/unit_tests/mcp_service/dataset/tool/test_update_dataset_metric.py
b/tests/unit_tests/mcp_service/dataset/tool/test_update_dataset_metric.py
new file mode 100644
index 00000000000..d533498b26f
--- /dev/null
+++ b/tests/unit_tests/mcp_service/dataset/tool/test_update_dataset_metric.py
@@ -0,0 +1,541 @@
+# 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.
+
+from collections.abc import Iterator
+from types import SimpleNamespace
+from typing import Any
+from unittest.mock import MagicMock, patch
+
+import pytest
+from fastmcp import Client, FastMCP
+from pydantic import ValidationError
+
+from superset.mcp_service.app import mcp
+from superset.mcp_service.dataset.schemas import UpdateDatasetMetricRequest
+from superset.mcp_service.utils.sanitization import (
+ LLM_CONTEXT_CLOSE_DELIMITER,
+ LLM_CONTEXT_OPEN_DELIMITER,
+)
+from superset.utils import json
+
+
+def _wrapped(value: str) -> str:
+ return
f"{LLM_CONTEXT_OPEN_DELIMITER}\n{value}\n{LLM_CONTEXT_CLOSE_DELIMITER}"
+
+
[email protected]
+def mcp_server() -> FastMCP:
+ """Provide the shared FastMCP app instance for the in-process test
client."""
+ return mcp
+
+
[email protected](autouse=True)
+def mock_auth() -> Iterator[MagicMock]:
+ """Mock authentication for all tests."""
+ from unittest.mock import Mock
+
+ 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
+
+
[email protected](autouse=True)
+def allow_ownership() -> Iterator[MagicMock]:
+ """Let ownership checks pass by default; override to simulate a non-owner.
+
+ Patches the security manager class method (rather than the module-level
+ proxy) so it resolves without an app context, as in CI.
+ """
+ with patch(
+ "superset.security.SupersetSecurityManager.raise_for_editorship",
+ return_value=None,
+ ) as mock_raise:
+ yield mock_raise
+
+
+def make_metric(
+ metric_id: int = 10,
+ metric_name: str = "count",
+ uuid: str = "a1b2c3d4-5678-90ab-cdef-1234567890ab",
+ **overrides: Any,
+) -> SimpleNamespace:
+ """Build a stand-in SqlMetric with all serialized attributes populated."""
+ metric = SimpleNamespace(
+ id=metric_id,
+ uuid=uuid,
+ metric_name=metric_name,
+ verbose_name=None,
+ expression="COUNT(*)",
+ description=None,
+ d3format=None,
+ metric_type=None,
+ currency=None,
+ warning_text=None,
+ )
+ for key, value in overrides.items():
+ setattr(metric, key, value)
+ return metric
+
+
+def make_dataset(
+ dataset_id: int = 1, metrics: list[SimpleNamespace] | None = None
+) -> MagicMock:
+ """Build a stand-in SqlaTable exposing an id, table_name, and metrics."""
+ dataset = MagicMock()
+ dataset.id = dataset_id
+ dataset.table_name = "my_table"
+ dataset.metrics = metrics if metrics is not None else [make_metric()]
+ return dataset
+
+
+# ---------------------------------------------------------------------------
+# Request schema validation
+# ---------------------------------------------------------------------------
+
+
+def test_request_requires_at_least_one_property() -> None:
+ with pytest.raises(ValidationError, match="At least one metric property"):
+ UpdateDatasetMetricRequest(dataset_id=1, metric="count")
+
+
+def test_request_rejects_empty_expression() -> None:
+ with pytest.raises(ValidationError, match="expression cannot be empty"):
+ UpdateDatasetMetricRequest(dataset_id=1, metric="count", expression="
")
+
+
+def test_request_rejects_null_metric_name() -> None:
+ with pytest.raises(ValidationError, match="metric_name cannot be empty"):
+ UpdateDatasetMetricRequest(dataset_id=1, metric="count",
metric_name=None)
+
+
+def test_request_updates_only_includes_provided_properties() -> None:
+ request = UpdateDatasetMetricRequest.model_validate(
+ {"dataset_id": 1, "metric": "count", "expression": "COUNT(1)"}
+ )
+ assert request.updates() == {"expression": "COUNT(1)"}
+
+
+def test_request_updates_keeps_explicit_nulls() -> None:
+ request = UpdateDatasetMetricRequest.model_validate(
+ {"dataset_id": 1, "metric": "count", "description": None, "d3format":
",.2f"}
+ )
+ assert request.updates() == {"description": None, "d3format": ",.2f"}
+
+
+def test_request_rejects_invalid_extra_json() -> None:
+ with pytest.raises(ValidationError, match="extra must be a valid JSON"):
+ UpdateDatasetMetricRequest(dataset_id=1, metric="count", extra="{not
json")
+
+
+def test_request_accepts_valid_extra_json() -> None:
+ request = UpdateDatasetMetricRequest.model_validate(
+ {"dataset_id": 1, "metric": "count", "extra": '{"warning_markdown":
"hi"}'}
+ )
+ assert request.updates() == {"extra": '{"warning_markdown": "hi"}'}
+
+
+def test_request_updates_dumps_nested_currency() -> None:
+ request = UpdateDatasetMetricRequest.model_validate(
+ {
+ "dataset_id": 1,
+ "metric": 10,
+ "currency": {"symbol": "USD", "symbolPosition": "prefix"},
+ }
+ )
+ assert request.updates() == {
+ "currency": {"symbol": "USD", "symbolPosition": "prefix"}
+ }
+
+
+# ---------------------------------------------------------------------------
+# Tool behavior
+# ---------------------------------------------------------------------------
+
+
[email protected]
+async def test_update_dataset_metric_success(mcp_server: FastMCP) -> None:
+ """Happy path: only the target metric is changed, others sent as stubs."""
+ target = make_metric(metric_id=10, metric_name="count")
+ other = make_metric(
+ metric_id=11,
+ metric_name="sum_revenue",
+ uuid="b2c3d4e5-6789-01bc-def0-234567890abc",
+ expression="SUM(revenue)",
+ )
+ dataset = make_dataset(dataset_id=1, metrics=[target, other])
+
+ updated_target = make_metric(
+ metric_id=10, metric_name="count", expression="COUNT(1)", d3format=",d"
+ )
+ updated_dataset = make_dataset(dataset_id=1, metrics=[updated_target,
other])
+
+ mock_command = MagicMock()
+ mock_command.run.return_value = updated_dataset
+
+ with (
+ patch(
+ "superset.daos.dataset.DatasetDAO.find_by_id",
+ return_value=dataset,
+ ),
+ patch(
+ "superset.commands.dataset.update.UpdateDatasetCommand",
+ return_value=mock_command,
+ ) as 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:
+ result = await client.call_tool(
+ "update_dataset_metric",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metric": "count",
+ "expression": "COUNT(1)",
+ "d3format": ",d",
+ }
+ },
+ )
+ data = json.loads(result.content[0].text)
+
+ command_cls.assert_called_once_with(
+ 1,
+ {
+ "metrics": [
+ {
+ "id": 10,
+ "metric_name": "count",
+ "expression": "COUNT(1)",
+ "d3format": ",d",
+ },
+ {"id": 11, "metric_name": "sum_revenue"},
+ ]
+ },
+ )
+ assert data["error"] is None
+ assert data["dataset_id"] == 1
+ assert data["dataset_name"] == "my_table"
+ assert data["updated_properties"] == ["d3format", "expression"]
+ assert data["metric"]["id"] == 10
+ assert data["metric"]["expression"] == _wrapped("COUNT(1)")
+ assert data["metric"]["d3format"] == ",d"
+ assert "datasource_id=1" in data["url"]
+
+
[email protected]
+async def test_update_dataset_metric_returns_extra(mcp_server: FastMCP) ->
None:
+ """`extra` is an updatable property, so it is echoed back in the
response."""
+ dataset = make_dataset(dataset_id=1, metrics=[make_metric(metric_id=10)])
+ updated = make_dataset(
+ dataset_id=1,
+ metrics=[make_metric(metric_id=10, extra='{"warning_markdown":
"note"}')],
+ )
+ mock_command = MagicMock()
+ mock_command.run.return_value = updated
+
+ with (
+ patch(
+ "superset.daos.dataset.DatasetDAO.find_by_id",
+ return_value=dataset,
+ ),
+ patch(
+ "superset.commands.dataset.update.UpdateDatasetCommand",
+ 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(
+ "update_dataset_metric",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metric": 10,
+ "extra": '{"warning_markdown": "note"}',
+ }
+ },
+ )
+ data = json.loads(result.content[0].text)
+
+ assert data["error"] is None
+ assert data["updated_properties"] == ["extra"]
+ assert data["metric"]["extra"] == _wrapped('{"warning_markdown": "note"}')
+
+
[email protected]
+async def test_update_dataset_metric_by_id_and_uuid(mcp_server: FastMCP) ->
None:
+ """The metric can be addressed by numeric ID (also as string) or UUID."""
+ target = make_metric(metric_id=10, metric_name="count")
+ dataset = make_dataset(dataset_id=1, metrics=[target])
+
+ mock_command = MagicMock()
+ mock_command.run.return_value = dataset
+
+ for identifier in (10, "10", "A1B2C3D4-5678-90ab-cdef-1234567890ab"):
+ mock_command.reset_mock()
+ with (
+ patch(
+ "superset.daos.dataset.DatasetDAO.find_by_id",
+ return_value=dataset,
+ ),
+ patch(
+ "superset.commands.dataset.update.UpdateDatasetCommand",
+ 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(
+ "update_dataset_metric",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metric": identifier,
+ "verbose_name": "Count",
+ }
+ },
+ )
+ data = json.loads(result.content[0].text)
+
+ assert data["error"] is None, identifier
+ mock_command.run.assert_called_once()
+
+
[email protected]
+async def test_update_dataset_metric_not_found_suggests_names(
+ mcp_server: FastMCP,
+) -> None:
+ dataset = make_dataset(
+ dataset_id=1, metrics=[make_metric(metric_id=10,
metric_name="sum_revenue")]
+ )
+
+ with patch(
+ "superset.daos.dataset.DatasetDAO.find_by_id",
+ return_value=dataset,
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "update_dataset_metric",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metric": "sum_revenu",
+ "expression": "SUM(net_revenue)",
+ }
+ },
+ )
+ data = json.loads(result.content[0].text)
+
+ assert data["metric"] is None
+ assert "not found" in data["error"]
+ assert "sum_revenue" in data["error"]
+
+
[email protected]
+async def test_update_dataset_metric_not_found_escapes_names(
+ mcp_server: FastMCP,
+) -> None:
+ """Metric names in the not-found error are escaped like the success
path."""
+ from superset.mcp_service.utils.sanitization import (
+ LLM_CONTEXT_ESCAPED_OPEN_DELIMITER,
+ )
+
+ hostile_name = f"{LLM_CONTEXT_OPEN_DELIMITER}evil"
+ dataset = make_dataset(
+ dataset_id=1, metrics=[make_metric(metric_id=10,
metric_name=hostile_name)]
+ )
+
+ with patch(
+ "superset.daos.dataset.DatasetDAO.find_by_id",
+ return_value=dataset,
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "update_dataset_metric",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metric": "missing",
+ "expression": "COUNT(1)",
+ }
+ },
+ )
+ data = json.loads(result.content[0].text)
+
+ assert data["metric"] is None
+ assert LLM_CONTEXT_OPEN_DELIMITER not in data["error"]
+ assert LLM_CONTEXT_ESCAPED_OPEN_DELIMITER in data["error"]
+
+
[email protected]
+async def test_update_dataset_metric_dataset_not_found(mcp_server: FastMCP) ->
None:
+ with patch(
+ "superset.daos.dataset.DatasetDAO.find_by_id",
+ return_value=None,
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "update_dataset_metric",
+ {
+ "request": {
+ "dataset_id": 999,
+ "metric": "count",
+ "description": "desc",
+ }
+ },
+ )
+ data = json.loads(result.content[0].text)
+
+ assert data["metric"] is None
+ assert "No dataset found" in data["error"]
+
+
[email protected]
+async def test_update_dataset_metric_forbidden_non_owner(
+ mcp_server: FastMCP, allow_ownership: MagicMock
+) -> None:
+ """A non-owner is rejected before the metric list is inspected, so the
+ error cannot be used to enumerate metric names."""
+ from superset.exceptions import SupersetSecurityException
+
+ allow_ownership.side_effect = SupersetSecurityException(
+ MagicMock(message="not an owner")
+ )
+
+ dataset = make_dataset(
+ dataset_id=1,
+ metrics=[make_metric(metric_id=10, metric_name="secret_metric")],
+ )
+ mock_command = MagicMock()
+
+ with (
+ patch(
+ "superset.daos.dataset.DatasetDAO.find_by_id",
+ return_value=dataset,
+ ),
+ patch(
+ "superset.commands.dataset.update.UpdateDatasetCommand",
+ return_value=mock_command,
+ ) as command_cls,
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "update_dataset_metric",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metric": "does_not_exist",
+ "expression": "COUNT(1)",
+ }
+ },
+ )
+ data = json.loads(result.content[0].text)
+
+ assert data["metric"] is None
+ assert "owner" in data["error"]
+ # No metric names leak, and the command is never constructed/run.
+ assert "secret_metric" not in data["error"]
+ command_cls.assert_not_called()
+
+
[email protected]
+async def test_update_dataset_metric_invalid_error(mcp_server: FastMCP) ->
None:
+ """Validation failures (e.g. duplicate name) surface as error responses."""
+ from superset.commands.dataset.exceptions import (
+ DatasetInvalidError,
+ DatasetMetricsDuplicateValidationError,
+ )
+
+ invalid_exc = DatasetInvalidError()
+ invalid_exc.append(DatasetMetricsDuplicateValidationError())
+
+ dataset = make_dataset()
+ mock_command = MagicMock()
+ mock_command.run.side_effect = invalid_exc
+
+ with (
+ patch(
+ "superset.daos.dataset.DatasetDAO.find_by_id",
+ return_value=dataset,
+ ),
+ patch(
+ "superset.commands.dataset.update.UpdateDatasetCommand",
+ return_value=mock_command,
+ ),
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "update_dataset_metric",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metric": "count",
+ "metric_name": "sum_revenue",
+ }
+ },
+ )
+ data = json.loads(result.content[0].text)
+
+ assert data["metric"] is None
+ assert data["error"] is not None
+
+
[email protected]
+async def test_update_dataset_metric_update_failed(mcp_server: FastMCP) ->
None:
+ """A persistence failure surfaces as a "Failed to update" error
response."""
+ from superset.commands.dataset.exceptions import DatasetUpdateFailedError
+
+ dataset = make_dataset()
+ mock_command = MagicMock()
+ mock_command.run.side_effect = DatasetUpdateFailedError()
+
+ with (
+ patch(
+ "superset.daos.dataset.DatasetDAO.find_by_id",
+ return_value=dataset,
+ ),
+ patch(
+ "superset.commands.dataset.update.UpdateDatasetCommand",
+ return_value=mock_command,
+ ),
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "update_dataset_metric",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metric": "count",
+ "expression": "COUNT(1)",
+ }
+ },
+ )
+ data = json.loads(result.content[0].text)
+
+ assert data["metric"] is None
+ assert "Failed to update" in data["error"]
