Skip to content

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

Draft
aminghadersohi wants to merge 6 commits intomasterfrom
mcp-chart-tools-rearch
Draft

feat(mcp): chart type plugin registry for extensible generate_chart#39922
aminghadersohi wants to merge 6 commits intomasterfrom
mcp-chart-tools-rearch

Conversation

@aminghadersohi
Copy link
Copy Markdown
Contributor

@aminghadersohi aminghadersohi commented May 6, 2026

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
Copy link
Copy Markdown

netlify Bot commented May 6, 2026

Deploy Preview for superset-docs-preview ready!

Name Link
🔨 Latest commit 1e2b541
🔍 Latest deploy log https://app.netlify.com/projects/superset-docs-preview/deploys/69fd0fe2ca5f1f0008dacace
😎 Deploy Preview https://deploy-preview-39922--superset-docs-preview.netlify.app
📱 Preview on mobile
Toggle QR Code...

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.

Comment thread superset/mcp_service/chart/chart_utils.py
Comment thread superset/mcp_service/chart/validation/dataset_validator.py
@bito-code-review
Copy link
Copy Markdown
Contributor

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.

Comment thread superset/mcp_service/chart/plugins/big_number.py Outdated
Comment thread superset/mcp_service/chart/plugins/big_number.py Outdated
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread superset/mcp_service/chart/registry.py
Comment thread superset/mcp_service/chart/validation/schema_validator.py
Comment thread superset/mcp_service/chart/validation/schema_validator.py
Comment thread superset/mcp_service/chart/validation/schema_validator.py Outdated
Comment thread superset/mcp_service/chart/validation/runtime/__init__.py
Comment thread superset/mcp_service/chart/validation/runtime/__init__.py
Comment thread superset/mcp_service/chart/plugins/xy.py
Comment thread superset/mcp_service/chart/tool/generate_chart.py
Comment thread superset/mcp_service/app.py Outdated
Comment thread superset/mcp_service/chart/plugins/pie.py Outdated
aminghadersohi and others added 6 commits May 7, 2026 22:19
…generation

Replaces four scattered dispatch locations (schema_validator, dataset_validator,
chart_utils, runtime validator) with a central ChartTypePlugin registry. Each of
the 7 supported chart types (xy, table, pie, pivot_table, mixed_timeseries,
handlebars, big_number) now owns its pre-validation, column extraction, form_data
mapping, post-map validation, column normalization, and runtime warnings in a
single plugin class.

Key changes:
- Add ChartTypePlugin protocol and BaseChartPlugin base class (plugin.py)
- Add ChartTypeRegistry with register/get/all_types helpers (registry.py)
- Add 7 chart type plugins under chart/plugins/ with full coverage
- Fix 5-type column validation gap: pie, pivot_table, mixed_timeseries, handlebars,
  and big_number now participate in dataset column validation (previously silently skipped)
- Move BigNumber trendline temporal check to BigNumberChartPlugin.post_map_validate()
- Add get_runtime_warnings() to plugin protocol; XYChartPlugin implements
  format/cardinality checks, removing isinstance(config, XYChartConfig) from RuntimeValidator
- Fix stale generate_chart.py docstring listing only 'xy' and 'table' chart types
- Add missing pie, pivot_table, mixed_timeseries handlers to _enhance_validation_error;
  refactor into a data-driven lookup table to stay within complexity limits
- Fix empty details fallback in Pydantic error handler
Each ChartTypePlugin now declares:
- display_name: human-readable label for the chart_type discriminator
  (e.g. "Line / Bar / Area / Scatter Chart", "Pivot Table")
- native_viz_types: dict mapping every Superset-internal viz_type the
  plugin produces to a user-friendly name
  (e.g. {"echarts_timeseries_line": "Line Chart", "echarts_area": "Area Chart"})

The registry gains display_name_for_viz_type(viz_type) which searches
all plugins' native_viz_types maps, replacing the need for a separate
viz_type_display_names.json or viz_type_names.py module.

ChartInfo gains a chart_type_display_name field populated via the registry,
so list_charts / get_chart_info return human-readable chart type names.
The MCP system instructions now reference display names rather than
internal viz_type identifiers.
H1: Delete 7 dead _pre_validate_* static methods from SchemaValidator
    — exact duplicates of plugin pre_validate() methods, never called
    after _pre_validate_chart_type() was updated to delegate to plugin.

H2: Inline DatasetValidator._normalize_xy_config/_normalize_table_config
    into XYChartPlugin/TableChartPlugin.normalize_column_refs() and delete
    both DatasetValidator helper methods. The 5 other plugins already
    called _get_canonical_column_name directly; XY and Table now match.

H3: Add generate_name()/resolve_viz_type() to ChartTypePlugin protocol
    and BaseChartPlugin, implement in all 7 plugins. Replace the 7-arm
    isinstance chain in generate_chart_name() and the 7-arm elif chain
    in _resolve_viz_type() with single-line registry dispatch.

H4: Add a sync comment above _CHART_TYPE_ERROR_HINTS to document that
    it must stay in sync with the plugin registry.

M4: Move logger=logging.getLogger(__name__) from inside
    XYChartPlugin.get_runtime_warnings() to module scope.
