feat(mcp): chart type plugin registry for extensible generate_chart

[aminghadersohi]

SUMMARY

Introduces a ChartTypePlugin protocol and central registry that consolidates all chart-type-specific logic into a single plugin class per chart type. Previously, adding a new chart type required synchronised edits to four separate dispatch locations. With this change it requires one new file.

The problem: 4 dispatch points that had to stay in sync

Before this PR, generate_chart logic was split across:

Location	Responsibility
schema_validator.py	7 _pre_validate_* static methods, one per type
dataset_validator.py	isinstance branches in _extract_column_references and normalize_column_names
chart_utils.py	if/elif chain in map_config_to_form_data, generate_chart_name, _resolve_viz_type
runtime/__init__.py	isinstance(config, XYChartConfig) hardcoded branch

Result: 5 of 7 chart types (pie, pivot_table, mixed_timeseries, handlebars, big_number) silently skipped column validation — a bug discovered only when charts failed at query time.

---

Architecture

New file layout

superset/mcp_service/chart/
├── plugin.py          # ChartTypePlugin Protocol + BaseChartPlugin base class
├── registry.py        # Central registry with lazy bootstrap
└── plugins/
    ├── __init__.py    # Imports and registers all 7 built-in plugins
    ├── xy.py
    ├── table.py
    ├── pie.py
    ├── pivot_table.py
    ├── mixed_timeseries.py
    ├── handlebars.py
    └── big_number.py

Request flow (before → after)

flowchart LR
    subgraph BEFORE["Before: 4 separate dispatch points"]
        direction TB
        A1[schema_validator.py\n7 pre_validate methods]
        A2[dataset_validator.py\nisinstance branches]
        A3[chart_utils.py\nif/elif chains]
        A4[runtime/__init__.py\nXY hardcoding]
    end

    subgraph AFTER["After: single registry lookup"]
        direction TB
        B1[registry.get chart_type] --> B2[plugin.pre_validate]
        B1 --> B3[plugin.extract_column_refs]
        B1 --> B4[plugin.to_form_data]
        B1 --> B5[plugin.post_map_validate]
        B1 --> B6[plugin.normalize_column_refs]
        B1 --> B7[plugin.get_runtime_warnings]
        B1 --> B8[plugin.generate_name]
        B1 --> B9[plugin.resolve_viz_type]
    end

    BEFORE -.->|refactored to| AFTER

Loading

Plugin protocol

classDiagram
    class ChartTypePlugin {
        <<Protocol>>
        +str chart_type
        +str display_name
        +dict native_viz_types
        +pre_validate(config) ChartGenerationError|None
        +extract_column_refs(config) list[ColumnRef]
        +to_form_data(config, dataset_id) dict
        +post_map_validate(config, form_data, dataset_id) ChartGenerationError|None
        +normalize_column_refs(config, dataset_context) Any
        +get_runtime_warnings(config, dataset_id) list[str]
        +generate_name(config, dataset_name) str
        +resolve_viz_type(config) str
    }

    class BaseChartPlugin {
        +pre_validate() None
        +extract_column_refs() []
        +to_form_data() NotImplementedError
        +post_map_validate() None
        +normalize_column_refs() config unchanged
        +get_runtime_warnings() []
        +generate_name() "Chart"
        +resolve_viz_type() "unknown"
        +_with_context(what, context)$ str
    }

    class XYChartPlugin
    class PieChartPlugin
    class TableChartPlugin
    class BigNumberChartPlugin
    class PivotTableChartPlugin
    class MixedTimeseriesChartPlugin
    class HandlebarsChartPlugin

    ChartTypePlugin <|.. BaseChartPlugin
    BaseChartPlugin <|-- XYChartPlugin
    BaseChartPlugin <|-- PieChartPlugin
    BaseChartPlugin <|-- TableChartPlugin
    BaseChartPlugin <|-- BigNumberChartPlugin
    BaseChartPlugin <|-- PivotTableChartPlugin
    BaseChartPlugin <|-- MixedTimeseriesChartPlugin
    BaseChartPlugin <|-- HandlebarsChartPlugin

Loading

Registry bootstrap (lazy, no circular imports)

sequenceDiagram
    participant Caller as Caller (test / tool)
    participant Registry as registry.py
    participant Plugins as plugins/__init__.py

    Caller->>Registry: get("xy")
    Registry->>Registry: _ensure_plugins_loaded()
    alt first call
        Registry->>Plugins: import (triggers register() calls)
        Plugins-->>Registry: 7 plugins registered
    end
    Registry-->>Caller: XYChartPlugin instance

Loading

The registry self-bootstraps on first lookup — no dependency on app.py being imported first, so isolated unit tests work without the full Flask app.

