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/10/22 08:36:10 UTC
(superset) branch master updated: fix(mcp): pass time_grain as a string and keep annotation DB error classification (#45484)
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 47300cf4c2f fix(mcp): pass time_grain as a string and keep annotation DB error classification (#45484)
47300cf4c2f is described below
commit 47300cf4c2f13aa3a4bef6c7f349f58397de0035
Author: Amin Ghadersohi <am...@gmail.com>
AuthorDate: Thu Oct 22 04:36:02 2026 -0400
fix(mcp): pass time_grain as a string and keep annotation DB error classification (#45484)
---
.../docs/using-superset/using-ai-with-superset.mdx | 20 +
superset/mcp_service/chart/chart_helpers.py | 23 +
superset/mcp_service/chart/chart_utils.py | 16 +-
superset/mcp_service/chart/compile.py | 176 ++++++--
.../mcp_service/chart/test_chart_helpers.py | 42 ++
tests/unit_tests/mcp_service/chart/test_compile.py | 464 +++++++++++++++++++++
.../test_review_query_and_preview_regressions.py | 7 +-
.../chart/test_time_grain_generic_axes.py | 171 ++++++++
8 files changed, 886 insertions(+), 33 deletions(-)
diff --git a/docs/docs/using-superset/using-ai-with-superset.mdx b/docs/docs/using-superset/using-ai-with-superset.mdx
index 1d475397d08..98d13ad60a3 100644
--- a/docs/docs/using-superset/using-ai-with-superset.mdx
+++ b/docs/docs/using-superset/using-ai-with-superset.mdx
@@ -106,6 +106,10 @@ comparison preserves other selected offsets, including custom comparisons.
Forecasts use the x-axis time grain, resolved query grain, or panel grain, with a
daily fallback.
+When chart compilation fails, MCP distinguishes database connection failures from
+query validation errors. For Mixed Timeseries charts, annotation failures on
+either query use the annotation charts' database engines for classification.
+
### Create Dashboards
Build dashboards from a collection of charts:
@@ -473,6 +477,22 @@ also have `data: null`. For each statement, `executed_sql` is `null` when it equ
- Ask for smaller result sets: use filters, reduce `page_size`, or request specific columns
- Example: "Show me the top 10 rows from the sales dataset" instead of "Show me all sales data"
+### Chart query and template errors
+
+MCP chart tools accept ISO-8601 time grains such as `P1M` for monthly
+buckets, including Big Number trendlines with saved metrics. Template-rendering
+and query-processing validation failures are reported as `compile_error`, with
+the underlying error in `details`, rather than as database connection failures.
+Check the chart configuration or dataset expressions before retrying.
+Connection failures in an annotation layer's chart query are still reported as
+`database_connection_error`, even when the annotation layer wraps the failure as
+a validation error. Typed database exceptions are recognized directly; other
+failures are matched against candidate annotation-chart database engines. If
+multiple engines are candidates or a chart cannot be located, only specific
+connection-error patterns count, not an engine's generic error fallback. If no
+annotation engine can be determined, the main dataset's engine supplies those
+patterns. The error message names the annotation layer as the source.
+
### AI doesn't see Superset tools
- Verify the connection in your AI client (e.g., the hammer icon in Claude Desktop)
diff --git a/superset/mcp_service/chart/chart_helpers.py b/superset/mcp_service/chart/chart_helpers.py
index 92306bb97ed..445f9592240 100644
--- a/superset/mcp_service/chart/chart_helpers.py
+++ b/superset/mcp_service/chart/chart_helpers.py
@@ -27,6 +27,8 @@ from __future__ import annotations
import logging
from collections.abc import Mapping
+from copy import deepcopy
+from enum import Enum
from typing import Any, TYPE_CHECKING
from urllib.parse import parse_qs, urlparse
@@ -1148,6 +1150,20 @@ def build_single_query_dict(
return qd
+def coerce_enum_values(value: Any) -> Any:
+ """Replace enum members with their values, in place for dicts and lists."""
+ if isinstance(value, Enum):
+ return value.value
+ if isinstance(value, dict):
+ for key, item in value.items():
+ value[key] = coerce_enum_values(item)
+ elif isinstance(value, list):
+ value[:] = [coerce_enum_values(item) for item in value]
+ elif isinstance(value, tuple):
+ return tuple(coerce_enum_values(item) for item in value)
+ return value
+
+
def build_query_dicts_from_form_data(
form_data: dict[str, Any],
datasource_id: Any,
@@ -1163,12 +1179,19 @@ def build_query_dicts_from_form_data(
``superset.common.form_data_query_context``. This production entry point
only performs server filter normalization and the legacy Deck.gl spatial
augmentation before delegating to that shared frontend-equivalent contract.
+ Normalization uses deep copies to preserve caller-owned form data.
"""
from superset.common.form_data_query_context import (
build_query_objects_from_form_data,
)
from superset.mcp_service.chart.registry import plugin_for_viz_type
+ # Typed chart configs carry enums such as TimeGrain. Every MCP query is
+ # built here, so unwrap them once: template rendering rejects non-primitive
+ # context values.
+ form_data = coerce_enum_values(deepcopy(form_data))
+ if extra_form_data is not None:
+ extra_form_data = coerce_enum_values(deepcopy(extra_form_data))
viz_type: str = (
form_data.get("viz_type")
or (getattr(chart, "viz_type", "") if chart else "")
diff --git a/superset/mcp_service/chart/chart_utils.py b/superset/mcp_service/chart/chart_utils.py
index 3085d19558e..668f01b1ed2 100644
--- a/superset/mcp_service/chart/chart_utils.py
+++ b/superset/mcp_service/chart/chart_utils.py
@@ -36,7 +36,7 @@ if TYPE_CHECKING:
from superset.connectors.sqla.models import SqlaTable
from superset.semantic_layers.models import SemanticView
-from superset.constants import NO_TIME_RANGE
+from superset.constants import NO_TIME_RANGE, TimeGrain
from superset.daos.exceptions import DatasourceNotFound, DatasourceValueIsIncorrect
from superset.exceptions import SupersetSecurityException
from superset.mcp_service.chart.schemas import (
@@ -1510,7 +1510,7 @@ def add_xy_sort_config(
def configure_temporal_handling(
form_data: Dict[str, Any],
x_is_temporal: bool,
- time_grain: str | None,
+ time_grain: TimeGrain | str | None,
) -> None:
"""Configure form_data based on whether x-axis column is temporal.
@@ -1523,7 +1523,9 @@ def configure_temporal_handling(
if x_is_temporal:
form_data["granularity_sqla"] = form_data.get("x_axis")
if time_grain:
- form_data["time_grain_sqla"] = time_grain
+ form_data["time_grain_sqla"] = (
+ time_grain.value if isinstance(time_grain, TimeGrain) else time_grain
+ )
else:
# Non-temporal column - disable temporal handling to prevent DATE_TRUNC
form_data["x_axis_sort_series_type"] = "name"
@@ -2051,7 +2053,7 @@ def map_sunburst_config(config: SunburstChartConfig) -> Dict[str, Any]:
if config.temporal_column is not None:
form_data["granularity_sqla"] = config.temporal_column
if config.time_grain is not None:
- form_data["time_grain_sqla"] = config.time_grain
+ form_data["time_grain_sqla"] = config.time_grain.value
_copy_sunburst_native_envelope(form_data, config)
@@ -2688,7 +2690,7 @@ def finalize_sunburst_update_form_data( # noqa: C901
merged["granularity_sqla"] = config.temporal_column
if grain_set and config.time_grain is not None:
- merged["time_grain_sqla"] = config.time_grain
+ merged["time_grain_sqla"] = config.time_grain.value
if "time_range" in fields_set and config.time_range is not None:
merged["time_range"] = config.time_range
@@ -2909,7 +2911,7 @@ def map_waterfall_config(config: WaterfallChartConfig) -> Dict[str, Any]:
# `x_axis || granularity_sqla`. Providing time_grain signals temporal
# intent; Superset ignores both for a non-temporal column.
if config.time_grain:
- form_data["time_grain_sqla"] = config.time_grain
+ form_data["time_grain_sqla"] = config.time_grain.value
form_data["granularity_sqla"] = config.x_axis.name
add_currency_format(form_data, config.currency_format)
_add_adhoc_filters(form_data, config.filters)
@@ -3034,7 +3036,7 @@ def map_big_number_config(
form_data["start_y_axis_at_zero"] = config.start_y_axis_at_zero
if config.time_grain:
- form_data["time_grain_sqla"] = config.time_grain
+ form_data["time_grain_sqla"] = config.time_grain.value
if config.compare_lag is not None:
form_data["compare_lag"] = config.compare_lag
diff --git a/superset/mcp_service/chart/compile.py b/superset/mcp_service/chart/compile.py
index 76e12f8867a..c0035ab83e0 100644
--- a/superset/mcp_service/chart/compile.py
+++ b/superset/mcp_service/chart/compile.py
@@ -40,6 +40,7 @@ from typing import Any, Dict, List, Literal
from sqlalchemy.exc import SQLAlchemyError
from superset.commands.exceptions import CommandException
+from superset.errors import SupersetErrorType
from superset.exceptions import QueryObjectValidationError
from superset.mcp_service.chart.chart_helpers import canonicalize_operation_form_data
from superset.mcp_service.chart.query_result import (
@@ -216,10 +217,12 @@ def _compile_chart( # noqa: C901
return CompileResult(success=True, warnings=warnings, row_count=row_count)
except (ChartDataQueryFailedError, ChartDataCacheLoadError) as exc:
- if (
- datasource_type == DatasourceType.TABLE.value
- and _classify_as_database_error(exc, dataset_id)
- ):
+ database_failure = (
+ _classify_as_database_error(exc, form_data, dataset_id)
+ if datasource_type == DatasourceType.TABLE.value
+ else None
+ )
+ if database_failure is not None:
logger.warning(
"Database connection error during chart compile check: %s: %s",
type(exc).__name__,
@@ -230,7 +233,9 @@ def _compile_chart( # noqa: C901
error=f"Database connection error: {exc}",
error_code="CHART_COMPILE_FAILED",
tier="compile",
- error_obj=_build_database_error(str(exc)),
+ error_obj=_build_database_error(
+ str(exc), annotation=database_failure == "annotation"
+ ),
)
return CompileResult(
success=False,
@@ -404,43 +409,166 @@ def _active_adhoc_filters(filters: list[Any]) -> list[dict[str, Any]]:
]
-def _classify_as_database_error(exc: BaseException, dataset_id: int) -> bool:
- """Use the dataset's DB engine spec to classify the error.
+@dataclass(frozen=True)
+class _DatabaseFailure:
+ """A database failure found in an exception's ``__cause__`` chain."""
- Walks the ``__cause__`` chain for direct ``SQLAlchemyError`` instances,
- then falls back to the engine spec's ``extract_errors`` regex patterns —
- the same classification the Superset UI uses.
+ error: BaseException
+ # SQLAlchemy errors are database failures without engine pattern matching.
+ typed: bool
+ # Found beneath a validation wrapper, as annotation query failures are.
+ wrapped: bool
+
+
+def _find_database_failure(exc: BaseException) -> _DatabaseFailure | None:
+ """Find the deepest database failure in *exc*'s ``__cause__`` chain.
+
+ Validation wrappers can carry an annotation chart's database failure, so
+ they are walked through but never classified themselves.
"""
- # Direct SQLAlchemy errors (unwrapped or in cause chain)
+ from superset.commands.chart.exceptions import (
+ ChartDataQueryFailedError,
+ ChartDataQueryValidationError,
+ )
+ from superset.db_engine_specs.exceptions import SupersetDBAPIError
+
+ # A non-validation outer error (e.g. a cache load failure) keeps the engine
+ # pattern fallback even when no typed database error is in its chain.
+ failure = (
+ None
+ if isinstance(exc, (ChartDataQueryValidationError, QueryObjectValidationError))
+ else _DatabaseFailure(exc, typed=False, wrapped=False)
+ )
+ wrapped = False
+ seen: set[int] = set()
current: BaseException | None = exc
- while current is not None:
+ while current is not None and id(current) not in seen:
+ seen.add(id(current))
if isinstance(current, SQLAlchemyError):
- return True
+ return _DatabaseFailure(current, typed=True, wrapped=wrapped)
+ if isinstance(
+ current, (ChartDataQueryValidationError, QueryObjectValidationError)
+ ):
+ wrapped = True
+ elif isinstance(current, (ChartDataQueryFailedError, SupersetDBAPIError)):
+ failure = _DatabaseFailure(current, typed=False, wrapped=wrapped)
current = current.__cause__
+ return failure
- # Use the dataset's engine spec to classify (same as the UI)
- try:
- from superset.daos.dataset import DatasetDAO
- dataset = DatasetDAO.find_by_id(dataset_id)
- if dataset and dataset.database and isinstance(exc, Exception):
- errors = dataset.database.db_engine_spec.extract_errors(exc)
- return any(e.error_type in CONNECTION_ERROR_TYPES for e in errors)
+def _chart_annotation_layers(form_data: dict[str, Any]) -> list[dict[str, Any]]:
+ """Collect chart-backed annotation layers from both Mixed Timeseries queries."""
+ from superset.utils.core import ANNOTATION_SOURCE_TYPES_WITH_CHART_REFERENCE
+
+ return [
+ layer
+ for key in ("annotation_layers", "annotation_layers_b")
+ if isinstance(layers := form_data.get(key), list)
+ for layer in layers
+ if isinstance(layer, dict)
+ and layer.get("sourceType") in ANNOTATION_SOURCE_TYPES_WITH_CHART_REFERENCE
+ ]
+
+
+def _annotation_engine_specs(form_data: dict[str, Any]) -> tuple[list[Any], bool]:
+ """Return the engine specs of the charts the annotation layers query.
+
+ The flag is True when the failing layer's engine is certain: every chart
+ layer resolved to the same engine. The exception does not name the layer
+ that raised it, so with several engines (or an unresolved chart) any of
+ them could be the failing one.
+ """
+ from superset.daos.chart import ChartDAO
+
+ specs: list[Any] = []
+ unresolved = False
+ for layer in _chart_annotation_layers(form_data):
+ if not isinstance(chart_id := layer.get("value"), (int, str)):
+ continue
+ chart = ChartDAO.find_by_id(chart_id)
+ database = getattr(getattr(chart, "datasource", None), "database", None)
+ spec = getattr(database, "db_engine_spec", None)
+ if spec is None:
+ unresolved = True
+ elif spec not in specs:
+ specs.append(spec)
+ return specs, len(specs) == 1 and not unresolved
+
+
+def _has_chart_annotation_layers(form_data: dict[str, Any]) -> bool:
+ """Whether the chart runs annotation queries against other charts."""
+ return bool(_chart_annotation_layers(form_data))
+
+
+def _classify_as_database_error(
+ exc: BaseException, form_data: dict[str, Any], dataset_id: int
+) -> Literal["dataset", "annotation"] | None:
+ """Classify *exc* as a database connection failure, or return None.
+
+ Returns whether the failure came from the chart's own dataset or from an
+ annotation layer's chart. Walks the ``__cause__`` chain for typed database
+ errors, then falls back to the engine spec's ``extract_errors`` regex
+ patterns, the same classification the Superset UI uses. An annotation
+ failure is matched against the annotation charts' engines. Unless a single
+ engine is certain to be the failing one, only specific connection patterns
+ count: an engine's generic fallback would turn any error into a connection
+ failure.
+ """
+ failure = _find_database_failure(exc)
+ if failure is None:
+ return None
+ source: Literal["dataset", "annotation"] = (
+ "annotation"
+ if failure.wrapped and _has_chart_annotation_layers(form_data)
+ else "dataset"
+ )
+ if failure.typed:
+ return source
+ if not isinstance(failure.error, Exception):
+ return None
+
+ try:
+ specs: list[Any] = []
+ exact = source != "annotation"
+ if source == "annotation":
+ specs, exact = _annotation_engine_specs(form_data)
+ if not specs:
+ from superset.daos.dataset import DatasetDAO
+
+ dataset = DatasetDAO.find_by_id(dataset_id)
+ if dataset and dataset.database:
+ specs = [dataset.database.db_engine_spec]
+ for spec in specs:
+ errors = spec.extract_errors(failure.error)
+ if any(
+ e.error_type in CONNECTION_ERROR_TYPES
+ and (exact or e.error_type != SupersetErrorType.GENERIC_DB_ENGINE_ERROR)
+ for e in errors
+ ):
+ return source
except Exception: # pylint: disable=broad-except
logger.debug(
- "Failed to classify error via engine spec for dataset %s: %s",
+ "Failed to classify %s error via candidate engine specs "
+ "(main dataset %s): %s",
+ source,
dataset_id,
exc,
)
- return False
+ return None
-def _build_database_error(message: str) -> ChartGenerationError:
+def _build_database_error(
+ message: str, *, annotation: bool = False
+) -> ChartGenerationError:
"""Wrap a database connection failure in the structured response envelope."""
return ChartGenerationError(
error_type="database_connection_error",
- message="Unable to connect to the database.",
+ message=(
+ "Unable to connect to the database of an annotation layer's chart."
+ if annotation
+ else "Unable to connect to the database."
+ ),
details=message or "",
suggestions=[
"Check that the database is online and reachable",
diff --git a/tests/unit_tests/mcp_service/chart/test_chart_helpers.py b/tests/unit_tests/mcp_service/chart/test_chart_helpers.py
index 64af1b546ed..7a3cdc805f3 100644
--- a/tests/unit_tests/mcp_service/chart/test_chart_helpers.py
+++ b/tests/unit_tests/mcp_service/chart/test_chart_helpers.py
@@ -2195,3 +2195,45 @@ def test_deck_path_metric_keeps_canonical_path_grouping(
sql = table.get_query_str(query_object.to_dict())
assert "GROUP BY path" in sql, sql
assert "COUNT(*)" in sql, sql
+
+
+def test_build_query_dicts_from_form_data_passes_enum_values_as_strings(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ """Enum members in form_data reach the query as plain strings.
+
+ Template rendering rejects non-primitive context values, so an enum a mapper
+ forgets to unwrap must not leak into queries. Caller-owned data is preserved.
+ """
+ from copy import deepcopy
+ from enum import Enum
+
+ from superset.constants import TimeGrain
+
+ class _FutureEnum(Enum):
+ VALUE = "future"
+
+ monkeypatch.setattr(
+ "superset.mcp_service.chart.chart_helpers.resolve_datasource_engine",
+ lambda datasource_id, datasource_type: "base",
+ )
+ form_data: dict[str, Any] = {
+ "viz_type": "table",
+ "query_mode": "aggregate",
+ "groupby": ["state"],
+ "metrics": ["count"],
+ "adhoc_filters": [],
+ "granularity_sqla": "ds",
+ "time_grain_sqla": TimeGrain.MONTH,
+ "future_option": [{"nested": _FutureEnum.VALUE}],
+ }
+
+ original_form_data = deepcopy(form_data)
+ queries = build_query_dicts_from_form_data(form_data, 1, "table")
+
+ assert form_data == original_form_data
+ grain = queries[0]["extras"]["time_grain_sqla"]
+ assert grain == "P1M"
+ assert type(grain) is str
+ assert form_data["time_grain_sqla"] is TimeGrain.MONTH
+ assert form_data["future_option"][0]["nested"] is _FutureEnum.VALUE
diff --git a/tests/unit_tests/mcp_service/chart/test_compile.py b/tests/unit_tests/mcp_service/chart/test_compile.py
index 7fe27a9a606..63dfaba9a25 100644
--- a/tests/unit_tests/mcp_service/chart/test_compile.py
+++ b/tests/unit_tests/mcp_service/chart/test_compile.py
@@ -29,6 +29,7 @@ from unittest.mock import Mock, patch
import numpy as np
import pandas as pd
import pytest
+from pytest_mock import MockerFixture
from superset.common.db_query_status import QueryStatus
from superset.mcp_service.chart.compile import (
@@ -1133,3 +1134,466 @@ def test_compile_chart_returns_structured_error_for_malformed_gantt_form_data(
assert result.error_obj.error_type == "compile_error"
assert "tooltip_columns must contain at most 50 entries" in (result.error or "")
mock_cmd_cls.assert_not_called()
+
+
+@pytest.mark.parametrize(
+ "message",
+ [
+ "Error in jinja expression in metric expression: "
+ "Unsafe template value for key time_grain: TimeGrain",
+ "Invalid post-processing operation",
+ ],
+)
+def test_compile_chart_validation_failure_is_not_connection_error(
+ mocker: MockerFixture, message: str
+) -> None:
+ """Typed template/processing failures must bypass generic DB classification."""
+ from superset.db_engine_specs.base import BaseEngineSpec
+ from superset.exceptions import (
+ QueryObjectValidationError,
+ QueryObjectValidationMessage,
+ )
+
+ context = Mock()
+ context.get_payload.return_value = {
+ "queries": [
+ {"error": QueryObjectValidationMessage(QueryObjectValidationError(message))}
+ ]
+ }
+ mocker.patch(
+ "superset.mcp_service.chart.chart_helpers.build_query_context_from_form_data",
+ return_value=context,
+ )
+ mocker.patch("superset.charts.data.form_data.set_query_context_form_data")
+ dataset = Mock()
+ dataset.database.db_engine_spec = BaseEngineSpec
+ mocker.patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=dataset)
+
+ result = _compile_chart({"metrics": ["count"]}, dataset_id=1)
+
+ assert not result.success
+ assert result.error_obj is not None
+ assert result.error_obj.error_type == "compile_error"
+ assert result.error_obj.error_code == "CHART_COMPILE_FAILED"
+ assert message in result.error_obj.details
+
+
+@pytest.mark.parametrize(
+ ("message", "validation_failure", "expected_error_type"),
+ [
+ (
+ "(psycopg2.OperationalError) connection to server failed: "
+ "Connection refused",
+ False,
+ "database_connection_error",
+ ),
+ (
+ "Unsafe template value for key time_grain: TimeGrain",
+ True,
+ "compile_error",
+ ),
+ ("Invalid post-processing operation", True, "compile_error"),
+ ],
+)
+def test_compile_chart_annotation_failure_preserves_error_classification(
+ mocker: MockerFixture,
+ message: str,
+ validation_failure: bool,
+ expected_error_type: str,
+) -> None:
+ """Annotation failures retain their classification through validation wraps."""
+ from superset.common.query_context_processor import QueryContextProcessor
+ from superset.db_engine_specs.base import BaseEngineSpec
+ from superset.exceptions import (
+ QueryObjectValidationError,
+ QueryObjectValidationMessage,
+ )
+
+ annotation_context = Mock(queries=[])
+ annotation_error = (
+ QueryObjectValidationMessage(QueryObjectValidationError(message))
+ if validation_failure
+ else message
+ )
+ annotation_context.get_payload.return_value = {
+ "queries": [{"error": annotation_error}]
+ }
+ chart = Mock()
+ chart.get_query_context.return_value = annotation_context
+ mocker.patch(
+ "superset.common.query_context_processor.ChartDAO.find_by_id",
+ return_value=chart,
+ )
+
+ def get_payload(**kwargs: Any) -> dict[str, Any]:
+ """Record the annotation error like the main query processor."""
+ try:
+ QueryContextProcessor.get_viz_annotation_data(
+ {"value": 2, "name": "annotation"}, force=False
+ )
+ except QueryObjectValidationError as exc:
+ return {"queries": [{"error": QueryObjectValidationMessage(exc)}]}
+ raise AssertionError("Expected the annotation query to fail")
+
+ context = Mock()
+ context.get_payload.side_effect = get_payload
+ mocker.patch(
+ "superset.mcp_service.chart.chart_helpers.build_query_context_from_form_data",
+ return_value=context,
+ )
+ mocker.patch("superset.charts.data.form_data.set_query_context_form_data")
+ dataset = Mock()
+ dataset.database.db_engine_spec = BaseEngineSpec
+ mocker.patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=dataset)
+
+ result = _compile_chart({"metrics": ["count"]}, dataset_id=1)
+
+ assert not result.success
+ assert result.error_obj is not None
+ assert result.error_obj.error_type == expected_error_type
+ assert message in result.error_obj.details
+
+
+_POSTGRES_CONNECTION_REFUSED = (
+ "(psycopg2.OperationalError) could not connect to server: Connection refused\n"
+ '\tIs the server running on host "db.example" (10.0.0.1) and accepting\n'
+ "\tTCP/IP connections on port 5432?"
+)
+
+
+def _annotation_compile_result(
+ mocker: MockerFixture,
+ annotation_failure: str | Exception,
+ *,
+ dataset_engine_spec: Any,
+ annotation_engine_spec: Any | None = None,
+ other_annotation_engine_spec: Any | None = None,
+ annotation_key: str = "annotation_layers",
+ other_annotation_key: str = "annotation_layers",
+) -> Any:
+ """Compile a chart whose first line annotation chart fails.
+
+ ``other_annotation_engine_spec`` adds a second, healthy annotation chart on
+ another engine, so the exception alone does not say which engine failed.
+
+ ``annotation_failure`` is either the annotation query's payload error text or
+ an exception raised while the annotation chart builds its payload.
+ """
+ from superset.common.query_context_processor import QueryContextProcessor
+ from superset.exceptions import (
+ QueryObjectValidationError,
+ QueryObjectValidationMessage,
+ )
+
+ annotation_context = Mock(queries=[])
+ if isinstance(annotation_failure, Exception):
+ annotation_context.get_payload.side_effect = annotation_failure
+ else:
+ annotation_context.get_payload.return_value = {
+ "queries": [{"error": annotation_failure}]
+ }
+ chart = Mock()
+ chart.get_query_context.return_value = annotation_context
+ if annotation_engine_spec is None:
+ chart.datasource = None
+ else:
+ chart.datasource.database.db_engine_spec = annotation_engine_spec
+ other_chart = Mock()
+ other_chart.datasource.database.db_engine_spec = other_annotation_engine_spec
+ mocker.patch(
+ "superset.daos.chart.ChartDAO.find_by_id",
+ side_effect=lambda chart_id: other_chart if chart_id == 3 else chart,
+ )
+ layer = {
+ "annotationType": "TIME_SERIES",
+ "name": "annotation",
+ "sourceType": "line",
+ "value": 2,
+ }
+ form_data: dict[str, Any] = {"metrics": ["count"], annotation_key: [layer]}
+ if annotation_key == "annotation_layers_b":
+ form_data["viz_type"] = "mixed_timeseries"
+ if other_annotation_engine_spec is not None:
+ form_data.setdefault(other_annotation_key, []).append(
+ {**layer, "name": "other annotation", "value": 3}
+ )
+
+ def get_payload(**kwargs: Any) -> dict[str, Any]:
+ """Record the annotation error like the main query processor."""
+ try:
+ QueryContextProcessor.get_viz_annotation_data(layer, force=False)
+ except QueryObjectValidationError as exc:
+ return {"queries": [{"error": QueryObjectValidationMessage(exc)}]}
+ raise AssertionError("Expected the annotation query to fail")
+
+ context = Mock()
+ context.get_payload.side_effect = get_payload
+ mocker.patch(
+ "superset.mcp_service.chart.chart_helpers.build_query_context_from_form_data",
+ return_value=context,
+ )
+ mocker.patch("superset.charts.data.form_data.set_query_context_form_data")
+ dataset = Mock()
+ dataset.database.db_engine_spec = dataset_engine_spec
+ mocker.patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=dataset)
+
+ return _compile_chart(form_data, dataset_id=1)
+
+
+@pytest.mark.parametrize(
+ ("message", "expected_error_type"),
+ [
+ (_POSTGRES_CONNECTION_REFUSED, "database_connection_error"),
+ (
+ '(psycopg2.errors.SyntaxError) syntax error at or near "FROMM"',
+ "compile_error",
+ ),
+ ],
+)
+def test_compile_chart_annotation_failure_uses_engine_error_patterns(
+ mocker: MockerFixture, message: str, expected_error_type: str
+) -> None:
+ """Only the engine's connection patterns make an annotation failure a DB error.
+
+ PostgresEngineSpec recognizes both texts, so a syntax error in the
+ annotation query stays a compile error instead of the generic fallback.
+ """
+ from superset.db_engine_specs.postgres import PostgresEngineSpec
+
+ result = _annotation_compile_result(
+ mocker,
+ message,
+ dataset_engine_spec=PostgresEngineSpec,
+ annotation_engine_spec=PostgresEngineSpec,
+ )
+
+ assert not result.success
+ assert result.error_obj is not None
+ assert result.error_obj.error_type == expected_error_type
+ assert message in result.error_obj.details
+
+
+def test_compile_chart_annotation_dbapi_error_is_connection_error(
+ mocker: MockerFixture,
+) -> None:
+ """A typed DB-API failure raised outside the annotation query stays a DB error.
+
+ The annotation chart raises it before any payload exists, so the chain holds
+ no chart-data query failure, only the typed engine error.
+ """
+ from superset.db_engine_specs.exceptions import SupersetDBAPIOperationalError
+ from superset.db_engine_specs.postgres import PostgresEngineSpec
+
+ result = _annotation_compile_result(
+ mocker,
+ SupersetDBAPIOperationalError(_POSTGRES_CONNECTION_REFUSED),
+ dataset_engine_spec=PostgresEngineSpec,
+ annotation_engine_spec=PostgresEngineSpec,
+ )
+
+ assert not result.success
+ assert result.error_obj is not None
+ assert result.error_obj.error_type == "database_connection_error"
+
+
+@pytest.mark.parametrize(
+ ("message", "expected_error_type"),
+ [
+ (
+ "(MySQLdb.OperationalError) (2003, \"Can't connect to MySQL server "
+ "on 'mysql.example' (111)\")",
+ "database_connection_error",
+ ),
+ (
+ "(MySQLdb.ProgrammingError) (1064, 'You have an error in your SQL "
+ "syntax; check the manual that corresponds to your MySQL server "
+ "version for the right syntax to use near 'FROMM x' at line 1')",
+ "compile_error",
+ ),
+ ],
+)
+def test_compile_chart_annotation_failure_uses_annotation_engine_spec(
+ mocker: MockerFixture, message: str, expected_error_type: str
+) -> None:
+ """An annotation chart's failure is classified by its own database's engine.
+
+ The main dataset is on Postgres, whose patterns match neither MySQL text and
+ so fall back to the generic engine error.
+ """
+ from superset.db_engine_specs.mysql import MySQLEngineSpec
+ from superset.db_engine_specs.postgres import PostgresEngineSpec
+
+ result = _annotation_compile_result(
+ mocker,
+ message,
+ dataset_engine_spec=PostgresEngineSpec,
+ annotation_engine_spec=MySQLEngineSpec,
+ )
+
+ assert not result.success
+ assert result.error_obj is not None
+ assert result.error_obj.error_type == expected_error_type
+ if expected_error_type == "database_connection_error":
+ assert "annotation layer" in result.error_obj.message
+
+
+def test_classify_as_database_error_stops_on_cyclic_cause_chain() -> None:
+ """A cyclic ``__cause__`` chain ends the walk instead of looping forever."""
+ from superset.commands.chart.exceptions import ChartDataQueryValidationError
+ from superset.exceptions import QueryObjectValidationError
+ from superset.mcp_service.chart.compile import _classify_as_database_error
+
+ outer = ChartDataQueryValidationError("Error: annotation failed")
+ cause_reads = 0
+
+ class BoundedValidationError(QueryObjectValidationError):
+ """Fail synchronously if the cycle guard stops bounding the walk."""
+
+ def __getattribute__(self, name: str) -> Any:
+ nonlocal cause_reads
+ if name == "__cause__":
+ cause_reads += 1
+ assert cause_reads <= 2, "cause-chain walk revisited the cycle"
+ return super().__getattribute__(name)
+
+ inner = BoundedValidationError("annotation failed")
+ outer.__cause__ = inner
+ inner.__cause__ = outer
+ assert _classify_as_database_error(outer, {}, 1) is None
+ assert cause_reads == 1
+
+
+def test_compile_chart_annotation_syntax_error_with_several_engines_is_not_connection(
+ mocker: MockerFixture,
+) -> None:
+ """Which annotation engine failed is unknown, so no generic fallback applies.
+
+ Two annotation charts sit on Postgres and MySQL. A MySQL syntax error matches
+ no specific connection pattern; an engine's generic fallback must not turn it
+ into a database connection error.
+ """
+ from superset.db_engine_specs.mysql import MySQLEngineSpec
+ from superset.db_engine_specs.postgres import PostgresEngineSpec
+
+ message = (
+ "(MySQLdb.ProgrammingError) (1064, 'You have an error in your SQL "
+ "syntax; check the manual that corresponds to your MySQL server "
+ "version for the right syntax to use near 'FROMM x' at line 1')"
+ )
+ result = _annotation_compile_result(
+ mocker,
+ message,
+ dataset_engine_spec=PostgresEngineSpec,
+ annotation_engine_spec=MySQLEngineSpec,
+ other_annotation_engine_spec=PostgresEngineSpec,
+ )
+
+ assert not result.success
+ assert result.error_obj is not None
+ assert result.error_obj.error_type == "compile_error"
+ assert message in result.error_obj.details
+
+
+def test_compile_chart_annotation_connection_error_with_several_engines(
+ mocker: MockerFixture,
+) -> None:
+ """A specific connection pattern still classifies with several engines."""
+ from superset.db_engine_specs.mysql import MySQLEngineSpec
+ from superset.db_engine_specs.postgres import PostgresEngineSpec
+
+ result = _annotation_compile_result(
+ mocker,
+ _POSTGRES_CONNECTION_REFUSED,
+ dataset_engine_spec=MySQLEngineSpec,
+ annotation_engine_spec=PostgresEngineSpec,
+ other_annotation_engine_spec=MySQLEngineSpec,
+ )
+
+ assert not result.success
+ assert result.error_obj is not None
+ assert result.error_obj.error_type == "database_connection_error"
+ assert "annotation layer's chart" in result.error_obj.message
+
+
+@pytest.mark.parametrize("with_query_a_layer", [False, True])
+@pytest.mark.parametrize(
+ ("message", "expected_error_type"),
+ [
+ (
+ "(MySQLdb.OperationalError) (2003, \"Can't connect to MySQL server "
+ "on 'mysql.example' (111)\")",
+ "database_connection_error",
+ ),
+ (
+ "(MySQLdb.ProgrammingError) (1064, 'You have an error in your SQL "
+ "syntax; check the manual that corresponds to your MySQL server "
+ "version for the right syntax to use near 'FROMM x' at line 1')",
+ "compile_error",
+ ),
+ ],
+)
+def test_compile_chart_mixed_timeseries_query_b_annotation_failure(
+ mocker: MockerFixture,
+ with_query_a_layer: bool,
+ message: str,
+ expected_error_type: str,
+) -> None:
+ """Query B contributes both the annotation source and possible failing engines."""
+ from superset.db_engine_specs.mysql import MySQLEngineSpec
+ from superset.db_engine_specs.postgres import PostgresEngineSpec
+
+ result = _annotation_compile_result(
+ mocker,
+ message,
+ dataset_engine_spec=PostgresEngineSpec,
+ annotation_engine_spec=MySQLEngineSpec,
+ other_annotation_engine_spec=(
+ PostgresEngineSpec if with_query_a_layer else None
+ ),
+ annotation_key="annotation_layers_b",
+ )
+
+ assert not result.success
+ assert result.error_obj is not None
+ assert result.error_obj.error_type == expected_error_type
+ assert message in result.error_obj.details
+ if expected_error_type == "database_connection_error":
+ assert "annotation layer's chart" in result.error_obj.message
+
+
+@pytest.mark.parametrize("source", ["dataset", "annotation"])
+def test_classification_debug_log_identifies_candidate_engine_source(
+ mocker: MockerFixture, source: str
+) -> None:
+ """Engine lookup failures identify the source, not an annotation dataset."""
+ from superset.commands.chart.exceptions import (
+ ChartDataQueryFailedError,
+ ChartDataQueryValidationError,
+ )
+ from superset.mcp_service.chart.compile import _classify_as_database_error
+
+ failure = ChartDataQueryFailedError("query failed")
+ exc: Exception = failure
+ form_data: dict[str, Any] = {}
+ spec = Mock()
+ spec.extract_errors.side_effect = RuntimeError("classification failed")
+ dataset = Mock()
+ dataset.database.db_engine_spec = spec
+ mocker.patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=dataset)
+ mocker.patch(
+ "superset.mcp_service.chart.compile._annotation_engine_specs",
+ return_value=([spec], True),
+ )
+ if source == "annotation":
+ exc = ChartDataQueryValidationError("annotation failed")
+ exc.__cause__ = failure
+ form_data["annotation_layers"] = [{"sourceType": "line", "value": 2}]
+ debug = mocker.patch("superset.mcp_service.chart.compile.logger.debug")
+
+ assert _classify_as_database_error(exc, form_data, 1) is None
+ debug.assert_called_once_with(
+ "Failed to classify %s error via candidate engine specs (main dataset %s): %s",
+ source,
+ 1,
+ exc,
+ )
diff --git a/tests/unit_tests/mcp_service/chart/test_review_query_and_preview_regressions.py b/tests/unit_tests/mcp_service/chart/test_review_query_and_preview_regressions.py
index 8a6435b82f6..b05a247deaa 100644
--- a/tests/unit_tests/mcp_service/chart/test_review_query_and_preview_regressions.py
+++ b/tests/unit_tests/mcp_service/chart/test_review_query_and_preview_regressions.py
@@ -15,6 +15,7 @@
# specific language governing permissions and limitations
# under the License.
+from copy import deepcopy
from datetime import datetime
from types import SimpleNamespace
from typing import Any
@@ -81,6 +82,7 @@ def test_table_inherited_offset_survives_filter_preparation(
}
if not saved_viz:
form_data.pop("viz_type")
+ original_form_data = deepcopy(form_data)
with patch(
"superset.mcp_service.chart.chart_helpers.resolve_datasource_engine",
return_value="base",
@@ -95,7 +97,7 @@ def test_table_inherited_offset_survives_filter_preparation(
else None,
)[0]
offset = request_offset or "1 year ago"
- assert "extra_form_data" not in form_data
+ assert form_data == original_form_data
assert query["time_offsets"] == [offset]
assert query["post_processing"][0]["options"]["compare_columns"] == [
f"revenue__{offset}"
@@ -349,6 +351,7 @@ def test_table_inherited_custom_offset_preserves_other_comparisons(
}
if not saved_viz:
form_data.pop("viz_type")
+ original_form_data = deepcopy(form_data)
with patch(
"superset.mcp_service.chart.chart_helpers.resolve_datasource_engine",
return_value="base",
@@ -367,7 +370,7 @@ def test_table_inherited_custom_offset_preserves_other_comparisons(
if request_offset == "3 weeks ago"
else ["1 year ago", "2 weeks ago"]
)
- assert "extra_form_data" not in form_data
+ assert form_data == original_form_data
assert query["time_offsets"] == offsets
assert query["post_processing"][0]["options"]["compare_columns"] == [
f"revenue__{offset}" for offset in offsets
diff --git a/tests/unit_tests/mcp_service/chart/test_time_grain_generic_axes.py b/tests/unit_tests/mcp_service/chart/test_time_grain_generic_axes.py
index f342ebc8ab7..736807a5b12 100644
--- a/tests/unit_tests/mcp_service/chart/test_time_grain_generic_axes.py
+++ b/tests/unit_tests/mcp_service/chart/test_time_grain_generic_axes.py
@@ -363,3 +363,174 @@ async def test_saved_chart_grain_acceptance(
sql_counts.append(len(rows))
assert sum(row.ct for row in rows) == 25
assert sql_counts == [expected_rows] * query_count
+
+
+@pytest.mark.parametrize(
+ "app", [{"FEATURE_FLAGS": {"ENABLE_TEMPLATE_PROCESSING": True}}], indirect=True
+)
+def test_big_number_time_grain_renders_plain_saved_metric(
+ mocker: MockerFixture,
+) -> None:
+ """MCP time grains must remain safe even when a metric contains no Jinja."""
+ from superset.common.form_data_query_context import (
+ build_query_objects_from_form_data,
+ )
+ from superset.connectors.sqla.models import SqlMetric
+ from superset.mcp_service.chart.chart_utils import map_config_to_form_data
+ from superset.mcp_service.chart.schemas import BigNumberChartConfig
+
+ table = _events_dataset(mocker)
+ table.metrics = [SqlMetric(metric_name="total_value", expression="SUM(val)")]
+ config = BigNumberChartConfig.model_validate(
+ {
+ "chart_type": "big_number",
+ "metric": {"name": "total_value", "saved_metric": True},
+ "show_trendline": True,
+ "temporal_column": "ts",
+ "time_grain": "P1M",
+ "aggregation": "sum",
+ }
+ )
+ form_data = map_config_to_form_data(config)
+ query = build_query_objects_from_form_data(form_data)[0]
+ sql_query = table.get_sqla_query(
+ metrics=query["metrics"],
+ extras=query["extras"],
+ granularity=query["granularity"],
+ is_timeseries=True,
+ )
+ assert "SUM(val)" in str(sql_query.sqla_query)
+ assert type(form_data["time_grain_sqla"]) is str
+
+
+@pytest.mark.parametrize(
+ "config_data",
+ [
+ {
+ "chart_type": "xy",
+ "x": {"name": "ts"},
+ "y": [{"name": "val", "aggregate": "SUM"}],
+ },
+ {
+ "chart_type": "mixed_timeseries",
+ "x": {"name": "ts"},
+ "y": [{"name": "val", "aggregate": "SUM"}],
+ "y_secondary": [{"name": "val", "aggregate": "AVG"}],
+ "time_grain_sqla_b": "P1Y",
+ },
+ {
+ "chart_type": "sunburst",
+ "hierarchy": [{"name": "ts"}],
+ "metric": {"name": "val", "aggregate": "SUM"},
+ },
+ {
+ "chart_type": "waterfall",
+ "x_axis": {"name": "ts"},
+ "metric": {"name": "val", "aggregate": "SUM"},
+ },
+ {
+ "chart_type": "interactive_pivot",
+ "rows": [{"name": "ts"}],
+ "metrics": [{"name": "val", "aggregate": "SUM"}],
+ },
+ ],
+ ids=["xy", "mixed_timeseries", "sunburst", "waterfall", "interactive_pivot"],
+)
+def test_chart_time_grain_form_data_contains_strings(
+ mocker: MockerFixture, config_data: dict[str, Any]
+) -> None:
+ """Every typed time-grain mapper must emit primitive form-data values."""
+ from pydantic import TypeAdapter
+
+ from superset.jinja_context import validate_context_types
+ from superset.mcp_service.chart.chart_utils import (
+ map_config_to_form_data,
+ merge_form_data_for_update,
+ )
+ from superset.mcp_service.chart.schemas import ChartConfig
+
+ mocker.patch(
+ "superset.mcp_service.chart.chart_utils.is_column_truly_temporal",
+ return_value=True,
+ )
+ config = TypeAdapter(ChartConfig).validate_python(
+ {**config_data, "time_grain": "P1M", "temporal_column": "ts"}
+ )
+ form_data = map_config_to_form_data(config, include_disabled=True)
+ if config.chart_type == "sunburst":
+ form_data = merge_form_data_for_update(form_data, form_data.copy(), config)
+ # Equality alone cannot distinguish a str enum from its value.
+ assert type(form_data["time_grain_sqla"]) is str
+ assert validate_context_types({"time_grain": form_data["time_grain_sqla"]}) == {
+ "time_grain": "P1M"
+ }
+ if config.chart_type == "mixed_timeseries":
+ assert type(form_data["time_grain_sqla_b"]) is str
+ assert form_data["time_grain_sqla_b"] == "P1Y"
+
+
+@pytest.mark.parametrize(
+ "app", [{"FEATURE_FLAGS": {"ENABLE_TEMPLATE_PROCESSING": True}}], indirect=True
+)
+@pytest.mark.parametrize(
+ "config_data",
+ [
+ {
+ "chart_type": "big_number",
+ "metric": {"name": "val", "aggregate": "SUM"},
+ "show_trendline": True,
+ "temporal_column": "ts",
+ "time_grain": "P1M",
+ "aggregation": "sum",
+ },
+ {
+ "chart_type": "xy",
+ "x": {"name": "ts"},
+ "y": [{"name": "val", "aggregate": "SUM"}],
+ "kind": "line",
+ "time_grain": "P1M",
+ },
+ ],
+ ids=["big_number", "xy"],
+)
+def test_time_grain_renders_on_virtual_dataset(
+ mocker: MockerFixture, config_data: dict[str, Any]
+) -> None:
+ """A time grain must reach a virtual dataset's Jinja context as a string.
+
+ A virtual dataset (``sql`` set) renders its SQL through the template
+ processor, which rejects any context value whose type is not a plain
+ builtin. An enum-valued ``time_grain`` fails there with ``Unsafe template
+ value for key time_grain: TimeGrain``.
+ """
+ from pydantic import TypeAdapter
+
+ from superset.common.form_data_query_context import (
+ build_query_objects_from_form_data,
+ )
+ from superset.mcp_service.chart.chart_utils import map_config_to_form_data
+ from superset.mcp_service.chart.schemas import ChartConfig
+
+ mocker.patch(
+ "superset.mcp_service.chart.chart_utils.is_column_truly_temporal",
+ return_value=True,
+ )
+ table = _events_dataset(mocker)
+ table.table_name = "virtual_events"
+ table.sql = "SELECT ts, val FROM events"
+
+ config = TypeAdapter(ChartConfig).validate_python(config_data)
+ form_data = map_config_to_form_data(config)
+ query = build_query_objects_from_form_data(form_data)[0]
+
+ sql_query = table.get_sqla_query(
+ metrics=query["metrics"],
+ columns=query.get("columns"),
+ extras=query["extras"],
+ granularity=query.get("granularity"),
+ is_timeseries=True,
+ )
+
+ assert "SELECT ts, val FROM events" in str(sql_query.sqla_query)
+ assert type(query["extras"]["time_grain_sqla"]) is str
+ assert query["extras"]["time_grain_sqla"] == "P1M"