bito-code-review[bot] commented on code in PR #43770:
URL: https://github.com/apache/superset/pull/43770#discussion_r3924022507
##########
tests/unit_tests/mcp_service/chart/tool/test_get_chart_preview.py:
##########
@@ -45,12 +53,1126 @@
_first_query_has_fields,
_no_query_fields_error,
ASCIIPreviewStrategy,
+ get_chart_preview,
PreviewFormatStrategy,
TablePreviewStrategy,
)
from superset.utils import json as utils_json
+def _query_context_stub(form_data: dict[str, Any] | None = None) -> Any:
+ """Return the minimal real-shaped context needed by Jinja form-data
seeding."""
+ return SimpleNamespace(form_data=form_data or {}, queries=[])
+
+
+def _entrypoint_preview(content: str) -> ChartPreview:
+ return ChartPreview(
+ chart_id=1,
+ chart_name="",
+ chart_type="bullet",
+ explore_url="",
+ content=ASCIIPreview(ascii_content=content, width=80, height=20),
+ chart_description="",
+ accessibility=AccessibilityMetadata(
+ color_blind_safe=True,
+ alt_text="",
+ high_contrast_available=False,
+ ),
+ performance=PerformanceMetadata(
+ query_duration_ms=0,
+ cache_status="miss",
+ optimization_suggestions=[],
+ ),
+ )
+
+
[email protected]
[email protected](
+ ("request_payload", "extra", "expected_type"),
+ [
+ ({"id": 1, "format": "ascii"}, 0, ChartPreview),
+ ({"form_data_key": "cached-preview", "format": "ascii"}, 1,
ChartError),
+ ],
+ ids=["identifier-alias-exact-limit", "cached-preview-limit-plus-one"],
+)
+async def
test_get_chart_preview_entrypoint_preflights_complete_exact_wire_response(
+ request_payload: dict[str, object], extra: int, expected_type: type[object]
+) -> None:
+ empty = _entrypoint_preview("")
+ filler = "x" * (
+ MAX_QUERY_RESULT_VALUE_BYTES - len(empty.model_dump_json().encode()) +
extra
+ )
+ candidate = _entrypoint_preview(filler)
+ request = GetChartPreviewRequest.model_validate(request_payload)
+ ctx = MagicMock()
+ ctx.info = AsyncMock()
+ ctx.debug = AsyncMock()
+ ctx.warning = AsyncMock()
+
+ user = MagicMock(id=1, username="admin", roles=[], groups=[])
+ with (
+ patch("superset.mcp_service.auth.get_user_from_request",
return_value=user),
+ patch(
+ "superset.mcp_service.chart.tool.get_chart_preview."
+ "_get_chart_preview_internal",
+ new=AsyncMock(return_value=candidate),
+ ),
+ ):
+ result = await get_chart_preview(request, ctx=ctx)
+
+ assert isinstance(result, expected_type)
+ if extra == 0:
+ assert len(candidate.model_dump_json().encode()) == (
+ MAX_QUERY_RESULT_VALUE_BYTES
+ )
+ assert result is candidate
+ else:
+ assert isinstance(result, ChartError)
+ assert result.error_type == "MalformedQueryResult"
+
+
[email protected]
[email protected]("format_", ["ascii", "vega_lite"])
+async def
test_bullet_numeric_and_temporal_categories_reach_real_mcp_entrypoint(
+ format_: str,
+) -> None:
+ from contextlib import nullcontext
+
+ from fastmcp import Client
+
+ from superset.mcp_service.app import mcp
+
+ preview_module = importlib.import_module(
+ "superset.mcp_service.chart.tool.get_chart_preview"
+ )
+ command_module = importlib.import_module(
+ "superset.commands.chart.data.get_data_command"
+ )
+ form_data = {
+ "viz_type": "bullet",
+ "metric": "Revenue",
+ "groupby": ["Category"],
+ }
+ chart = SimpleNamespace(
+ id=121,
+ slice_name="Number boundaries",
+ viz_type="bullet",
+ datasource_id=1,
+ datasource_type="table",
+ params=utils_json.dumps(form_data),
+ )
+ rows = [
+ {"Category": 9007199254740993, "Revenue": 1},
+ {"Category": Decimal("1.0000000000000001"), "Revenue": 2},
+ {"Category": Decimal("1.7976931348623159e308"), "Revenue": 3},
+ {"Category": date(2026, 9, 2), "Revenue": 4},
+ {
+ "Category": datetime(2026, 9, 2, 3, 4, 5, tzinfo=timezone.utc),
+ "Revenue": 5,
+ },
+ {
+ "Category": datetime(
+ 2023,
+ 11,
+ 5,
+ 1,
+ 30,
+ tzinfo=ZoneInfo("America/New_York"),
+ fold=0,
+ ),
+ "Revenue": 6,
+ },
+ {
+ "Category": datetime(
+ 2023,
+ 11,
+ 5,
+ 1,
+ 30,
+ tzinfo=ZoneInfo("America/New_York"),
+ fold=1,
+ ),
+ "Revenue": 7,
+ },
+ ]
+
+ class _Command:
+ def __init__(self, _query_context: Any) -> None: ...
+
+ def validate(self) -> None: ...
+
+ def run(self) -> dict[str, Any]:
+ return {
+ "queries": [
+ {
+ "data": rows,
+ "colnames": ["Category", "Revenue"],
+ }
+ ]
+ }
+
+ query_context = SimpleNamespace(
+ form_data={},
+ queries=[SimpleNamespace(metrics=["Revenue"], columns=["Category"])],
+ )
+ user = MagicMock(id=1, username="admin", roles=[], groups=[])
+ with (
+ patch("superset.mcp_service.auth.get_user_from_request",
return_value=user),
+ patch("superset.mcp_service.auth.check_tool_permission",
return_value=True),
+ patch.object(preview_module, "find_chart_by_identifier",
return_value=chart),
+ patch.object(preview_module.db.session, "refresh", return_value=None),
+ patch.object(
+ preview_module,
+ "validate_chart_dataset",
+ return_value=SimpleNamespace(is_valid=True, warnings=[],
error=None),
+ ),
+ patch.object(
+ preview_module.event_logger,
+ "log_context",
+ side_effect=lambda **_kwargs: nullcontext(),
+ ),
+ patch.object(
+ preview_module,
+ "build_query_context_from_form_data",
+ return_value=query_context,
+ ),
+ patch.object(preview_module, "set_query_context_form_data",
return_value=None),
+ patch.object(command_module, "ChartDataCommand", _Command),
+ patch.object(
+ preview_module, "get_superset_base_url",
return_value="http://localhost";
+ ),
+ ):
+ async with Client(mcp) as client:
+ result = await client.call_tool(
+ "get_chart_preview",
+ {"request": {"id": 121, "format": format_}},
+ )
+
+ payload = utils_json.loads(result.content[0].text)
+ if format_ == "ascii":
+ content = payload["content"]["ascii_content"]
+ assert "9007199254740992" in content
+ assert "Infinity" in content
+ assert "1788307200000" in content
+ assert "1699162200000" in content
+ assert "1699165800000" in content
+ else:
+ specification = payload["content"]["specification"]
+ bar = next(
+ layer for layer in specification["layer"] if layer["mark"]["type"]
== "bar"
+ )
+ category_field = bar["encoding"]["y"]["field"]
+ assert [row[category_field] for row in
specification["data"]["values"]] == [
+ "9007199254740992",
+ "1",
+ "Infinity",
+ "1788307200000",
+ "1788318245000",
+ "1699162200000",
+ "1699165800000",
+ ]
+ assert [row["Category"] for row in
specification["data"]["values"][3:]] == [
+ 1788307200000.0,
+ 1788318245000.0,
+ 1699162200000.0,
+ 1699165800000.0,
+ ]
+ assert bar["encoding"]["tooltip"][0]["field"] == category_field
+ assert "transform" not in specification
+
+
[email protected]
[email protected]("format_", ["ascii", "vega_lite"])
+async def test_bullet_timestamp_categories_from_dataframe_reach_fastmcp(
+ format_: str,
+) -> None:
+ from contextlib import nullcontext
+
+ from fastmcp import Client
+
+ from superset.dataframe import df_to_records
+ from superset.mcp_service.app import mcp
+
+ preview_module = importlib.import_module(
+ "superset.mcp_service.chart.tool.get_chart_preview"
+ )
+ command_module = importlib.import_module(
+ "superset.commands.chart.data.get_data_command"
+ )
+ zoneinfo_tz = ZoneInfo("America/New_York")
+ pytz_tz = pytz.timezone("America/New_York")
+ dateutil_dublin = dateutil_tz.gettz("Europe/Dublin")
+ dateutil_new_york = dateutil_tz.gettz("America/New_York")
+ assert dateutil_dublin is not None
+ assert dateutil_new_york is not None
+ dublin_fold = datetime(
+ 2024,
+ 10,
+ 27,
+ 1,
+ 30,
+ 0,
+ 123456,
+ tzinfo=dateutil_dublin,
+ fold=1,
+ )
+ new_york_gap = datetime(2024, 3, 10, 2, 30, 0, 123456,
tzinfo=dateutil_new_york)
+ source_values = [
+ pd.Timestamp("2024-01-02 03:04:05.123456789"),
+ pd.Timestamp("2024-01-02 08:34:05.123456789+05:30"),
+ pd.Timestamp(datetime(2024, 11, 3, 1, 30, tzinfo=zoneinfo_tz, fold=0)),
+ pd.Timestamp(datetime(2024, 11, 3, 1, 30, tzinfo=zoneinfo_tz, fold=1)),
+ pd.Timestamp(pytz_tz.localize(datetime(2024, 11, 3, 1, 30),
is_dst=True)),
+ pd.Timestamp(pytz_tz.localize(datetime(2024, 11, 3, 1, 30),
is_dst=False)),
+ pd.Timestamp("1969-12-31 23:59:59.999999999"),
+ date(2024, 1, 2),
+ pd.NaT,
+ dublin_fold,
+ new_york_gap,
+ datetime(2040, 7, 1, 12, 0, 0, 123456, tzinfo=dateutil_new_york),
+ datetime(
+ 2024,
+ 3,
+ 10,
+ 2,
+ 30,
+ 0,
+ 123456,
+ tzinfo=dateutil_tz.tzoffset("EDT", -4 * 3600),
+ ),
+ datetime(2024, 3, 10, 6, 30, 0, 123456, tzinfo=timezone.utc),
+ pd.Timestamp(dublin_fold),
+ pd.Timestamp(new_york_gap),
+ ]
+ rows = df_to_records(
+ pd.DataFrame(
+ {
+ "Category": pd.Series(source_values, dtype=object),
+ "Revenue": range(1, len(source_values) + 1),
+ }
+ ),
+ convert_big_integers=False,
+ )
+ if format_ == "ascii":
+ # ASCII intentionally displays at most ten categories. Keep every
+ # dateutil named-zone edge in this public-format invocation.
+ rows = rows[9:12]
+ expected = [
+ "1704164645123.456",
+ "1704164645123.456",
+ "1730611800000",
+ "1730615400000",
+ "1730611800000",
+ "1730615400000",
+ "-0.0010000000000287557",
+ "1704153600000",
+ "null",
+ "1729989000123.456",
+ "1710052200123.456",
+ "2224774800123.456",
+ "1710052200123.456",
+ "1710052200123.456",
+ "1729989000123.456",
+ "1710052200123.456",
+ ]
+ form_data = {
+ "viz_type": "bullet",
+ "metric": "Revenue",
+ "groupby": ["Category"],
+ }
+ chart = SimpleNamespace(
+ id=122,
+ slice_name="Timestamp categories",
+ viz_type="bullet",
+ datasource_id=1,
+ datasource_type="table",
+ params=utils_json.dumps(form_data),
+ )
+
+ class _Command:
+ def __init__(self, _query_context: Any) -> None: ...
+
+ def validate(self) -> None: ...
+
+ def run(self) -> dict[str, Any]:
+ return {
+ "queries": [
+ {
+ "data": rows,
+ "colnames": ["Category", "Revenue"],
+ }
+ ]
+ }
+
+ query_context = SimpleNamespace(
+ form_data={},
+ queries=[SimpleNamespace(metrics=["Revenue"], columns=["Category"])],
+ )
+ user = MagicMock(id=1, username="admin", roles=[], groups=[])
+ with (
+ patch("superset.mcp_service.auth.get_user_from_request",
return_value=user),
+ patch("superset.mcp_service.auth.check_tool_permission",
return_value=True),
+ patch.object(preview_module, "find_chart_by_identifier",
return_value=chart),
+ patch.object(preview_module.db.session, "refresh", return_value=None),
+ patch.object(
+ preview_module,
+ "validate_chart_dataset",
+ return_value=SimpleNamespace(is_valid=True, warnings=[],
error=None),
+ ),
+ patch.object(
+ preview_module.event_logger,
+ "log_context",
+ side_effect=lambda **_kwargs: nullcontext(),
+ ),
+ patch.object(
+ preview_module,
+ "build_query_context_from_form_data",
+ return_value=query_context,
+ ),
+ patch.object(preview_module, "set_query_context_form_data",
return_value=None),
+ patch.object(command_module, "ChartDataCommand", _Command),
+ patch.object(
+ preview_module, "get_superset_base_url",
return_value="http://localhost";
+ ),
+ ):
+ async with Client(mcp) as client:
+ result = await client.call_tool(
+ "get_chart_preview",
+ {"request": {"id": 122, "format": format_}},
+ )
+
+ payload = utils_json.loads(result.content[0].text)
+ if format_ == "ascii":
+ content = payload["content"]["ascii_content"]
+ for category in set(expected[9:12]):
+ assert category[:20] in content
+ else:
+ specification = payload["content"]["specification"]
+ bar = next(
+ layer for layer in specification["layer"] if layer["mark"]["type"]
== "bar"
+ )
+ category_field = bar["encoding"]["y"]["field"]
+ assert [row[category_field] for row in
specification["data"]["values"]] == (
+ expected
+ )
+ assert bar["encoding"]["tooltip"][0]["field"] == category_field
+ assert [row["Category"] for row in specification["data"]["values"]] ==
[
+ 1704164645123.456,
+ 1704164645123.456,
+ 1730611800000.0,
+ 1730615400000.0,
+ 1730611800000.0,
+ 1730615400000.0,
+ -0.0010000000000287557,
+ 1704153600000.0,
+ None,
+ 1729989000123.456,
+ 1710052200123.456,
+ 2224774800123.456,
+ 1710052200123.456,
+ 1710052200123.456,
+ 1729989000123.456,
+ 1710052200123.456,
+ ]
+
+
[email protected]
+async def test_transitionless_dateutil_dataframe_reaches_fastmcp_preview() ->
None:
+ from contextlib import nullcontext
+
+ from fastmcp import Client
+
+ from superset.commands.chart.data.get_data_command import (
+ ChartDataCommand as ProducerChartDataCommand,
+ )
+ from superset.common.chart_data import ChartDataResultType
+ from superset.dataframe import df_to_records
+ from superset.mcp_service.app import mcp
+ from superset.mcp_service.chart.preview_utils import
_javascript_number_string
+ from superset.utils.json import json_int_dttm_ser
+
+ preview_module = importlib.import_module(
+ "superset.mcp_service.chart.tool.get_chart_preview"
+ )
+ command_module = importlib.import_module(
+ "superset.commands.chart.data.get_data_command"
+ )
+ names = [
+ "UTC",
+ "GMT",
+ "Universal",
+ "Zulu",
+ "EST",
+ "HST",
+ "MST",
+ "Etc/GMT+1",
+ "Etc/GMT-2",
+ ]
+ values = []
+ for getter in (dateutil_tz.gettz, get_zonefile_instance().get):
+ for name in names:
+ timezone_value = getter(name)
+ assert timezone_value is not None
+ values.append(
+ datetime(2040, 7, 1, 12, 34, 56, 123456, tzinfo=timezone_value)
+ )
+ rows = df_to_records(
+ pd.DataFrame(
+ {
+ "Category": pd.Series(values, dtype=object),
+ "Revenue": range(1, len(values) + 1),
+ }
+ ),
+ convert_big_integers=False,
+ )
+
+ class _ProducerContext:
+ result_type = ChartDataResultType.FULL
+
+ def get_payload(self, **_kwargs: Any) -> dict[str, Any]:
+ return {
+ "queries": [
+ {
+ "data": rows,
+ "colnames": ["Category", "Revenue"],
+ "rowcount": len(rows),
+ }
+ ]
+ }
+
+ producer_result = ProducerChartDataCommand(
+ _ProducerContext() # type: ignore[arg-type]
+ ).run()
+ expected_numbers = [json_int_dttm_ser(value) for value in values]
+ expected_categories = [
+ _javascript_number_string(float(value)) for value in expected_numbers
+ ]
+ form_data = {
+ "viz_type": "bullet",
+ "metric": "Revenue",
+ "groupby": ["Category"],
+ }
+ chart = SimpleNamespace(
+ id=123,
+ slice_name="Transitionless timestamps",
+ viz_type="bullet",
+ datasource_id=1,
+ datasource_type="table",
+ params=utils_json.dumps(form_data),
+ )
+
+ class _Command:
+ def __init__(self, _query_context: Any) -> None: ...
+
+ def validate(self) -> None: ...
+
+ def run(self) -> dict[str, Any]:
+ return producer_result
+
+ query_context = SimpleNamespace(
+ form_data={},
+ queries=[SimpleNamespace(metrics=["Revenue"], columns=["Category"])],
+ )
+ user = MagicMock(id=1, username="admin", roles=[], groups=[])
+ with (
+ patch("superset.mcp_service.auth.get_user_from_request",
return_value=user),
+ patch("superset.mcp_service.auth.check_tool_permission",
return_value=True),
+ patch.object(preview_module, "find_chart_by_identifier",
return_value=chart),
+ patch.object(preview_module.db.session, "refresh", return_value=None),
+ patch.object(
+ preview_module,
+ "validate_chart_dataset",
+ return_value=SimpleNamespace(is_valid=True, warnings=[],
error=None),
+ ),
+ patch.object(
+ preview_module.event_logger,
+ "log_context",
+ side_effect=lambda **_kwargs: nullcontext(),
+ ),
+ patch.object(
+ preview_module,
+ "build_query_context_from_form_data",
+ return_value=query_context,
+ ),
+ patch.object(preview_module, "set_query_context_form_data",
return_value=None),
+ patch.object(command_module, "ChartDataCommand", _Command),
+ patch.object(
+ preview_module, "get_superset_base_url",
return_value="http://localhost";
+ ),
+ ):
+ async with Client(mcp) as client:
+ result = await client.call_tool(
+ "get_chart_preview",
+ {"request": {"id": 123, "format": "vega_lite"}},
+ )
+
+ payload = utils_json.loads(result.content[0].text)
+ specification = payload["content"]["specification"]
+ bar = next(
+ layer for layer in specification["layer"] if layer["mark"]["type"] ==
"bar"
+ )
+ category_field = bar["encoding"]["y"]["field"]
+ assert [row["Category"] for row in specification["data"]["values"]] == (
+ expected_numbers
+ )
+ assert [row[category_field] for row in specification["data"]["values"]] ==
(
+ expected_categories
+ )
+ assert bar["encoding"]["tooltip"][0]["field"] == category_field
+
+
[email protected]
[email protected]("format_", ["ascii", "vega_lite"])
+async def test_duration_dataframe_reaches_fastmcp_bullet_preview(
+ format_: str,
+) -> None:
+ from contextlib import nullcontext
+
+ import numpy as np
+ from fastmcp import Client
+
+ from superset.commands.chart.data.get_data_command import (
+ ChartDataCommand as ProducerChartDataCommand,
+ )
+ from superset.common.chart_data import ChartDataResultType
+ from superset.dataframe import df_to_records
+ from superset.mcp_service.app import mcp
+
+ preview_module = importlib.import_module(
+ "superset.mcp_service.chart.tool.get_chart_preview"
+ )
+ command_module = importlib.import_module(
+ "superset.commands.chart.data.get_data_command"
+ )
+ source_values = [
+ timedelta(0),
+ timedelta(days=1, seconds=2, microseconds=3),
+ pd.Timedelta(-1, unit="ns"),
+ np.timedelta64(123456789, "ns"),
+ np.timedelta64("NaT"),
+ ]
+ rows = df_to_records(
+ pd.DataFrame(
+ {
+ "Duration": pd.Series(source_values, dtype=object),
+ "Revenue": [50, 120, 350, 0, 200],
+ }
+ ),
+ convert_big_integers=False,
+ )
+ expected = [
+ "null"
+ if row["Duration"] is None
+ else utils_json.loads(
+ utils_json.dumps(row["Duration"],
default=utils_json.json_int_dttm_ser)
+ )
+ for row in rows
+ ]
Review Comment:
<div>
<div id="suggestion">
<div id="issue"><b>Test crashes on np.timedelta64</b></div>
<div id="fix">
The `expected` comprehension calls `utils_json.dumps(row["Duration"],
default=json_int_dttm_ser)`, which falls through to `base_json_conv`. That
function handles `timedelta`/`pd.Timedelta` but not `np.timedelta64` (not a
`timedelta` subclass), so it raises `TypeError` for the two `np.timedelta64`
source values and the test crashes before the tool call. Also `df_to_records`
leaves `np.timedelta64("NaT")` as-is (not None), so the `"null"` branch never
fires for NaT. Normalize `np.timedelta64` (NaT→None, else `pd.Timedelta`) to
mirror `_chart_data_duration_text`.
</div>
</div>
<small><i>Code Review Run #d53793</i></small>
</div>
---
Should Bito avoid suggestions like this for future reviews? (<a
href=https://alpha.bito.ai/home/ai-agents/review-rules>Manage Rules</a>)
- [ ] Yes, avoid them
##########
tests/unit_tests/mcp_service/chart/test_bullet_chart.py:
##########
@@ -0,0 +1,3380 @@
+# 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.
+
+"""Product-path coverage for typed ECharts Bullet MCP support."""
+
+import math
+from datetime import date, datetime, time, timedelta, timezone, tzinfo
+from decimal import Decimal
+from enum import Enum, IntEnum, StrEnum
+from types import SimpleNamespace
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock, patch
+from uuid import UUID
+from zoneinfo import ZoneInfo
+
+import numpy as np
+import pandas as pd
+import pytest
+import pytz
+from dateutil import tz as dateutil_tz
+from dateutil.zoneinfo import get_zonefile_instance
+from pydantic import TypeAdapter, ValidationError
+
+from superset.mcp_service.chart.chart_helpers import (
+ build_query_dicts_from_form_data,
+)
+from superset.mcp_service.chart.chart_utils import (
+ analyze_chart_capabilities,
+ map_bullet_config,
+ map_config_to_form_data,
+ MCP_DASHBOARD_TIME_FILTER_SUBJECT,
+ merge_update_form_data,
+ validate_merged_bullet_form_data,
+)
+from superset.mcp_service.chart.compile import _compile_chart
+from superset.mcp_service.chart.preview_utils import (
+ _generate_ascii_preview_from_data,
+ _generate_vega_lite_preview_from_data,
+ _javascript_number_string,
+ BulletOutputError,
+ generate_preview_from_form_data,
+ resolve_bullet_render_model,
+)
+from superset.mcp_service.chart.query_result import (
+ _chart_data_duration_text,
+ _chart_data_temporal_number,
+)
+from superset.mcp_service.chart.schemas import (
+ ASCIIPreview,
+ BulletChartConfig,
+ ChartConfig,
+ ChartError,
+ DataColumn,
+ GenerateChartRequest,
+ GetChartPreviewRequest,
+ UpdateChartPreviewRequest,
+ UpdateChartRequest,
+ VegaLitePreview,
+ XYChartConfig,
+)
+from superset.mcp_service.chart.tool.generate_chart import generate_chart
+from superset.mcp_service.chart.tool.get_chart_data import (
+ _candidates_single_numeric,
+ _VIZ_CATEGORY,
+)
+from superset.mcp_service.chart.tool.get_chart_preview import (
+ ASCIIPreviewStrategy,
+ TablePreviewStrategy,
+ VegaLitePreviewStrategy,
+)
+from superset.mcp_service.chart.tool.get_chart_type_schema import (
+ _get_chart_type_schema_impl,
+ VALID_CHART_TYPES,
+)
+from superset.mcp_service.chart.tool.update_chart import (
+ _build_preview_form_data,
+ _build_update_payload,
+ update_chart,
+)
+from superset.mcp_service.chart.tool.update_chart_preview import
update_chart_preview
+from superset.mcp_service.chart.validation.dataset_validator import
DatasetValidator
+from superset.mcp_service.common.error_schemas import DatasetContext
+from superset.utils.json import json_int_dttm_ser
+
+
+def _reject_scalar_conversion(*_args: object, **_kwargs: object) -> Any:
+ raise AssertionError("hostile query scalar method must not run")
+
+
+class _PathHostileStr(str):
+ __getitem__ = _reject_scalar_conversion
+ __str__ = _reject_scalar_conversion
+
+
+class _PathHostileEnum(str, Enum):
+ FAILED = "warehouse unavailable"
+
+ @property
+ def value(self) -> str:
+ """Reject the public descriptor while preserving Enum's stored
value."""
+ return _reject_scalar_conversion()
+
+ __getitem__ = _reject_scalar_conversion
+ __str__ = _reject_scalar_conversion
+
+
+class _OutputHostileInt(int):
+ __float__ = _reject_scalar_conversion
+ __str__ = _reject_scalar_conversion
+
+
+class _OutputHostileFloat(float):
+ __float__ = _reject_scalar_conversion
+ __str__ = _reject_scalar_conversion
+
+
+class _OutputHostileStr(str):
+ __str__ = _reject_scalar_conversion
+ strip = _reject_scalar_conversion
+
+
+class _OutputHostileDecimal(Decimal):
+ __float__ = _reject_scalar_conversion
+ __str__ = _reject_scalar_conversion
+
+
+class _OutputSafeIntEnum(IntEnum):
+ VALUE = 12
+
+
+class _OutputSafeStrEnum(StrEnum):
+ VALUE = "12.5"
+
+
+_OutputSafeIntEnum.__float__ = _reject_scalar_conversion # type:
ignore[method-assign]
+_OutputSafeIntEnum.__str__ = _reject_scalar_conversion # type:
ignore[method-assign]
+_OutputSafeStrEnum.__float__ = _reject_scalar_conversion # type:
ignore[attr-defined]
+_OutputSafeStrEnum.__str__ = _reject_scalar_conversion # type:
ignore[method-assign]
+
+
+def _simple_metric(name: str = "revenue") -> dict[str, str]:
+ return {"name": name, "aggregate": "SUM"}
+
+
+def _tool_user() -> SimpleNamespace:
+ return SimpleNamespace(id=1, username="admin", roles=[], groups=[])
+
+
+def _orm_dataset() -> SimpleNamespace:
+ def column(
+ name: str, type_: str, *, temporal: bool = False, numeric: bool = False
+ ) -> SimpleNamespace:
+ return SimpleNamespace(
+ column_name=name,
+ type=type_,
+ is_temporal=temporal,
+ is_numeric=numeric,
+ is_dttm=temporal,
+ python_date_format=None,
+ )
+
+ return SimpleNamespace(
+ id=7,
+ table_name="sales",
+ schema=None,
+ main_dttm_col="OrderDate",
+ database=SimpleNamespace(database_name="main", db_engine_spec=None),
+ columns=[
+ column("Revenue", "NUMERIC", numeric=True),
+ column("Region", "VARCHAR"),
+ column("Team", "VARCHAR"),
+ column("Status", "VARCHAR"),
+ column("OrderDate", "TIMESTAMP", temporal=True),
+ column("EventDate", "TIMESTAMP", temporal=True),
+ ],
+ metrics=[
+ SimpleNamespace(
+ metric_name="SavedRevenue",
+ expression="SUM(Revenue)",
+ description=None,
+ )
+ ],
+ )
+
+
+def test_bullet_discriminated_union_uses_exact_tag() -> None:
+ config = TypeAdapter(ChartConfig).validate_python(
+ {"chart_type": "bullet", "metric": _simple_metric()}
+ )
+ assert isinstance(config, BulletChartConfig)
+ with pytest.raises(ValidationError):
+ TypeAdapter(ChartConfig).validate_python(
+ {"chart_type": "bullet_chart", "metric": _simple_metric()}
+ )
+
+
+def test_bullet_equal_dimension_aliases_are_order_independent_and_round_trip()
-> None:
+ for payload in (
+ {
+ "dimensions": [{"name": "Region", "label": "Market"}, "Team"],
+ "groupby": ["Region", {"column_name": "Team"}],
+ },
+ {
+ "groupby": ["Region", {"column": "Team"}],
+ "dimensions": [{"name": "Region", "label": "Market"}, "Team"],
+ },
+ ):
+ config = BulletChartConfig.model_validate(
+ {"metric": _simple_metric(), **payload}
+ )
+ assert [dimension.name for dimension in config.dimensions or []] == [
+ "Region",
+ "Team",
+ ]
+ mapped = map_bullet_config(config)
+ assert mapped["groupby"] == ["Region", "Team"]
+ round_trip = BulletChartConfig.model_validate(mapped)
+ assert [dimension.name for dimension in round_trip.dimensions or []]
== [
+ "Region",
+ "Team",
+ ]
+
+
[email protected](
+ "request_payload",
+ [
+ {"dataset_id": 7},
+ {"identifier": 9},
+ {"dataset_id": 7, "form_data_key": "preview"},
+ ],
+)
[email protected]("reverse", [False, True])
+def test_bullet_request_models_reject_conflicting_dimension_aliases(
+ request_payload: dict[str, object], reverse: bool
+) -> None:
+ aliases = [
+ ("dimensions", [{"name": "Region"}, {"name": "Team"}]),
+ ("groupby", ["Team", "Region"]),
+ ]
+ if reverse:
+ aliases.reverse()
+ config = {"chart_type": "bullet", "metric": _simple_metric(),
**dict(aliases)}
+ payload = {**request_payload, "config": config}
+ request_type = (
+ UpdateChartRequest
+ if "identifier" in request_payload
+ else (
+ UpdateChartPreviewRequest
+ if "form_data_key" in request_payload
+ else GenerateChartRequest
+ )
+ )
+ with pytest.raises(ValidationError, match="Conflicting Bullet dimension
aliases"):
+ request_type.model_validate(payload)
+
+
[email protected](
+ "metric",
+ [
+ {"name": "revenue", "aggregate": "SUM", "label": "Revenue"},
+ {"name": "saved_revenue", "saved_metric": True},
+ {"sql_expression": "SUM(revenue) / COUNT(*)", "label": "Average"},
+ ],
+)
+def test_bullet_accepts_simple_saved_and_sql_metrics(metric: dict[str,
object]) -> None:
+ config = BulletChartConfig(metric=metric)
+ form_data = map_bullet_config(config)
+ assert form_data["viz_type"] == "bullet"
+ assert form_data["metric"]
+
+
+def test_bullet_native_form_data_round_trip_is_semantically_stable() -> None:
+ native = {
+ "viz_type": "bullet",
+ "datasource": "7__table",
+ "metric": {
+ "aggregate": "SUM",
+ "column": {"column_name": "Revenue"},
+ "expressionType": "SIMPLE",
+ "label": "Total Revenue",
+ },
+ "groupby": ["Region", "Team"],
+ "ranges": "100,250,500",
+ "range_labels": "Minimum,Target,Stretch",
+ "markers": "300",
+ "marker_labels": "Plan",
+ "marker_lines": "400",
+ "marker_line_labels": "Forecast",
+ "y_axis_format": "$,.0f",
+ "show_labels": False,
+ "show_legend": True,
+ "row_limit": 250,
+ "orderby": [["Region", True], ["Total Revenue", False]],
+ "adhoc_filters": [
+ {
+ "clause": "WHERE",
+ "expressionType": "SIMPLE",
+ "subject": "Status",
+ "operator": "==",
+ "comparator": "Active",
+ }
+ ],
+ }
+ config = BulletChartConfig.model_validate(native)
+ mapped = map_bullet_config(config)
+
+ assert mapped["metric"]["label"] == "Total Revenue"
+ assert mapped["groupby"] == ["Region", "Team"]
+ assert mapped["ranges"] == "100,250,500"
+ assert mapped["range_labels"] == "Minimum,Target,Stretch"
+ assert mapped["markers"] == "300"
+ assert mapped["marker_lines"] == "400"
+ assert mapped["show_labels"] is False
+ assert mapped["show_legend"] is True
+ assert mapped["orderby"][0] == ["Region", True]
+ assert mapped["orderby"][1][0]["label"] == "Total Revenue"
+ assert mapped["adhoc_filters"][0]["subject"] == "Status"
+
+
+def test_bullet_presentation_numbers_use_shortest_round_trip_safe_tokens() ->
None:
+ ranges = [1.2345678901234567, 1.7976931348623157e308]
+ markers = [5e-324, -0.0]
+ marker_lines = [9.876543210987654e-200]
+ config = BulletChartConfig(
+ metric=_simple_metric(),
+ ranges=ranges,
+ markers=markers,
+ marker_lines=marker_lines,
+ show_legend=True,
+ )
+ mapped = map_bullet_config(config)
+
+ for key, expected in (
+ ("ranges", ranges),
+ ("markers", markers),
+ ("marker_lines", marker_lines),
+ ):
+ tokens = mapped[key].split(",")
+ assert [float(token) for token in tokens] == expected
+ assert all(
+ float(token).hex() == value.hex()
+ for token, value in zip(tokens, expected, strict=True)
+ )
+
+ round_trip = BulletChartConfig.model_validate(mapped)
+ assert round_trip.ranges == ranges
+ assert round_trip.markers == markers
+ assert round_trip.marker_lines == marker_lines
+
+ model = resolve_bullet_render_model(
+ [{"SUM(revenue)": 1.0}],
+ mapped,
+ )
+ assert model.ranges == ranges
+ assert model.markers == markers
+ assert model.marker_lines == marker_lines
+ assert (
+ "1.7976931348623157e+308"
+ in _generate_ascii_preview_from_data(
+ [{"SUM(revenue)": 1.0}], mapped
+ ).ascii_content
+ )
+ vega = _generate_vega_lite_preview_from_data([{"SUM(revenue)": 1.0}],
mapped)
+ assert vega.specification["layer"]
+
+
+def test_bullet_native_saved_metric_and_legacy_metric_aliases() -> None:
+ saved = BulletChartConfig.model_validate(
+ {"viz_type": "bullet", "metric": "saved_revenue"}
+ )
+ legacy = BulletChartConfig.model_validate(
+ {"viz_type": "bullet", "metric": "sum__revenue"}
+ )
+ assert saved.metric.saved_metric is True
+ assert saved.metric.name == "saved_revenue"
+ assert legacy.metric.name == "sum__revenue"
+ assert legacy.metric.saved_metric is True
+
+
[email protected]("metric_name", ["sum__num", "sum__SP_POP_TOTL"])
[email protected](
+ ("request_type", "request_fields"),
+ [
+ (GenerateChartRequest, {"dataset_id": 7}),
+ (UpdateChartRequest, {"identifier": 9}),
+ (UpdateChartPreviewRequest, {"dataset_id": 7}),
+ ],
+)
+def
test_bullet_repository_metric_names_round_trip_as_saved_metrics_on_all_requests(
+ metric_name: str,
+ request_type: type[
+ GenerateChartRequest | UpdateChartRequest | UpdateChartPreviewRequest
+ ],
+ request_fields: dict[str, object],
+) -> None:
+ request = request_type.model_validate(
+ {**request_fields, "config": {"chart_type": "bullet", "metric":
metric_name}}
+ )
+ config = request.config
+ assert isinstance(config, BulletChartConfig)
+ assert config.metric.saved_metric is True
+ assert map_bullet_config(config)["metric"] == metric_name
+
+
+def test_bullet_legacy_label_only_saved_metric_adapter_is_strict_and_bounded()
-> None:
+ config = BulletChartConfig.model_validate(
+ {"viz_type": "bullet", "metric": {"label": "sum__num"}}
+ )
+ assert config.metric.saved_metric is True
+ assert map_bullet_config(config)["metric"] == "sum__num"
+
+ with pytest.raises(ValidationError):
+ BulletChartConfig.model_validate(
+ {
+ "viz_type": "bullet",
+ "metric": {"label": "sum__num", "aggregate": "SUM"},
+ }
+ )
+ with pytest.raises(ValidationError, match="at most 255"):
+ BulletChartConfig.model_validate(
+ {"viz_type": "bullet", "metric": {"label": "m" * 256}}
+ )
+
+
[email protected](
+ "metric",
+ [
+ "SavedRevenue",
+ {
+ "aggregate": "SUM",
+ "column": {"column_name": "Revenue"},
+ "expressionType": "SIMPLE",
+ "label": "Simple Revenue",
+ },
+ {
+ "aggregate": None,
+ "column": None,
+ "expressionType": "SQL",
+ "sqlExpression": "SUM(Revenue)",
+ "label": "SQL Revenue",
+ },
+ ],
+)
+def test_bullet_all_metric_shapes_round_trip_full_native_presentation(
+ metric: object,
+) -> None:
+ native = {
+ "viz_type": "bullet",
+ "metric": metric,
+ "groupby": ["Region"],
+ "ranges": "50,100",
+ "range_labels": "Low,High",
+ "markers": "75",
+ "marker_labels": "Plan",
+ "marker_lines": "90",
+ "marker_line_labels": "Forecast",
+ "y_axis_format": "$,.0f",
+ "show_labels": True,
+ "show_legend": True,
+ }
+ mapped = map_bullet_config(BulletChartConfig.model_validate(native))
+ assert validate_merged_bullet_form_data(mapped) is not None
+ assert mapped["groupby"] == ["Region"]
+ assert mapped["ranges"] == "50,100"
+ assert mapped["marker_line_labels"] == "Forecast"
+
+
+def test_bullet_rejects_invalid_roles_and_output_collisions() -> None:
+ with pytest.raises(ValidationError, match="physical dimension"):
+ BulletChartConfig(
+ metric=_simple_metric(),
+ dimensions=[{"name": "region", "aggregate": "COUNT"}],
+ )
+ with pytest.raises(ValidationError, match="Duplicate Bullet dimension"):
+ BulletChartConfig(
+ metric=_simple_metric(),
+ dimensions=[{"name": "Region"}, {"name": "region"}],
+ )
+ with pytest.raises(ValidationError, match="conflicts with a dimension"):
+ BulletChartConfig(
+ metric={"name": "revenue", "aggregate": "SUM", "label": "Region"},
+ dimensions=[{"name": "region", "label": "Region"}],
+ )
+
+
+def test_bullet_rejects_misaligned_labels_and_bad_order_target() -> None:
+ with pytest.raises(ValidationError, match="one label per ranges"):
+ BulletChartConfig(
+ metric=_simple_metric(), ranges=[1, 2], range_labels=["Only one"]
+ )
+ with pytest.raises(ValidationError, match="unknown: not_a_role"):
+ BulletChartConfig(metric=_simple_metric(), order_by=[{"column":
"not_a_role"}])
+
+
+def test_bullet_dimension_labels_are_input_aliases_not_result_aliases() ->
None:
+ config = BulletChartConfig(
+ metric={"name": "Revenue", "aggregate": "SUM", "label": "Total"},
+ dimensions=[
+ {"name": "Team", "label": "Region"},
+ {"name": "Region", "label": "Market"},
+ ],
+ # The exact physical Region must win over Team's display label.
+ order_by=[
+ {"column": "Region", "ascending": True},
+ {"column": "Revenue", "ascending": False},
+ ],
+ )
+ form_data = map_bullet_config(config)
+ assert form_data["groupby"] == ["Team", "Region"]
+ assert form_data["orderby"] == [
+ ["Region", True],
+ [form_data["metric"], False],
+ ]
+
+ label_order = map_bullet_config(
+ BulletChartConfig(
+ metric=config.metric,
+ dimensions=config.dimensions,
+ order_by=[{"column": "Market"}],
+ )
+ )
+ assert label_order["orderby"] == [["Region", False]]
+
+ model = resolve_bullet_render_model(
+ [{"Team": "Blue", "Region": "North", "Total": 10}], form_data
+ )
+ assert model.dimensions == ["Team", "Region"]
+ assert [model.rows[0][name] for name in model.dimensions] == ["Blue",
"North"]
+
+
+def test_bullet_rejects_ambiguous_display_alias_for_ordering() -> None:
+ with pytest.raises(ValidationError, match="ambiguous display alias"):
+ BulletChartConfig(
+ metric=_simple_metric(),
+ dimensions=[
+ {"name": "Region", "label": "Area"},
+ {"name": "Team", "label": "area"},
+ ],
+ order_by=[{"column": "AREA"}],
+ )
+
+
[email protected](
+ ("metric", "order_target", "output"),
+ [
+ (
+ {"name": "SavedRevenue", "saved_metric": True, "label":
"Friendly"},
+ "Friendly",
+ "SavedRevenue",
+ ),
+ (
+ {"name": "Revenue", "aggregate": "SUM", "label": "Simple Total"},
+ "Revenue",
+ "Simple Total",
+ ),
+ (
+ {"sql_expression": "SUM(Revenue)", "label": "SQL Total"},
+ "SQL Total",
+ "SQL Total",
+ ),
+ ],
+)
+def test_bullet_metric_shapes_share_physical_dimension_output_contract(
+ metric: dict[str, object], order_target: str, output: str
+) -> None:
+ config = BulletChartConfig(
+ metric=metric,
+ dimensions=[{"name": "Region", "label": "Market"}],
+ order_by=[{"column": order_target}],
+ )
+ form_data = map_bullet_config(config)
+ assert form_data["groupby"] == ["Region"]
+ assert form_data["orderby"] == [[form_data["metric"], False]]
+ model = resolve_bullet_render_model([{"Region": "North", output: 12}],
form_data)
+ assert model.metric_field == output
+ assert model.dimensions == ["Region"]
+
+
+def test_bullet_mapper_preserves_omission_and_honors_explicit_values() -> None:
+ omitted = map_bullet_config(BulletChartConfig(metric=_simple_metric()))
+ explicit = map_bullet_config(
+ BulletChartConfig(
+ metric=_simple_metric(),
+ dimensions=[],
+ filters=[],
+ ranges=[],
+ show_labels=False,
+ show_legend=False,
+ row_limit=42,
+ time_range=None,
+ )
+ )
+ for key in (
+ "groupby",
+ "adhoc_filters",
+ "ranges",
+ "show_labels",
+ "show_legend",
+ "row_limit",
+ "time_range",
+ ):
+ assert key not in omitted
+ assert explicit["groupby"] == []
+ assert explicit["adhoc_filters"] == []
+ assert explicit["ranges"] == ""
+ assert explicit["show_labels"] is False
+ assert explicit["show_legend"] is False
+ assert explicit["row_limit"] == 42
+ assert explicit["time_range"] is None
+
+
+def test_bullet_registry_schema_and_recommendation_metadata() -> None:
+ from superset.mcp_service.app import get_default_instructions
+ from superset.mcp_service.chart.registry import display_name_for_viz_type,
get
+
+ plugin = get("bullet")
+ assert plugin is not None
+ assert plugin.resolve_viz_type(None) == "bullet"
+ assert display_name_for_viz_type("bullet") == "Bullet Chart"
+ assert "bullet" in VALID_CHART_TYPES
+ discovered = _get_chart_type_schema_impl("bullet")
+ assert discovered["chart_type"] == "bullet"
+ assert discovered["examples"][0]["ranges"] == [100000, 250000, 500000]
+ assert _VIZ_CATEGORY["bullet"] == "bullet"
+ candidates = _candidates_single_numeric(
+ DataColumn(
+ name="Revenue",
+ display_name="Revenue",
+ data_type="numeric",
+ sample_values=[1],
+ null_count=0,
+ unique_count=1,
+ ),
+ row_count=1,
+ )
+ assert "bullet chart" in candidates
+ guidance = get_default_instructions()
+ assert 'chart_type="bullet": Bullet Chart' in guidance
+ assert "waterfall, bullet, and interactive_pivot" in guidance
+
+
+def test_bullet_dataset_normalization_canonicalizes_every_reference() -> None:
+ from superset.mcp_service.chart.registry import get
+
+ context = DatasetContext(
+ id=7,
+ table_name="sales",
+ schema=None,
+ database_name="main",
+ available_columns=[
+ {"name": "Revenue", "type": "NUMERIC", "is_numeric": True},
+ {"name": "Region", "type": "VARCHAR"},
+ {"name": "OrderDate", "type": "TIMESTAMP", "is_temporal": True},
+ {"name": "Status", "type": "VARCHAR"},
+ ],
+ available_metrics=[],
+ )
+ config = BulletChartConfig(
+ metric={"name": "revenue", "aggregate": "SUM"},
+ dimensions=[{"name": "region"}],
+ temporal_column="orderdate",
+ filters=[{"column": "status", "op": "=", "value": "active"}],
+ order_by=[{"column": "region", "ascending": True}],
+ )
+ plugin = get("bullet")
+ assert plugin is not None
+ normalized = plugin.normalize_column_refs(config, context)
+ assert normalized.metric.name == "Revenue"
+ assert normalized.dimensions[0].name == "Region"
+ assert normalized.temporal_column == "OrderDate"
+ assert normalized.filters[0].column == "Status"
+ assert normalized.order_by[0].column == "Region"
+ assert normalized.model_fields_set == config.model_fields_set
+
+
+def test_bullet_dataset_normalization_rejects_ambiguous_casefold_candidates()
-> None:
+ from superset.mcp_service.chart.registry import get
+
+ context = DatasetContext(
+ id=7,
+ table_name="sales",
+ schema=None,
+ database_name="main",
+ available_columns=[
+ {"name": "Revenue", "type": "NUMERIC", "is_numeric": True},
+ {"name": "revenue", "type": "NUMERIC", "is_numeric": True},
+ ],
+ available_metrics=[],
+ )
+ plugin = get("bullet")
+ assert plugin is not None
+ config = BulletChartConfig(metric={"name": "REVENUE", "aggregate": "SUM"})
+ with pytest.raises(ValueError, match="Revenue, revenue"):
+ plugin.normalize_column_refs(config, context)
+
+
+def test_bullet_numeric_output_constraint_rejects_text_min() -> None:
+ from superset.mcp_service.chart.registry import get
+
+ context = DatasetContext(
+ id=7,
+ table_name="sales",
+ schema=None,
+ database_name="main",
+ available_columns=[{"name": "status", "type": "VARCHAR"}],
+ available_metrics=[],
+ )
+ plugin = get("bullet")
+ assert plugin is not None
+ config = BulletChartConfig(metric={"name": "status", "aggregate": "MIN"})
+ with patch.object(DatasetValidator, "_get_dataset_context",
return_value=context):
+ error = plugin.post_map_validate(config, {}, dataset_id=7)
+ assert error is not None
+ assert error.error_type == "non_numeric_bullet_metric"
+
+
[email protected]("reverse_metadata", [False, True])
+def test_bullet_exact_case_type_and_role_resolution_is_order_independent(
+ reverse_metadata: bool,
+) -> None:
+ from superset.mcp_service.chart.registry import get
+
+ columns = [
+ {"name": "Revenue", "type": "NUMERIC", "is_numeric": True},
+ {"name": "revenue", "type": "VARCHAR", "is_numeric": False},
+ {"name": "Region", "type": "VARCHAR"},
+ ]
+ if reverse_metadata:
+ columns.reverse()
+ context = DatasetContext(
+ id=7,
+ table_name="sales",
+ schema=None,
+ database_name="main",
+ available_columns=columns,
+ available_metrics=[],
+ )
+ plugin = get("bullet")
+ assert plugin is not None
+
+ numeric = BulletChartConfig(metric={"name": "Revenue", "aggregate": "SUM"})
+ text = BulletChartConfig(metric={"name": "revenue", "aggregate": "MIN"})
+ with patch.object(DatasetValidator, "_get_dataset_context",
return_value=context):
+ assert plugin.post_map_validate(numeric, {}, dataset_id=7) is None
+ error = plugin.post_map_validate(text, {}, dataset_id=7)
+ assert error is not None
+ assert error.error_type == "non_numeric_bullet_metric"
+
+ roles = BulletChartConfig(
+ metric={"name": "Revenue", "aggregate": "SUM"},
+ dimensions=[{"name": "revenue"}, {"name": "Region"}],
+ filters=[{"column": "revenue", "op": "=", "value": "retail"}],
+ order_by=[{"column": "revenue", "ascending": True}],
+ )
+ normalized = plugin.normalize_column_refs(roles, context)
+ assert normalized.metric.name == "Revenue"
+ assert [dimension.name for dimension in normalized.dimensions or []] == [
+ "revenue",
+ "Region",
+ ]
+ assert normalized.filters
+ assert normalized.filters[0].column == "revenue"
+ assert normalized.order_by[0].column == "revenue"
+
+ ambiguous = BulletChartConfig(metric={"name": "REVENUE", "aggregate":
"SUM"})
+ with pytest.raises(ValueError, match="Ambiguous"):
+ plugin.normalize_column_refs(ambiguous, context)
+
+
[email protected]("reverse_metadata", [False, True])
+def test_generic_aggregation_validation_uses_exact_case_before_type(
+ reverse_metadata: bool,
+) -> None:
+ from superset.mcp_service.chart.schemas import PieChartConfig
+
+ columns = [
+ {"name": "Amount", "type": "BIGINT", "is_numeric": True},
+ {"name": "amount", "type": "VARCHAR", "is_numeric": False},
+ ]
+ if reverse_metadata:
+ columns.reverse()
+ context = DatasetContext(
+ id=7,
+ table_name="sales",
+ schema=None,
+ database_name="main",
+ available_columns=columns,
+ available_metrics=[],
+ )
+
+ assert (
+ DatasetValidator._validate_aggregations(
+ [BulletChartConfig(metric={"name": "Amount", "aggregate":
"SUM"}).metric],
+ context,
+ )
+ == []
+ )
+ errors = DatasetValidator._validate_aggregations(
+ [BulletChartConfig(metric={"name": "amount", "aggregate":
"SUM"}).metric],
+ context,
+ )
+ assert errors
+ assert errors[0].error_type == "invalid_aggregation"
+
+ ambiguous = DatasetValidator._validate_aggregations(
+ [BulletChartConfig(metric={"name": "AMOUNT", "aggregate":
"SUM"}).metric],
+ context,
+ )
+ assert ambiguous
+ assert ambiguous[0].error_type == "ambiguous_column_reference"
+
+ valid, error = DatasetValidator.validate_against_dataset(
+ PieChartConfig(
+ dimension={"name": "amount"},
+ metric={"name": "Amount", "aggregate": "SUM"},
+ ),
+ 7,
+ dataset_context=context,
+ )
+ assert valid is True
+ assert error is None
+
+
[email protected](
+ ("metric", "field"),
+ [
+ ({"name": "Revenue", "aggregate": "SUM", "label": "Simple"}, "Simple"),
+ ({"name": "SavedRevenue", "saved_metric": True}, "SavedRevenue"),
+ ({"sql_expression": "SUM(Revenue)", "label": "SQL Total"}, "SQL
Total"),
+ ],
+)
+def test_bullet_result_roles_are_exact_for_every_metric_shape(
+ metric: dict[str, object], field: str
+) -> None:
+ form_data = map_bullet_config(BulletChartConfig(metric=metric))
+ model = resolve_bullet_render_model([{field.swapcase(): "12.5"}],
form_data)
+ assert model.metric_field == field.swapcase()
+ assert model.measures == [12.5]
+
+
[email protected](
+ "rows, message",
+ [
+ ([{"other": 123}], "missing"),
+ ([{"Revenue": "not a number"}], "non-numeric text"),
+ ([{"Revenue": math.nan}], "NaN or infinite"),
+ ([{"Revenue": math.inf}], "NaN or infinite"),
+ ([{"Revenue": 1}, {}], "row 1.*missing"),
+ ([{"REVENUE": 1, "revenue": 2}], "ambiguous"),
+ ],
+)
+def test_bullet_result_validation_rejects_malformed_rows(
+ rows: list[dict[str, object]], message: str
+) -> None:
+ form_data = map_bullet_config(
+ BulletChartConfig(
+ metric={"name": "amount", "aggregate": "SUM", "label": "Revenue"}
+ )
+ )
+ with pytest.raises(BulletOutputError, match=message):
+ resolve_bullet_render_model(rows, form_data)
+
+
+def test_bullet_result_validation_accepts_null_and_numeric_strings() -> None:
+ form_data = map_bullet_config(
+ BulletChartConfig(
+ metric={"name": "amount", "aggregate": "SUM", "label": "Revenue"},
+ dimensions=[{"name": "Region"}],
+ )
+ )
+ model = resolve_bullet_render_model(
+ [
+ {"Region": "North", "Revenue": None},
+ {"Region": "South", "Revenue": " 4.25 "},
+ ],
+ form_data,
+ )
+ assert model.measures == [0.0, 4.25]
+
+
[email protected](
+ ("presentation", "message"),
+ [
+ ({"ranges": "10,nope"}, r"ranges\[1\].*not numeric"),
+ ({"markers": "NaN"}, r"markers\[0\].*NaN or infinite"),
+ (
+ {"ranges": "10,20", "range_labels": "Only one"},
+ "one label per value",
+ ),
+ ],
+)
+def test_bullet_result_validation_rejects_malformed_presentation(
+ presentation: dict[str, object], message: str
+) -> None:
+ form_data = {
+ **map_bullet_config(
+ BulletChartConfig(metric={"name": "amount", "aggregate": "SUM"})
+ ),
+ **presentation,
+ }
+ with pytest.raises(BulletOutputError, match=message):
+ resolve_bullet_render_model([{"SUM(amount)": 1}], form_data)
+
+
+def test_bullet_compile_accepts_empty_ungrouped_result() -> None:
+ form_data = map_bullet_config(
+ BulletChartConfig(
+ metric={"name": "Revenue", "aggregate": "SUM", "label": "Revenue"}
+ )
+ )
+ factory = MagicMock()
+ factory.create.return_value = MagicMock()
+ command = MagicMock()
+ command.run.return_value = {"queries": [{"data": []}]}
+ with (
+ patch(
+ "superset.common.query_context_factory.QueryContextFactory",
+ return_value=factory,
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand",
+ return_value=command,
+ ),
+ ):
+ result = _compile_chart(form_data, 7)
+ assert result.success is True
+ assert result.row_count == 0
+
+
+def test_bullet_compile_inspects_top_level_and_query_error_envelopes() -> None:
+ form_data = map_bullet_config(
+ BulletChartConfig(metric={"name": "Revenue", "aggregate": "SUM"})
+ )
+ factory = MagicMock()
+ factory.create.return_value = MagicMock()
+ command = MagicMock()
+ command.run.return_value = {
+ "status": "success",
+ "queries": [{"status": "failed", "message": "warehouse timeout"}],
+ }
+ with (
+ patch(
+ "superset.common.query_context_factory.QueryContextFactory",
+ return_value=factory,
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand",
+ return_value=command,
+ ),
+ ):
+ result = _compile_chart(form_data, 7)
+ assert result.success is False
+ assert "warehouse timeout" in (result.error or "")
+
+
+_MALFORMED_QUERY_ENVELOPES: list[object] = [
+ None,
+ [],
+ {},
+ {"queries": None},
+ {"queries": []},
+ {"queries": [None]},
+ {"queries": [{}]},
+ {"queries": [{"data": None}]},
+ {"queries": [{"data": []}, {"data": "not-an-array"}]},
+ {
+ "queries": [
+ {
+ "data": [{"Revenue": 12}],
+ "colnames": ["Revenue"],
+ "coltypes": [],
+ }
+ ]
+ },
+]
+
+
+def _compile_bullet_with_result(result: object) -> Any:
+ form_data = map_bullet_config(
+ BulletChartConfig(
+ metric={"name": "Revenue", "aggregate": "SUM", "label": "Revenue"}
+ )
+ )
+ factory = MagicMock()
+ factory.create.return_value = MagicMock()
+ command = MagicMock()
+ command.run.return_value = result
+ with (
+ patch(
+ "superset.common.query_context_factory.QueryContextFactory",
+ return_value=factory,
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand",
+ return_value=command,
+ ),
+ ):
+ return _compile_chart(form_data, 7)
+
+
[email protected]("envelope", _MALFORMED_QUERY_ENVELOPES)
+def test_bullet_compile_returns_stable_error_for_malformed_envelopes(
+ envelope: object,
+) -> None:
+ result = _compile_bullet_with_result(envelope)
+ assert result.success is False
+ assert result.error_code == "CHART_COMPILE_FAILED"
+ assert result.error_obj is not None
+ assert result.error_obj.error_type == "compile_error"
+
+
[email protected](
+ ("data", "expected_code", "expected_type"),
+ [
+ ([1], "CHART_COMPILE_FAILED", "compile_error"),
+ (
+ [{"Revenue": 10**10000}],
+ "CHART_COMPILE_FAILED",
+ "compile_error",
+ ),
+ ],
+)
+def test_bullet_compile_returns_malformed_output_for_bad_rows(
+ data: list[object],
+ expected_code: str,
+ expected_type: str,
+) -> None:
+ result = _compile_bullet_with_result({"queries": [{"data": data}]})
+ assert result.success is False
+ assert result.error_code == expected_code
+ assert result.error_obj is not None
+ assert result.error_obj.error_type == expected_type
+
+
+def test_bullet_shared_query_builder_matches_frontend_build_query() -> None:
+ metric = map_bullet_config(
+ BulletChartConfig(
+ metric={"name": "Revenue", "aggregate": "SUM", "label": "Revenue"},
+ dimensions=[{"name": "Region"}],
+ row_limit=25,
+ order_by=[{"column": "Revenue", "ascending": False}],
+ )
+ )
+ with patch(
+ "superset.mcp_service.chart.chart_helpers.resolve_datasource_engine",
+ return_value="base",
+ ):
+ queries = build_query_dicts_from_form_data(metric, 7, "table")
+ assert len(queries) == 1
+ assert queries[0]["columns"] == ["Region"]
+ assert queries[0]["metrics"] == [metric["metric"]]
+ assert queries[0]["orderby"] == metric["orderby"]
+ assert queries[0]["row_limit"] == 25
+
+
+def test_bullet_compile_path_uses_groupby_metric_orderby_and_usable_result()
-> None:
+ form_data = map_bullet_config(
+ BulletChartConfig(
+ metric={"name": "Revenue", "aggregate": "SUM", "label": "Revenue"},
+ dimensions=[{"name": "Region"}],
+ order_by=[{"column": "Revenue", "ascending": False}],
+ )
+ )
+ context = MagicMock()
+ factory = MagicMock()
+ factory.create.return_value = context
+ command = MagicMock()
+ command.run.return_value = {
+ "queries": [{"data": [{"Region": "North", "Revenue": "12.5"}]}]
+ }
+ with (
+ patch(
+ "superset.common.query_context_factory.QueryContextFactory",
+ return_value=factory,
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand",
+ return_value=command,
+ ),
+ ):
+ result = _compile_chart(form_data, 7)
+ query = factory.create.call_args.kwargs["queries"][0]
+ assert query["columns"] == ["Region"]
+ assert query["metrics"] == [form_data["metric"]]
+ assert query["orderby"] == form_data["orderby"]
+ assert result.success is True
+ assert result.row_count == 1
+
+
+def test_bullet_compile_projects_dataframe_timestamp_before_validation() ->
None:
+ from superset.dataframe import df_to_records
+
+ dublin = dateutil_tz.gettz("Europe/Dublin")
+ new_york = dateutil_tz.gettz("America/New_York")
+ assert dublin is not None
+ assert new_york is not None
+ source_values = [
+ pd.Timestamp("2024-01-02 03:04:05.123456789"),
+ datetime(2024, 10, 27, 1, 30, tzinfo=dublin, fold=1),
+ datetime(2024, 3, 10, 2, 30, tzinfo=new_york),
+ datetime(2040, 7, 1, 12, tzinfo=new_york),
+ ]
+ form_data = map_bullet_config(
+ BulletChartConfig(
+ metric={"name": "Revenue", "aggregate": "SUM", "label": "Revenue"},
+ dimensions=[{"name": "Category"}],
+ )
+ )
+ rows = df_to_records(
+ pd.DataFrame(
+ {
+ "Category": pd.Series(source_values, dtype=object),
+ "Revenue": range(1, len(source_values) + 1),
+ }
+ ),
+ convert_big_integers=False,
+ )
+ factory = MagicMock()
+ factory.create.return_value = MagicMock()
+ command = MagicMock()
+ command.run.return_value = {"queries": [{"data": rows}]}
+ captured: list[list[dict[str, Any]]] = []
+
+ def validate(data: list[dict[str, Any]], config: dict[str, Any]) -> Any:
+ captured.append(data)
+ return resolve_bullet_render_model(data, config)
+
+ with (
+ patch(
+ "superset.common.query_context_factory.QueryContextFactory",
+ return_value=factory,
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand",
+ return_value=command,
+ ),
+ patch(
+
"superset.mcp_service.chart.preview_utils.resolve_bullet_render_model",
+ side_effect=validate,
+ ),
+ ):
+ result = _compile_chart(form_data, 7)
+
+ assert result.success is True
+ assert [row["Category"] for row in captured[0]] == [
+ 1704164645123.456,
+ 1729989000000.0,
+ 1710052200000.0,
+ 2224774800000.0,
+ ]
+
+
+def test_bullet_compile_projects_real_dataframe_durations_to_chart_data_wire()
-> None:
+ from superset.commands.chart.data.get_data_command import ChartDataCommand
+ from superset.common.chart_data import ChartDataResultType
+ from superset.dataframe import df_to_records
+ from superset.utils import json
+
+ source_values = [
+ timedelta(0),
+ timedelta(days=1, seconds=2, microseconds=3),
+ pd.Timedelta(-1, unit="ns"),
+ pd.Timedelta("1 days 00:00:02.000003004"),
+ np.timedelta64(123456789, "ns"),
+ np.timedelta64("NaT"),
+ ]
+ rows = df_to_records(
+ pd.DataFrame(
+ {
+ "Duration": pd.Series(source_values, dtype=object),
+ "Revenue": range(1, len(source_values) + 1),
+ }
+ ),
+ convert_big_integers=False,
+ )
+ expected = [
+ None
+ if row["Duration"] is None
+ else json.loads(json.dumps(row["Duration"],
default=json.json_int_dttm_ser))
+ for row in rows
+ ]
Review Comment:
<div>
<div id="suggestion">
<div id="issue"><b>Test crashes on np.timedelta64</b></div>
<div id="fix">
`json.json_int_dttm_ser` cannot serialize `np.timedelta64` (falls through
`base_json_conv` and raises `TypeError`). Since `source_values` always includes
`np.timedelta64(123456789, "ns")` and `np.timedelta64("NaT")`, and
`df_to_records` leaves them as-is, this `expected` computation crashes the
test. Use `_chart_data_duration_text` (imported here) to match the real wire
projection, and remove the now-unused `json` import.
</div>
<details>
<summary>
<b>Code suggestion</b>
</summary>
<blockquote>Check the AI-generated fix before applying</blockquote>
<div id="code">
````suggestion
)
expected = []
for row in rows:
text, reason = _chart_data_duration_text(row["Duration"])
assert reason is None
expected.append(text)
````
</div>
</details>
</div>
<small><i>Code Review Run #d53793</i></small>
</div>
---
Should Bito avoid suggestions like this for future reviews? (<a
href=https://alpha.bito.ai/home/ai-agents/review-rules>Manage Rules</a>)
- [ ] Yes, avoid them
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