---

How to add a new chart type

Everything goes in one new file. Here is a minimal working example for a hypothetical funnel chart:

Step 1 — Add the Pydantic schema (chart/schemas.py):

class FunnelChartConfig(ChartConfig):
    chart_type: Literal["funnel"] = "funnel"
    metric: ColumnRef
    groupby: ColumnRef
    filters: list[FilterConfig] | None = None

Step 2 — Create the plugin (chart/plugins/funnel.py):

from superset.mcp_service.chart.plugin import BaseChartPlugin
from superset.mcp_service.chart.schemas import ColumnRef
from superset.mcp_service.common.error_schemas import ChartGenerationError


class FunnelChartPlugin(BaseChartPlugin):
    chart_type = "funnel"
    display_name = "Funnel Chart"
    native_viz_types = {"funnel": "Funnel Chart"}

    def pre_validate(self, config: dict) -> ChartGenerationError | None:
        if "metric" not in config or "groupby" not in config:
            return ChartGenerationError(
                error_type="missing_funnel_fields",
                message="Funnel chart requires 'metric' and 'groupby'",
                suggestions=["Add metric: {'name': 'col', 'aggregate': 'SUM'}"],
                error_code="MISSING_FUNNEL_FIELDS",
            )
        return None

    def extract_column_refs(self, config) -> list[ColumnRef]:
        from superset.mcp_service.chart.schemas import FunnelChartConfig
        if not isinstance(config, FunnelChartConfig):
            return []
        refs = [config.metric, config.groupby]
        if config.filters:
            refs.extend(ColumnRef(name=f.column) for f in config.filters)
        return refs

    def to_form_data(self, config, dataset_id=None) -> dict:
        return {
            "viz_type": "funnel",
            "metric": {"column": {"column_name": config.metric.name},
                       "aggregate": config.metric.aggregate},
            "groupby": [config.groupby.name],
        }

    def resolve_viz_type(self, config) -> str:
        return "funnel"

    def generate_name(self, config, dataset_name=None) -> str:
        return self._with_context(
            f"Funnel: {config.metric.aggregate}({config.metric.name})",
            dataset_name,
        )

    def normalize_column_refs(self, config, dataset_context):
        from superset.mcp_service.chart.schemas import FunnelChartConfig
        from superset.mcp_service.chart.validation.dataset_validator import DatasetValidator
        d = config.model_dump()
        get = DatasetValidator._get_canonical_column_name
        d["metric"]["name"] = get(d["metric"]["name"], dataset_context)
        d["groupby"]["name"] = get(d["groupby"]["name"], dataset_context)
        DatasetValidator._normalize_filters(d, dataset_context)
        return FunnelChartConfig.model_validate(d)

Step 3 — Register it (chart/plugins/__init__.py):

from superset.mcp_service.chart.plugins.funnel import FunnelChartPlugin
# ...existing imports...
register(FunnelChartPlugin())

Step 4 — Add a Pydantic error hint (chart/validation/schema_validator.py):

# In _CHART_TYPE_ERROR_HINTS — keep in sync with plugins/
"funnel": {
    "metric": "Add metric: {'name': 'col', 'aggregate': 'SUM'}",
    "groupby": "Add groupby: {'name': 'stage_column'}",
},

That is the complete change set — no edits to chart_utils.py, dataset_validator.py, or runtime/__init__.py dispatch logic.

---

Bug fixes included

• 5-type column validation gap: pie, pivot_table, mixed_timeseries, handlebars, and big_number charts previously skipped dataset column validation entirely (silent false-pass). All 7 types now run column existence and aggregation checks.
• BigNumber trendline check: Moved out of the monolithic dispatcher into BigNumberChartPlugin.post_map_validate().
• Runtime validator decoupling: Removed isinstance(config, XYChartConfig) hardcoding from RuntimeValidator; XY-specific format/cardinality checks now live in XYChartPlugin.get_runtime_warnings().
• Stale docstring: generate_chart.py listed only 'xy' and 'table' as valid chart types; updated to list all 7.
• Pydantic error handling: Added missing pie, pivot_table, mixed_timeseries cases to _enhance_validation_error; refactored to a data-driven lookup table to reduce cyclomatic complexity.
• Registry bootstrap: Registry self-bootstraps on first access so unit tests running without app.py get a fully populated registry.

BEFORE/AFTER SCREENSHOTS OR ANIMATED GIF

N/A — backend-only change, no UI modifications.

TESTING INSTRUCTIONS

# Run the MCP unit tests for chart functionality
pytest tests/unit_tests/mcp_service/chart/ -x -q

# Spot-check column validation now works for previously-unvalidated types
pytest tests/unit_tests/mcp_service/chart/validation/ -x -q

