diff --git a/superset/commands/chart/exceptions.py b/superset/commands/chart/exceptions.py index 72ef71f466e8..03c7a0cbdaa5 100644 --- a/superset/commands/chart/exceptions.py +++ b/superset/commands/chart/exceptions.py @@ -103,6 +103,19 @@ def __init__(self) -> None: ) +class ChartQueryContextDatasourceMismatchValidationError(ValidationError): + """ + Raised when a query-context-only update carries a datasource that does not + match the chart's own datasource. + """ + + def __init__(self) -> None: + super().__init__( + _("The query context datasource does not match the chart datasource"), + field_name="query_context", + ) + + class ChartNotFoundError(CommandException): message = "Chart not found." diff --git a/superset/commands/chart/update.py b/superset/commands/chart/update.py index e1a86732d3d6..c75192efa11e 100644 --- a/superset/commands/chart/update.py +++ b/superset/commands/chart/update.py @@ -29,6 +29,7 @@ ChartForbiddenError, ChartInvalidError, ChartNotFoundError, + ChartQueryContextDatasourceMismatchValidationError, ChartUpdateFailedError, DashboardsForbiddenError, DashboardsNotFoundValidationError, @@ -41,6 +42,7 @@ from superset.models.dashboard import Dashboard from superset.models.slice import Slice from superset.tags.models import ObjectType +from superset.utils import json from superset.utils.decorators import on_error, transaction logger = logging.getLogger(__name__) @@ -101,6 +103,52 @@ def _validate_new_dashboard_access( if not security_manager.is_owner(dash): raise DashboardsForbiddenError() + def _validate_query_context_datasource( + self, exceptions: list[ValidationError] + ) -> None: + """ + Ensure a query-context-only update keeps the chart's own datasource. + + The submitted query context is only verified when it carries a parseable + ``datasource`` object; a payload that references a different datasource than + the chart's persisted one is rejected. Payloads without a datasource fall + back to the chart's datasource at execution time and need no check. + """ + if not self._model: + return + + raw_query_context = self._properties.get("query_context") + if not raw_query_context: + return + + try: + query_context = json.loads(raw_query_context) + except (TypeError, ValueError): + # An unparseable payload cannot be verified or replayed; leave it for + # downstream handling rather than guessing at its intent. + return + + datasource = ( + query_context.get("datasource") if isinstance(query_context, dict) else None + ) + if not isinstance(datasource, dict): + return + + try: + ids_match = int(datasource["id"]) == self._model.datasource_id + except (KeyError, TypeError, ValueError): + ids_match = False + + # A datasource object must carry a type that matches the chart's own. + # Treating a missing type as valid would let an id-only payload through, + # and query-context loading reads datasource["type"] directly, so that + # payload raises KeyError when the saved context is later replayed. + datasource_type = datasource.get("type") + types_match = str(datasource_type) == self._model.datasource_type + + if not ids_match or not types_match: + exceptions.append(ChartQueryContextDatasourceMismatchValidationError()) + def validate(self) -> None: # noqa: C901 exceptions: list[ValidationError] = [] dashboard_ids = self._properties.get("dashboards") @@ -144,6 +192,9 @@ def validate(self) -> None: # noqa: C901 security_manager.raise_for_access(chart=self._model) except SupersetSecurityException as ex: raise ChartForbiddenError() from ex + # Keep the refreshed payload bound to the chart's own datasource so it + # cannot be repointed at an unrelated one. + self._validate_query_context_datasource(exceptions) # validate tags try: diff --git a/tests/unit_tests/commands/chart/update_test.py b/tests/unit_tests/commands/chart/update_test.py index eb678c718b0b..d27c1089c0b2 100644 --- a/tests/unit_tests/commands/chart/update_test.py +++ b/tests/unit_tests/commands/chart/update_test.py @@ -17,10 +17,11 @@ import pytest from pytest_mock import MockerFixture -from superset.commands.chart.exceptions import ChartForbiddenError +from superset.commands.chart.exceptions import ChartForbiddenError, ChartInvalidError from superset.commands.chart.update import UpdateChartCommand from superset.errors import ErrorLevel, SupersetError, SupersetErrorType from superset.exceptions import SupersetSecurityException +from superset.utils import json def _ownership_exc() -> SupersetSecurityException: @@ -151,3 +152,78 @@ def test_update_chart_owner_can_perform_regular_update( find_by_id.assert_called_once_with(1) raise_for_ownership.assert_called_once() + + +def _query_context_payload(datasource: object) -> dict[str, object]: + """Build a query-context-only update payload targeting ``datasource``.""" + return { + "query_context": json.dumps({"datasource": datasource, "queries": []}), + "query_context_generation": True, + } + + +def test_update_chart_query_context_matching_datasource_is_allowed( + mocker: MockerFixture, +) -> None: + """A query context that targets the chart's own datasource is accepted.""" + find_by_id = mocker.patch("superset.commands.chart.update.ChartDAO.find_by_id") + find_by_id.return_value = mocker.MagicMock( + id=1, tags=[], dashboards=[], datasource_id=42, datasource_type="table" + ) + mocker.patch("superset.commands.chart.update.security_manager.raise_for_ownership") + mocker.patch("superset.commands.chart.update.security_manager.raise_for_access") + + UpdateChartCommand( + 1, _query_context_payload({"id": 42, "type": "table"}) + ).validate() + + +@pytest.mark.parametrize( + "datasource", + [ + {"id": 99, "type": "table"}, # different id + {"id": 42, "type": "query"}, # different type + {"id": "99", "type": "table"}, # different id as string + {"id": 42}, # matching id but missing type + ], +) +def test_update_chart_query_context_mismatched_datasource_is_rejected( + mocker: MockerFixture, + datasource: dict[str, object], +) -> None: + """A query context pointing at a different datasource is rejected with a 4xx.""" + find_by_id = mocker.patch("superset.commands.chart.update.ChartDAO.find_by_id") + find_by_id.return_value = mocker.MagicMock( + id=1, tags=[], dashboards=[], datasource_id=42, datasource_type="table" + ) + mocker.patch("superset.commands.chart.update.security_manager.raise_for_ownership") + mocker.patch("superset.commands.chart.update.security_manager.raise_for_access") + + with pytest.raises(ChartInvalidError): + UpdateChartCommand(1, _query_context_payload(datasource)).validate() + + +@pytest.mark.parametrize( + "query_context", + [ + "{}", # no datasource key + '{"datasource": null}', # null datasource + "not-json", # unparseable payload + ], +) +def test_update_chart_query_context_without_datasource_is_allowed( + mocker: MockerFixture, + query_context: str, +) -> None: + """Payloads with no verifiable datasource fall back to the chart's own.""" + find_by_id = mocker.patch("superset.commands.chart.update.ChartDAO.find_by_id") + find_by_id.return_value = mocker.MagicMock( + id=1, tags=[], dashboards=[], datasource_id=42, datasource_type="table" + ) + mocker.patch("superset.commands.chart.update.security_manager.raise_for_ownership") + mocker.patch("superset.commands.chart.update.security_manager.raise_for_access") + + UpdateChartCommand( + 1, + {"query_context": query_context, "query_context_generation": True}, + ).validate()