fix(OAuth2): Support OAuth2 exception with legacy endpoint (#39897)

superset/tasks/async_queries.py

@@ -183,6 +183,11 @@ def load_explore_json_into_cache(  # pylint: disable=too-many-locals
         except Exception as ex:
             if isinstance(ex, SupersetVizException):
                 errors = ex.errors
+            # Extract SIP-40 style errors when available
+            elif isinstance(ex, SupersetErrorException):
+                errors = [dataclasses.asdict(ex.error)]  # type: ignore
+            elif isinstance(ex, SupersetErrorsException):
+                errors = [dataclasses.asdict(error) for error in ex.errors]  # type: ignore
             else:
                 error = ex.message if hasattr(ex, "message") else str(ex)
                 errors = [error]  # type: ignore

superset/views/core.py

@@ -68,6 +68,7 @@ from superset.daos.datasource import DatasourceDAO
 from superset.dashboards.permalink.exceptions import DashboardPermalinkGetFailedError
 from superset.exceptions import (
     CacheLoadError,
+    SupersetErrorException,
     SupersetException,
     SupersetSecurityException,
 )
@@ -266,6 +267,10 @@ class Superset(BaseSupersetView):
             )
 
             return self.generate_json(viz_obj, response_type)
+        except SupersetErrorException:
+            # Let structured Superset errors (e.g. OAuth2RedirectError) propagate
+            # so the global Flask error handler serializes them.
+            raise
         except SupersetException as ex:
             return json_error_response(utils.error_msg_from_exception(ex), 400)
 
@@ -290,7 +295,7 @@ class Superset(BaseSupersetView):
     @etag_cache()
     @check_resource_permissions(check_datasource_perms)
     @deprecated(eol_version="5.0.0")
-    def explore_json(
+    def explore_json(  # noqa: C901
         self, datasource_type: str | None = None, datasource_id: int | None = None
     ) -> FlaskResponse:
         """Serves all request that GET or POST form_data
@@ -377,6 +382,10 @@ class Superset(BaseSupersetView):
             )
 
             return self.generate_json(viz_obj, response_type)
+        except SupersetErrorException:
+            # Let structured Superset errors (e.g. OAuth2RedirectError) propagate
+            # so the global Flask error handler serializes them.
+            raise
         except SupersetException as ex:
             return json_error_response(utils.error_msg_from_exception(ex), 400)
 

superset/viz.py

@@ -51,6 +51,7 @@ from superset.exceptions import (
     NullValueException,
     QueryObjectValidationError,
     SpatialException,
+    SupersetErrorException,
     SupersetSecurityException,
 )
 from superset.extensions import cache_manager, security_manager
@@ -612,6 +613,10 @@ class BaseViz:  # pylint: disable=too-many-public-methods
                 )
                 self.errors.append(error)
                 self.status = QueryStatus.FAILED
+            except SupersetErrorException:
+                # Let structured Superset errors (e.g. OAuth2RedirectError) propagate
+                # so the global Flask error handler serializes them.
+                raise
             except Exception as ex:  # pylint: disable=broad-except
                 logger.exception(ex)
 

tests/integration_tests/core_tests.py

@@ -41,7 +41,7 @@ from superset.common.db_query_status import QueryStatus
 from superset.connectors.sqla.models import SqlaTable
 from superset.db_engine_specs.base import BaseEngineSpec
 from superset.db_engine_specs.mssql import MssqlEngineSpec
-from superset.exceptions import SupersetException
+from superset.exceptions import OAuth2RedirectError, SupersetException
 from superset.extensions import cache_manager
 from superset.models import core as models
 from superset.models.dashboard import Dashboard
@@ -458,6 +458,59 @@ class TestCore(SupersetTestCase):
         assert rv.status_code == 404
         assert data["error"] == "Cached data not found"
 
+    @pytest.mark.usefixtures("load_world_bank_dashboard_with_slices")
+    @mock.patch("superset.viz.BaseViz.get_df")
+    def test_explore_json_propagates_oauth2_redirect_error(
+        self, mock_get_df: mock.Mock
+    ) -> None:
+        """
+        SupersetErrorException exceptions bubble up properly.
+        """
+        mock_get_df.side_effect = OAuth2RedirectError(
+            url="https://accounts.example.com/o/oauth2/v2/auth?...",
+            tab_id="tab-123",
+            redirect_uri="https://superset.example.com/oauth2/redirect",
+        )
+
+        self.login(ADMIN_USERNAME)
+        slc = self.get_slice("Life Expectancy VS Rural %")
+        rv = self.client.post(
+            f"/superset/explore_json/{slc.datasource_type}/{slc.datasource_id}/",
+            data={"form_data": json.dumps(slc.form_data)},
+        )
+        data = json.loads(rv.data.decode("utf-8"))
+
+        assert "errors" in data, data
+        assert data["errors"][0]["error_type"] == "OAUTH2_REDIRECT"
+        assert data["errors"][0]["extra"] == {
+            "url": "https://accounts.example.com/o/oauth2/v2/auth?...",
+            "tab_id": "tab-123",
+            "redirect_uri": "https://superset.example.com/oauth2/redirect",
+        }
+
+    @pytest.mark.usefixtures("load_world_bank_dashboard_with_slices")
+    @mock.patch("superset.viz.BaseViz.get_df")
+    def test_explore_json_generic_exception_still_returns_viz_get_df_error(
+        self, mock_get_df: mock.Mock
+    ) -> None:
+        """
+        Non-Superset exceptions raised by ``get_df`` are reported as the
+        generic ``VIZ_GET_DF_ERROR``.
+        """
+        mock_get_df.side_effect = RuntimeError("boom")
+
+        self.login(ADMIN_USERNAME)
+        slc = self.get_slice("Life Expectancy VS Rural %")
+        rv = self.client.post(
+            f"/superset/explore_json/{slc.datasource_type}/{slc.datasource_id}/",
+            data={"form_data": json.dumps(slc.form_data)},
+        )
+        data = json.loads(rv.data.decode("utf-8"))
+
+        assert "errors" in data, data
+        assert data["errors"][0]["error_type"] == "VIZ_GET_DF_ERROR"
+        assert data["errors"][0]["message"] == "boom"
+
     def test_results_default_deserialization(self):
         use_new_deserialization = False
         data = [("a", 4, 4.0, "2019-08-18T16:39:16.660000")]

tests/unit_tests/tasks/test_async_queries.py

@@ -21,7 +21,12 @@ from flask_babel import lazy_gettext as _
 
 from superset.commands.chart.exceptions import ChartDataQueryFailedError
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
-from superset.exceptions import SupersetErrorException, SupersetErrorsException
+from superset.exceptions import (
+    OAuth2RedirectError,
+    SupersetErrorException,
+    SupersetErrorsException,
+    SupersetVizException,
+)
 
 
 @mock.patch("superset.tasks.async_queries.security_manager")
@@ -149,3 +154,170 @@ def test_load_chart_data_into_cache_with_superset_errors_exception(
     assert errors[1]["message"] == "Table not found"
     assert errors[1]["error_type"] == SupersetErrorType.TABLE_DOES_NOT_EXIST_ERROR
     assert errors[1]["level"] == ErrorLevel.WARNING
+
+
+@mock.patch("superset.tasks.async_queries.security_manager")
+@mock.patch("superset.tasks.async_queries.async_query_manager")
+@mock.patch("superset.tasks.async_queries.get_viz")
+@mock.patch("superset.tasks.async_queries.get_datasource_info")
+def test_load_explore_json_into_cache_preserves_oauth2_redirect_error(
+    mock_get_datasource_info,
+    mock_get_viz,
+    mock_async_query_manager,
+    mock_security_manager,
+):
+    """
+    OAuth2RedirectError raised by ``viz_obj.get_payload`` must reach the async
+    job's errors list as a structured SIP-40 envelope so the frontend can
+    render the OAuth2 banner identically to the sync legacy path.
+    """
+    from superset.tasks.async_queries import load_explore_json_into_cache
+
+    job_metadata = {"user_id": 1}
+    form_data: dict = {}
+
+    mock_get_datasource_info.return_value = (1, "table")
+    mock_security_manager.get_user_by_id.return_value = mock.MagicMock()
+    mock_async_query_manager.STATUS_ERROR = "error"
+
+    viz_obj = mock.MagicMock()
+    viz_obj.get_payload.side_effect = OAuth2RedirectError(
+        url="https://accounts.example.com/o/oauth2/v2/auth?...",
+        tab_id="tab-123",
+        redirect_uri="https://superset.example.com/oauth2/redirect",
+    )
+    mock_get_viz.return_value = viz_obj
+
+    with pytest.raises(OAuth2RedirectError):
+        load_explore_json_into_cache(job_metadata, form_data)
+
+    call_args = mock_async_query_manager.update_job.call_args
+    assert call_args[0] == (job_metadata, "error")
+    errors = call_args[1]["errors"]
+    assert len(errors) == 1
+    assert errors[0]["error_type"] == SupersetErrorType.OAUTH2_REDIRECT
+    assert errors[0]["extra"] == {
+        "url": "https://accounts.example.com/o/oauth2/v2/auth?...",
+        "tab_id": "tab-123",
+        "redirect_uri": "https://superset.example.com/oauth2/redirect",
+    }
+
+
+@mock.patch("superset.tasks.async_queries.security_manager")
+@mock.patch("superset.tasks.async_queries.async_query_manager")
+@mock.patch("superset.tasks.async_queries.get_viz")
+@mock.patch("superset.tasks.async_queries.get_datasource_info")
+def test_load_explore_json_into_cache_preserves_superset_errors_exception(
+    mock_get_datasource_info,
+    mock_get_viz,
+    mock_async_query_manager,
+    mock_security_manager,
+):
+    """SupersetErrorsException must be preserved as a list of SIP-40 dicts."""
+    from superset.tasks.async_queries import load_explore_json_into_cache
+
+    job_metadata = {"user_id": 1}
+    form_data: dict = {}
+
+    mock_get_datasource_info.return_value = (1, "table")
+    mock_security_manager.get_user_by_id.return_value = mock.MagicMock()
+    mock_async_query_manager.STATUS_ERROR = "error"
+
+    viz_obj = mock.MagicMock()
+    viz_obj.get_payload.side_effect = SupersetErrorsException(
+        [
+            SupersetError(
+                message="Column not found",
+                error_type=SupersetErrorType.COLUMN_DOES_NOT_EXIST_ERROR,
+                level=ErrorLevel.ERROR,
+            ),
+            SupersetError(
+                message="Table not found",
+                error_type=SupersetErrorType.TABLE_DOES_NOT_EXIST_ERROR,
+                level=ErrorLevel.WARNING,
+            ),
+        ]
+    )
+    mock_get_viz.return_value = viz_obj
+
+    with pytest.raises(SupersetErrorsException):
+        load_explore_json_into_cache(job_metadata, form_data)
+
+    errors = mock_async_query_manager.update_job.call_args[1]["errors"]
+    assert len(errors) == 2
+    assert errors[0]["error_type"] == SupersetErrorType.COLUMN_DOES_NOT_EXIST_ERROR
+    assert errors[1]["error_type"] == SupersetErrorType.TABLE_DOES_NOT_EXIST_ERROR
+
+
+@mock.patch("superset.tasks.async_queries.security_manager")
+@mock.patch("superset.tasks.async_queries.async_query_manager")
+@mock.patch("superset.tasks.async_queries.get_viz")
+@mock.patch("superset.tasks.async_queries.get_datasource_info")
+def test_load_explore_json_into_cache_preserves_superset_viz_exception(
+    mock_get_datasource_info,
+    mock_get_viz,
+    mock_async_query_manager,
+    mock_security_manager,
+):
+    """
+    Test that SupersetVizException passes ``ex.errors`` straight through.
+    """
+    from superset.tasks.async_queries import load_explore_json_into_cache
+
+    job_metadata = {"user_id": 1}
+    form_data: dict = {}
+
+    mock_get_datasource_info.return_value = (1, "table")
+    mock_security_manager.get_user_by_id.return_value = mock.MagicMock()
+    mock_async_query_manager.STATUS_ERROR = "error"
+
+    payload_errors = [
+        {
+            "message": "Bad column",
+            "error_type": SupersetErrorType.VIZ_GET_DF_ERROR,
+            "level": ErrorLevel.ERROR,
+        }
+    ]
+    viz_obj = mock.MagicMock()
+    viz_obj.get_payload.return_value = {"errors": payload_errors}
+    viz_obj.has_error.return_value = True
+    mock_get_viz.return_value = viz_obj
+
+    with pytest.raises(SupersetVizException):
+        load_explore_json_into_cache(job_metadata, form_data)
+
+    errors = mock_async_query_manager.update_job.call_args[1]["errors"]
+    assert errors == payload_errors
+
+
+@mock.patch("superset.tasks.async_queries.security_manager")
+@mock.patch("superset.tasks.async_queries.async_query_manager")
+@mock.patch("superset.tasks.async_queries.get_viz")
+@mock.patch("superset.tasks.async_queries.get_datasource_info")
+def test_load_explore_json_into_cache_falls_back_to_string_for_generic_exception(
+    mock_get_datasource_info,
+    mock_get_viz,
+    mock_async_query_manager,
+    mock_security_manager,
+):
+    """
+    Test that Non-Superset exception are passed as plain-string error.
+    """
+    from superset.tasks.async_queries import load_explore_json_into_cache
+
+    job_metadata = {"user_id": 1}
+    form_data: dict = {}
+
+    mock_get_datasource_info.return_value = (1, "table")
+    mock_security_manager.get_user_by_id.return_value = mock.MagicMock()
+    mock_async_query_manager.STATUS_ERROR = "error"
+
+    viz_obj = mock.MagicMock()
+    viz_obj.get_payload.side_effect = RuntimeError("boom")
+    mock_get_viz.return_value = viz_obj
+
+    with pytest.raises(RuntimeError):
+        load_explore_json_into_cache(job_metadata, form_data)
+
+    errors = mock_async_query_manager.update_job.call_args[1]["errors"]
+    assert errors == ["boom"]

tests/unit_tests/test_viz_get_df_payload.py

@@ -0,0 +1,111 @@
+# 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.
+from typing import Any
+from unittest.mock import patch
+
+import pytest
+
+from superset import viz
+from superset.common.db_query_status import QueryStatus
+from superset.connectors.sqla.models import SqlaTable
+from superset.errors import SupersetErrorType
+from superset.exceptions import (
+    OAuth2RedirectError,
+    QueryObjectValidationError,
+)
+from superset.models.core import Database
+
+QUERY_OBJ: dict[str, Any] = {"row_limit": 100, "from_dttm": None, "to_dttm": None}
+
+
+def _viz() -> viz.BaseViz:
+    database = Database(database_name="d", sqlalchemy_uri="sqlite://")
+    datasource = SqlaTable(
+        table_name="t",
+        columns=[],
+        metrics=[],
+        main_dttm_col=None,
+        database=database,
+    )
+    # ``force=True`` skips the data cache lookup branch so ``get_df`` is always
+    # invoked, which is what we want to assert error-handling against.
+    return viz.BaseViz(
+        datasource=datasource,
+        form_data={"viz_type": "table"},
+        force=True,
+    )
+
+
+def test_get_df_payload_propagates_oauth2_redirect_error() -> None:
+    """
+    OAuth2RedirectError (a SupersetErrorException) must propagate out of
+    ``get_df_payload`` so the global Flask error handler can serialize it.
+    """
+    obj = _viz()
+    oauth_exc = OAuth2RedirectError(
+        url="https://accounts.example.com/o/oauth2/v2/auth?...",
+        tab_id="tab-123",
+        redirect_uri="https://superset.example.com/oauth2/redirect",
+    )
+
+    with patch.object(viz.BaseViz, "get_df", side_effect=oauth_exc):
+        with pytest.raises(OAuth2RedirectError) as exc_info:
+            obj.get_df_payload(QUERY_OBJ)
+
+    assert exc_info.value.error.error_type == SupersetErrorType.OAUTH2_REDIRECT
+    assert exc_info.value.error.extra == {
+        "url": "https://accounts.example.com/o/oauth2/v2/auth?...",
+        "tab_id": "tab-123",
+        "redirect_uri": "https://superset.example.com/oauth2/redirect",
+    }
+
+
+def test_get_df_payload_captures_generic_exception_as_viz_get_df_error() -> None:
+    """
+    Non-Superset exception raised by ``get_df`` are downgraded to a
+    ``VIZ_GET_DF_ERROR`` entry on ``self.errors``.
+    """
+    obj = _viz()
+
+    with patch.object(viz.BaseViz, "get_df", side_effect=RuntimeError("boom")):
+        payload = obj.get_df_payload(QUERY_OBJ)
+
+    assert obj.status == QueryStatus.FAILED
+    assert payload["status"] == QueryStatus.FAILED
+    assert len(obj.errors) == 1
+    assert obj.errors[0]["error_type"] == SupersetErrorType.VIZ_GET_DF_ERROR
+    assert obj.errors[0]["message"] == "boom"
+
+
+def test_get_df_payload_captures_query_object_validation_error() -> None:
+    """
+    ``QueryObjectValidationError`` is reported as ``VIZ_GET_DF_ERROR``.
+    """
+    obj = _viz()
+
+    with patch.object(
+        viz.BaseViz,
+        "get_df",
+        side_effect=QueryObjectValidationError("bad query"),
+    ):
+        payload = obj.get_df_payload(QUERY_OBJ)
+
+    assert obj.status == QueryStatus.FAILED
+    assert payload["status"] == QueryStatus.FAILED
+    assert len(obj.errors) == 1
+    assert obj.errors[0]["error_type"] == SupersetErrorType.VIZ_GET_DF_ERROR
+    assert obj.errors[0]["message"] == "bad query"