…xes, test repairs

On top of the dead-code elimination in the previous commit:
- Add lazy _ensure_plugins_loaded() bootstrap to ChartTypeRegistry so the
  registry is populated even without importing app.py (fixes isolated test runs)
- Delegate _RegistryProxy methods to module-level functions so bootstrap runs
- Guard register() against empty chart_type strings
- Add generate_name + resolve_viz_type to ChartTypePlugin Protocol and
  BaseChartPlugin; delegate generate_chart_name/_resolve_viz_type in
  chart_utils to the plugin registry
- Add _with_context static helper to BaseChartPlugin (shared by all plugins)
- Fix stale 'five methods' → 'eight methods' docstring in plugin.py
- Add TypeVar _C to normalize_column_names so mypy infers correct return type
- Fix broken tests: update _pre_validate_big_number_config → _pre_validate_chart_type,
  remove deleted TestNormalizeXYConfig/TestNormalizeTableConfig classes,
  update runtime validator tests for removed _validate_format_compatibility /
  _validate_cardinality methods, add x is not None narrowing guards

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…nal corrections, cardinality suggestions

- Remove redundant local imports from BigNumberChartPlugin.post_map_validate()
  now that BigNumberChartConfig and is_column_truly_temporal are at top level
- Add explanatory comments on the two remaining local get_registry imports in
  chart_utils.py and dataset_validator.py (circular import prevention)
- Fix schema_validator.py and generate_chart.py docstring: XY 'x' field is
  optional (defaults to dataset primary datetime column), not required
- Propagate cardinality suggestions alongside warnings in XYChartPlugin
- Clarify app.py instructions: chart_type_display_name is null for viz_types
  outside the 7 generate_chart-supported types

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
All per-method local imports in the 7 chart plugins were moved to module-level.
None of them 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.

- big_number: add map_big_number_config, _big_number_chart_what,
  _summarize_filters, DatasetValidator to top-level imports
- pie: add map_pie_config, _pie_chart_what, _summarize_filters, PieChartConfig,
  DatasetValidator to top-level imports
- xy: add map_xy_config, _xy_chart_what/context, XYChartConfig, DatasetValidator,
  FormatTypeValidator, CardinalityValidator to top-level imports
- table: add map_table_config, _table_chart_what, _summarize_filters,
  TableChartConfig, DatasetValidator to top-level imports
- pivot_table: add map_pivot_table_config, _pivot_table_what, _summarize_filters,
  PivotTableChartConfig, DatasetValidator to top-level imports
- mixed_timeseries: add map_mixed_timeseries_config, _mixed_timeseries_what,
  _summarize_filters, MixedTimeseriesChartConfig, DatasetValidator to top-level
- handlebars: add map_handlebars_config, _handlebars_chart_what, _summarize_filters,
  HandlebarsChartConfig, DatasetValidator to top-level imports

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@aminghadersohi aminghadersohi force-pushed the mcp-chart-tools-rearch branch from 891676f to 1e2b541 Compare May 7, 2026 22:19
@codecov
Copy link
Copy Markdown

codecov Bot commented May 7, 2026

Codecov Report

❌ Patch coverage is 34.41781% with 383 lines in your changes missing coverage. Please review.
✅ Project coverage is 63.90%. Comparing base (ad5e317) to head (1e2b541).

Files with missing lines Patch % Lines
superset/mcp_service/chart/plugins/xy.py 29.33% 53 Missing ⚠️
...rset/mcp_service/chart/plugins/mixed_timeseries.py 26.98% 46 Missing ⚠️
superset/mcp_service/chart/plugins/big_number.py 29.03% 44 Missing ⚠️
superset/mcp_service/chart/plugins/handlebars.py 28.81% 42 Missing ⚠️
superset/mcp_service/chart/plugins/pivot_table.py 31.48% 37 Missing ⚠️
superset/mcp_service/chart/plugins/pie.py 38.63% 27 Missing ⚠️
superset/mcp_service/chart/registry.py 45.45% 22 Missing and 2 partials ⚠️
superset/mcp_service/chart/plugins/table.py 42.50% 23 Missing ⚠️
superset/mcp_service/chart/chart_utils.py 0.00% 20 Missing ⚠️
.../mcp_service/chart/validation/dataset_validator.py 14.28% 18 Missing ⚠️
... and 4 more
Additional details and impacted files
@@            Coverage Diff             @@
##           master   #39922      +/-   ##
==========================================
+ Coverage   63.88%   63.90%   +0.02%     
==========================================
  Files        2583     2593      +10     
  Lines      136618   136938     +320     
  Branches    31502    31513      +11     
==========================================
+ Hits        87272    87504     +232     
- Misses      47830    47916      +86     
- Partials     1516     1518       +2     
Flag Coverage Δ
hive 39.54% <34.41%> (+0.16%) ⬆️
mysql 59.10% <34.41%> (+0.06%) ⬆️
postgres 59.18% <34.41%> (+0.06%) ⬆️
presto 41.23% <34.41%> (+0.16%) ⬆️
python 60.62% <34.41%> (+0.06%) ⬆️
sqlite 58.82% <34.41%> (+0.07%) ⬆️
unit 100.00% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants