bito-code-review[bot] commented on code in PR #43573: URL: https://github.com/apache/superset/pull/43573#discussion_r4130332314
########## superset/mcp_service/chart/plugins/sankey.py: ########## @@ -0,0 +1,148 @@ +# 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. + +"""Sankey chart type plugin.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any, ClassVar + +from superset.mcp_service.chart.chart_utils import ( + _sankey_chart_what, + _summarize_filters, + map_sankey_config, +) +from superset.mcp_service.chart.plugin import BaseChartPlugin +from superset.mcp_service.chart.schemas import ColumnRef, SankeyChartConfig +from superset.mcp_service.chart.validation.dataset_validator import DatasetValidator +from superset.mcp_service.common.error_schemas import ChartGenerationError + + +class SankeyChartPlugin(BaseChartPlugin): + """Plugin for sankey chart type.""" + + chart_type = "sankey_v2" + display_name = "Sankey Diagram" + native_viz_types: ClassVar[Mapping[str, str]] = { + "sankey_v2": "Sankey Diagram", + } + + def pre_validate( + self, + config: dict[str, Any], + ) -> ChartGenerationError | None: + missing_fields = [] + + if "source" not in config: + missing_fields.append("'source' (origin node column)") + if "target" not in config: + missing_fields.append("'target' (destination node column)") + if "metric" not in config: + missing_fields.append("'metric' (edge weight)") + + if missing_fields: + return ChartGenerationError( + error_type="missing_sankey_fields", + message=( + f"Sankey chart missing required fields: {', '.join(missing_fields)}" + ), + details=( + "Sankey diagrams draw weighted flows from a source node to " + "a target node; the metric sets each edge's width" + ), + suggestions=[ + "Add 'source': {'name': 'from_stage'}", + "Add 'target': {'name': 'to_stage'}", + "Add 'metric': {'name': 'users', 'aggregate': 'SUM'}", + "Example: {'chart_type': 'sankey_v2', " + "'source': {'name': 'from_stage'}, " + "'target': {'name': 'to_stage'}, " + "'metric': {'name': 'users', 'aggregate': 'SUM'}}", + ], + error_code="MISSING_SANKEY_FIELDS", + ) + + return None + + def extract_column_refs(self, config: Any) -> list[ColumnRef]: + if not isinstance(config, SankeyChartConfig): + return [] + refs: list[ColumnRef] = [config.source, config.target, config.metric] + if config.filters: + for f in config.filters: + refs.append(ColumnRef(name=f.column)) + return refs + + def to_form_data( + self, config: Any, dataset_id: int | str | None = None + ) -> dict[str, Any]: + return map_sankey_config(config) + + def generate_name(self, config: Any, dataset_name: str | None = None) -> str: + what = _sankey_chart_what(config) + context = _summarize_filters(config.filters) + return self._with_context(what, context) + + def resolve_viz_type(self, config: Any) -> str: + return "sankey_v2" + + def normalize_column_refs(self, config: Any, dataset_context: Any) -> Any: + config_dict = config.model_dump(exclude_unset=True) + + for key in ("source", "target"): + col = config_dict.get(key) + if col and not col.get("sql_expression") and not col.get("saved_metric"): + col["name"] = DatasetValidator.get_canonical_column_name( + col["name"], dataset_context + ) + if config_dict.get("metric"): + if config_dict["metric"].get("sql_expression"): + pass + elif config_dict["metric"].get("saved_metric"): + config_dict["metric"]["name"] = ( + DatasetValidator.get_canonical_metric_name( + config_dict["metric"]["name"], dataset_context + ) + ) + else: + config_dict["metric"]["name"] = ( + DatasetValidator.get_canonical_column_name( + config_dict["metric"]["name"], dataset_context Review Comment: <div> <div id="suggestion"> <div id="issue"><b>Duplicated metric normalization</b></div> <div id="fix"> The metric-normalization chain (sql_expression → skip, saved_metric → `get_canonical_metric_name`, else `get_canonical_column_name`) is copied verbatim from `pie.py`/`treemap.py` and appears in 13 sibling plugins; `BaseChartPlugin` hosts no shared helper. Extract one helper (e.g. on `DatasetValidator` beside `get_canonical_*`) and call it here, so the next canonicalization fix lands once instead of 13 times. </div> </div> <small><i>Code Review Run #d00ad7</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_sankey_chart.py: ########## @@ -0,0 +1,270 @@ +# 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. + +"""Tests for the sankey chart type plugin. + +Schema validation, form_data mapping (matching the frontend Sankey buildQuery +contract for viz_type ``sankey_v2`` — a ``source`` and ``target`` column plus +one ``metric`` weighting each edge), and registry integration. +""" + +from typing import Any + +import pytest +from pydantic import TypeAdapter, ValidationError + +from superset.common.form_data_query_context import columns_from_form_data +from superset.mcp_service.chart.chart_helpers import resolve_groupby +from superset.mcp_service.chart.chart_utils import map_sankey_config +from superset.mcp_service.chart.schemas import ChartConfig, SankeyChartConfig + + +class TestSankeyChartConfigSchema: + """SankeyChartConfig schema validation.""" + + def test_basic_sankey_config(self) -> None: + config = SankeyChartConfig( + chart_type="sankey_v2", + source={"name": "from_stage"}, + target={"name": "to_stage"}, + metric={"name": "users", "aggregate": "SUM"}, + ) + assert config.source.name == "from_stage" + assert config.target.name == "to_stage" + assert config.sort_by_metric is True # shared control default Review Comment: <div> <div id="suggestion"> <div id="issue"><b>Unannotated test locals</b></div> <div id="fix"> Locals such as `cfg` (line 50), `config`, `form_data`, `queries`, and `orderby` are assigned without type annotations throughout the file. BITO rule 13153 asks for explicit annotations on test-file locals even when inferable; `Any` is already imported for this purpose. </div> </div> <small><i>Code Review Run #d00ad7</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_sankey_chart.py: ########## @@ -0,0 +1,270 @@ +# 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. + +"""Tests for the sankey chart type plugin. + +Schema validation, form_data mapping (matching the frontend Sankey buildQuery +contract for viz_type ``sankey_v2`` — a ``source`` and ``target`` column plus +one ``metric`` weighting each edge), and registry integration. +""" + +from typing import Any + +import pytest +from pydantic import TypeAdapter, ValidationError + +from superset.common.form_data_query_context import columns_from_form_data +from superset.mcp_service.chart.chart_helpers import resolve_groupby +from superset.mcp_service.chart.chart_utils import map_sankey_config +from superset.mcp_service.chart.schemas import ChartConfig, SankeyChartConfig + + +class TestSankeyChartConfigSchema: + """SankeyChartConfig schema validation.""" + + def test_basic_sankey_config(self) -> None: Review Comment: <div> <div id="suggestion"> <div id="issue"><b>Missing test docstrings</b></div> <div id="fix"> 16 of the 19 test functions (e.g. `test_basic_sankey_config`, `test_group_by_and_order_by`, `test_pre_validate_missing_fields`) have no docstring; only `test_sankey_source_rejects_aggregate` and the two query-builder tests do. BITO rule 12148 requires a docstring on every newly added test function documenting scenario and expected outcome. </div> </div> <small><i>Code Review Run #d00ad7</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_sankey_chart.py: ########## @@ -0,0 +1,270 @@ +# 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. + +"""Tests for the sankey chart type plugin. + +Schema validation, form_data mapping (matching the frontend Sankey buildQuery +contract for viz_type ``sankey_v2`` — a ``source`` and ``target`` column plus +one ``metric`` weighting each edge), and registry integration. +""" + +from typing import Any + +import pytest +from pydantic import TypeAdapter, ValidationError + +from superset.common.form_data_query_context import columns_from_form_data +from superset.mcp_service.chart.chart_helpers import resolve_groupby +from superset.mcp_service.chart.chart_utils import map_sankey_config +from superset.mcp_service.chart.schemas import ChartConfig, SankeyChartConfig + + +class TestSankeyChartConfigSchema: + """SankeyChartConfig schema validation.""" + + def test_basic_sankey_config(self) -> None: + config = SankeyChartConfig( + chart_type="sankey_v2", + source={"name": "from_stage"}, + target={"name": "to_stage"}, + metric={"name": "users", "aggregate": "SUM"}, + ) + assert config.source.name == "from_stage" + assert config.target.name == "to_stage" + assert config.sort_by_metric is True # shared control default + + @pytest.mark.parametrize("missing", ["source", "target", "metric"]) + def test_sankey_missing_required(self, missing: str) -> None: + cfg = { + "chart_type": "sankey_v2", + "source": {"name": "from_stage"}, + "target": {"name": "to_stage"}, + "metric": {"name": "users", "aggregate": "SUM"}, + } + del cfg[missing] + with pytest.raises(ValidationError): + SankeyChartConfig(**cfg) + + def test_sankey_rejects_extra_fields(self) -> None: + with pytest.raises(ValidationError): + SankeyChartConfig( + chart_type="sankey_v2", + source={"name": "from_stage"}, + target={"name": "to_stage"}, + metric={"name": "users", "aggregate": "SUM"}, + bogus=1, + ) + + def test_sankey_source_rejects_saved_metric(self) -> None: + with pytest.raises(ValidationError): + SankeyChartConfig( + chart_type="sankey_v2", + source={"name": "count", "saved_metric": True}, + target={"name": "to_stage"}, + metric={"name": "users", "aggregate": "SUM"}, + ) + + def test_sankey_target_rejects_saved_metric(self) -> None: + with pytest.raises(ValidationError): + SankeyChartConfig( + chart_type="sankey_v2", + source={"name": "from_stage"}, + target={"name": "count", "saved_metric": True}, + metric={"name": "users", "aggregate": "SUM"}, + ) + + def test_sankey_source_rejects_aggregate(self) -> None: + """An aggregate makes source metric-like; source is a node dimension.""" + with pytest.raises(ValidationError): + SankeyChartConfig( + chart_type="sankey_v2", + source={"name": "amount", "aggregate": "SUM"}, + target={"name": "to_stage"}, + metric={"name": "users", "aggregate": "SUM"}, + ) + + def test_sankey_target_rejects_aggregate(self) -> None: + with pytest.raises(ValidationError): + SankeyChartConfig( + chart_type="sankey_v2", + source={"name": "from_stage"}, + target={"name": "amount", "aggregate": "SUM"}, + metric={"name": "users", "aggregate": "SUM"}, + ) + + def test_chart_config_union_dispatches_sankey(self) -> None: + config = TypeAdapter(ChartConfig).validate_python( + { + "chart_type": "sankey_v2", + "source": {"name": "from_stage"}, + "target": {"name": "to_stage"}, + "metric": {"name": "users", "aggregate": "SUM"}, + } + ) + assert isinstance(config, SankeyChartConfig) + + +class TestMapSankeyConfig: + """form_data mapping must match the frontend Sankey buildQuery.""" + + def test_basic_sankey_form_data(self) -> None: + config = SankeyChartConfig( + chart_type="sankey_v2", + source={"name": "from_stage"}, + target={"name": "to_stage"}, + metric={"name": "users", "aggregate": "SUM"}, + ) + form_data = map_sankey_config(config) + assert form_data["viz_type"] == "sankey_v2" + assert form_data["source"] == "from_stage" + assert form_data["target"] == "to_stage" + assert form_data["groupby"] == ["from_stage", "to_stage"] + assert form_data["metric"]["label"] == "SUM(users)" + assert form_data["sort_by_metric"] is True + # orderby is applied by the query-dict builder, not stashed in form_data + # (a top-level form_data['orderby'] is a no-op on the MCP path) + assert "orderby" not in form_data + + def test_sankey_form_data_with_filters_and_no_sort(self) -> None: + config = SankeyChartConfig( + chart_type="sankey_v2", + source={"name": "from_stage"}, + target={"name": "to_stage"}, + metric={"name": "users", "aggregate": "SUM"}, + sort_by_metric=False, + filters=[{"column": "year", "op": "=", "value": 2026}], + ) + form_data = map_sankey_config(config) + assert form_data["sort_by_metric"] is False + assert "orderby" not in form_data # no metric ordering when unset + assert form_data["adhoc_filters"], "filters must map to adhoc_filters" + + def test_sankey_saved_metric_maps_to_name_string(self) -> None: + config = SankeyChartConfig( + chart_type="sankey_v2", + source={"name": "from_stage"}, + target={"name": "to_stage"}, + metric={"name": "flow_volume", "saved_metric": True}, + ) + assert map_sankey_config(config)["metric"] == "flow_volume" + + +class TestSankeySourceTargetReachQueryBuilders: + """The emitted form_data must group by the edge's node columns. + + ``source``/``target`` are the frontend Sankey buildQuery field names; the + backend query builders derive grouping columns from ``groupby`` and have no + alias for them (unlike ``entity``/``series``). Without an explicit + ``groupby`` the query aggregates the whole dataset into a single row instead + of one row per edge, so assert against the real consumers rather than the + emitted key alone. + """ + + @staticmethod + def _form_data() -> dict[str, Any]: + return map_sankey_config( + SankeyChartConfig( + chart_type="sankey_v2", + source={"name": "from_stage"}, + target={"name": "to_stage"}, + metric={"name": "users", "aggregate": "SUM"}, + ) + ) + + def test_resolve_groupby_returns_the_node_columns(self) -> None: + """The ``get_chart_data`` path groups by source and target.""" + assert resolve_groupby(self._form_data()) == ["from_stage", "to_stage"] + + def test_columns_from_form_data_returns_the_node_columns(self) -> None: + """The chart-preview path groups by source and target.""" + assert columns_from_form_data(self._form_data()) == [ + "from_stage", + "to_stage", + ] + + +class TestSankeyQueryContext: + """The built query must GROUP BY source+target and ORDER BY the metric. + + Both transforms live in the frontend buildQuery; the MCP path rebuilds the + query dict directly, so they must be re-derived or the query collapses to a + single unordered aggregate row. + """ + + def test_group_by_and_order_by(self, monkeypatch) -> None: Review Comment: <div> <div id="suggestion"> <div id="issue"><b>Unannotated fixture parameter</b></div> <div id="fix"> `monkeypatch` is a pytest fixture but is the only unannotated parameter in the file. BITO rule 12101 requires explicit annotations on fixture-injected parameters; `pytest.MonkeyPatch` is available in the repo's pytest 7.4.0. </div> </div> <small><i>Code Review Run #d00ad7</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 ########## superset/mcp_service/chart/schemas.py: ########## @@ -1673,6 +1673,64 @@ def record_implicit_metric_aggregate(self) -> "BubbleChartConfig": return self +class SankeyChartConfig(BaseChartConfig): + """Config for sankey charts (viz_type ``sankey_v2``). + + Matches the frontend Sankey buildQuery contract: a ``source`` and a + ``target`` column define the edges of the flow diagram and one ``metric`` + weights each edge. When ``sort_by_metric`` is set the query orders by the + metric descending. + """ + + model_config = ConfigDict(extra="ignore", populate_by_name=True) + + chart_type: Literal["sankey_v2"] = "sankey_v2" + source: ColumnRef = Field( + ..., + description="Column used as the source (origin) node of each edge", + ) + target: ColumnRef = Field( + ..., + description="Column used as the target (destination) node of each edge", + ) + metric: ColumnRef = Field( + ..., + description="Metric weighting each edge (use aggregate e.g. SUM, " + "COUNT for ad-hoc, or set saved_metric=True for a saved dataset metric)", + ) Review Comment: <div> <div id="suggestion"> <div id="issue"><b>Missing implicit-aggregate recording</b></div> <div id="fix"> Unlike `BubbleChartConfig.record_implicit_metric_aggregate`, no validator records an implicit SUM for a bare `metric` ref. `create_metric_object` defaults to SUM at map time, while `DatasetValidator._validate_aggregations` skips refs without `aggregate` — so SUM over a text column fails in the database instead of with a clear validation message. Mirror the Bubble validator. </div> </div> <small><i>Code Review Run #d00ad7</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 ########## superset/mcp_service/chart/schemas.py: ########## @@ -1673,6 +1673,64 @@ def record_implicit_metric_aggregate(self) -> "BubbleChartConfig": return self +class SankeyChartConfig(BaseChartConfig): + """Config for sankey charts (viz_type ``sankey_v2``). + + Matches the frontend Sankey buildQuery contract: a ``source`` and a + ``target`` column define the edges of the flow diagram and one ``metric`` + weights each edge. When ``sort_by_metric`` is set the query orders by the + metric descending. + """ + + model_config = ConfigDict(extra="ignore", populate_by_name=True) + + chart_type: Literal["sankey_v2"] = "sankey_v2" + source: ColumnRef = Field( + ..., + description="Column used as the source (origin) node of each edge", + ) + target: ColumnRef = Field( + ..., + description="Column used as the target (destination) node of each edge", + ) + metric: ColumnRef = Field( + ..., + description="Metric weighting each edge (use aggregate e.g. SUM, " + "COUNT for ad-hoc, or set saved_metric=True for a saved dataset metric)", + ) + sort_by_metric: bool = Field( + True, + description="Order edges by the metric descending (frontend default)", + ) + row_limit: int = Field(10000, description="Max edges queried", ge=1, le=100000) + filters: List[FilterConfig] | None = Field( + None, + description="Structured filters (column/op/value). " + "Do NOT use adhoc_filters or raw SQL expressions.", + ) + color_scheme: str | None = Field( + None, + description=( + "Superset color scheme ID (e.g. 'supersetColors', 'lyftColors', " + "'googleCategory10c', 'd3Category10'). Defaults to 'supersetColors'." + ), + max_length=100, + ) + + @model_validator(mode="after") + def reject_metric_style_nodes(self) -> "SankeyChartConfig": + """source and target are node dimensions, not metrics.""" + for col, name in ((self.source, "source"), (self.target, "target")): + _reject_sql_expression_on_dimension(col, name) + if col and col.is_metric: Review Comment: <div> <div id="suggestion"> <div id="issue"><b>Dead truthiness guard</b></div> <div id="fix"> `source`/`target` are required, non-optional `ColumnRef` fields, and Pydantic v2 `BaseModel` instances are always truthy (no `__bool__`/`__len__`), so `col` in `col and col.is_metric` is always true — a dead guard. Sibling `BubbleChartConfig.reject_metric_style_dimensions` omits it; drop the `col and` conjunct for consistency. </div> </div> <small><i>Code Review Run #d00ad7</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]