Manual verification:

1. Start the MCP service: python -m superset.mcp_service.server
2. Create a pie chart with a non-existent dimension column → should get a column_not_found error with fuzzy suggestions (previously: silent pass, runtime failure)
3. Create a big_number chart with show_trendline=true but no temporal_column → should error at validation
4. Run generate_chart for all 7 chart types to confirm no regressions

ADDITIONAL INFORMATION

• Has associated issue:
• Required feature flags:
• Changes UI
• Includes DB Migration
• Introduces new feature or API
• Removes existing feature or API

[netlify]

✅ Deploy Preview for superset-docs-preview ready!

Name	Link
🔨 Latest commit	8f656ee
🔍 Latest deploy log	https://app.netlify.com/projects/superset-docs-preview/deploys/69fbc7b02ff27400092aa051
😎 Deploy Preview	https://deploy-preview-39922--superset-docs-preview.netlify.app
📱 Preview on mobile	Toggle QR Code... Use your smartphone camera to open QR code link.
🤖 Make changes	Run an agent on this branch

---

To edit notification comments on pull requests, go to your Netlify project configuration.

[bito-code-review]

The inline context is insufficient to determine what 'this' refers to in the question. No diff hunk or suggestion snippet is provided for the thread in superset/mcp_service/chart/chart_utils.py.

[Copilot]

Pull request overview

This PR introduces a chart-type plugin architecture for the MCP chart pipeline, centralizing per-chart-type behavior (pre-validation, column extraction/normalization, form_data mapping, runtime warnings, naming, viz_type resolution) into one plugin per supported chart_type. This reduces the previous need to update multiple dispatch sites when adding or modifying chart types.

Changes:

• Added a ChartTypePlugin protocol + BaseChartPlugin, a global plugin registry, and 7 built-in chart plugins (xy, table, pie, pivot_table, mixed_timeseries, handlebars, big_number).
• Refactored schema/dataset/runtime validation and form_data mapping to dispatch via the registry instead of isinstance / if/elif chains.
• Enhanced tool/docs output by expanding generate_chart docstring chart types and adding chart_type_display_name to ChartInfo.

Reviewed changes

Copilot reviewed 17 out of 17 changed files in this pull request and generated 9 comments.

Show a summary per file

File	Description
superset/mcp_service/chart/validation/schema_validator.py	Switch pre-validation dispatch to registry/plugins; refactor Pydantic error enhancement.
superset/mcp_service/chart/validation/runtime/init.py	Replace XY-only runtime checks with per-plugin runtime warning dispatch.
superset/mcp_service/chart/validation/dataset_validator.py	Delegate column extraction/normalization to plugins; apply aggregation validation for all chart types.
superset/mcp_service/chart/tool/generate_chart.py	Update tool docstring to list all supported chart_type values.
superset/mcp_service/chart/schemas.py	Add chart_type_display_name to ChartInfo; resolve display name via registry mapping.
superset/mcp_service/chart/registry.py	Introduce global registry for chart plugins + viz_type display-name lookup.
superset/mcp_service/chart/plugin.py	Add plugin protocol and base class defining the plugin responsibilities.
superset/mcp_service/chart/plugins/init.py	Register the built-in plugins on import.
superset/mcp_service/chart/plugins/xy.py	Implement XY plugin (pre-validate, refs, mapping, naming, viz_type, runtime warnings).
superset/mcp_service/chart/plugins/table.py	Implement Table plugin.
superset/mcp_service/chart/plugins/pie.py	Implement Pie plugin.
superset/mcp_service/chart/plugins/pivot_table.py	Implement Pivot Table plugin.
superset/mcp_service/chart/plugins/mixed_timeseries.py	Implement Mixed Timeseries plugin.
superset/mcp_service/chart/plugins/handlebars.py	Implement Handlebars plugin.
superset/mcp_service/chart/plugins/big_number.py	Implement Big Number plugin including post-map trendline temporal validation.
superset/mcp_service/chart/chart_utils.py	Replace mapping/name/viz_type dispatch with registry/plugin calls.
superset/mcp_service/app.py	Import plugins at MCP app startup; update default instructions about chart type naming.

[aminghadersohi]

do these have to be local? can you check the rest of the pr and move up imports that can be moved up

[aminghadersohi]

Done in 891676f. All per-method local imports across all 7 plugin files have been moved to top level. None create circular imports — schemas.py, chart_utils.py, and dataset_validator.py are safe to import at plugin load time because those modules guard their own registry lookups with local imports (the two cases documented in prior replies). The only remaining local import in the whole plugin tree is import superset.mcp_service.chart.plugins inside _ensure_plugins_loaded() in registry.py — that one has to stay local since it's the bootstrap itself.