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

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"