test: mock ChartDataCommand.run for the ChartDataQueryFailedError no-reraise test (SC-118140)

Automated PR review (bito-code-review) correctly flagged that this test
mocked the schema loader's side effect instead of ChartDataCommand.run,
so it never actually reached the run() call the exception is meant to
simulate failing at. The except clause still catches the exception either
way (same try block), so the assertion was never wrong, but mocking at
the real trigger point matches the sibling ChartDataCacheLoadError test
and the corrected integration test, and is more representative of the
actual failure path.

Co-Authored-By: Claude <noreply@anthropic.com>

superset/tasks/async_queries.py

@@ -106,6 +106,10 @@ def load_chart_data_into_cache(
 ) -> None:
     # pylint: disable=import-outside-toplevel
     from superset.commands.chart.data.get_data_command import ChartDataCommand
+    from superset.commands.chart.exceptions import (
+        ChartDataCacheLoadError,
+        ChartDataQueryFailedError,
+    )
 
     with override_user(_load_user_from_job_metadata(job_metadata), force=False):
         try:
@@ -123,6 +127,20 @@ def load_chart_data_into_cache(
         except SoftTimeLimitExceeded as ex:
             _handle_soft_time_limit(job_metadata, ex, "loading chart data")
             raise
+        except (ChartDataCacheLoadError, ChartDataQueryFailedError) as ex:
+            # These map to 422/400 in the synchronous chart/data endpoint (see
+            # ChartDataRestApi._get_data_response) - expected, client-facing
+            # validation failures (e.g. a chart still referencing columns a
+            # customer has since dropped from the dataset), not application
+            # bugs. The failure is already delivered to the client via
+            # update_job below; re-raising would only surface it a second
+            # time as an unhandled Celery task exception.
+            logger.info("Chart data query failed while loading into cache: %s", ex)
+            async_query_manager.update_job(
+                job_metadata,
+                async_query_manager.STATUS_ERROR,
+                errors=sanitize_error_dicts([{"message": str(ex.message)}]),
+            )
         except Exception as ex:
             # Extract SIP-40 style errors when available
             if isinstance(ex, SupersetErrorException):

tests/integration_tests/tasks/async_queries_tests.py

@@ -110,8 +110,11 @@ class TestAsyncQueries(SupersetTestCase):
             "status": "pending",
             "errors": [],
         }
-        with pytest.raises(ChartDataQueryFailedError):
-            load_chart_data_into_cache(job_metadata, query_context)
+        # ChartDataQueryFailedError mirrors the synchronous chart/data endpoint's
+        # 400 (see ChartDataRestApi._get_data_response) - an expected validation
+        # failure, not a bug, so the task reports it via update_job and does not
+        # re-raise (see superset/tasks/async_queries.py).
+        load_chart_data_into_cache(job_metadata, query_context)
 
         mock_run_command.assert_called_once_with(cache=True)
         errors = [{"message": "Error: foo"}]

tests/unit_tests/tasks/test_async_queries.py

@@ -21,7 +21,10 @@ import pytest
 from celery.exceptions import SoftTimeLimitExceeded
 from flask_babel import lazy_gettext as _
 
-from superset.commands.chart.exceptions import ChartDataQueryFailedError
+from superset.commands.chart.exceptions import (
+    ChartDataCacheLoadError,
+    ChartDataQueryFailedError,
+)
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
 from superset.exceptions import (
     OAuth2RedirectError,
@@ -43,7 +46,7 @@ def test_load_chart_data_into_cache_with_error(
     job_metadata = {"user_id": 1}
     form_data = {}
     err_message = "Something went wrong"
-    err = ChartDataQueryFailedError(_(err_message))
+    err = RuntimeError(err_message)
 
     mock_user = mock.MagicMock()
     mock_query_context_schema = mock.MagicMock()
@@ -54,7 +57,7 @@ def test_load_chart_data_into_cache_with_error(
 
     mock_query_context_schema.load.side_effect = err
 
-    with pytest.raises(ChartDataQueryFailedError):
+    with pytest.raises(RuntimeError):
         load_chart_data_into_cache(job_metadata, form_data)
 
     expected_errors = [{"message": err_message}]
@@ -64,6 +67,76 @@ def test_load_chart_data_into_cache_with_error(
     )
 
 
+@mock.patch("superset.tasks.async_queries.security_manager")
+@mock.patch("superset.tasks.async_queries.async_query_manager")
+@mock.patch("superset.commands.chart.data.get_data_command.ChartDataCommand")
+@mock.patch("superset.tasks.async_queries.ChartDataQueryContextSchema")
+def test_load_chart_data_into_cache_with_query_failed_error_does_not_reraise(
+    mock_query_context_schema_cls: mock.MagicMock,
+    mock_command_cls: mock.MagicMock,
+    mock_async_query_manager: mock.MagicMock,
+    mock_security_manager: mock.MagicMock,
+) -> None:
+    """
+    ChartDataQueryFailedError maps to a 400 in the synchronous chart/data
+    endpoint (see ChartDataRestApi._get_data_response) - an expected,
+    client-facing validation failure (e.g. a chart still referencing columns
+    a customer has since dropped from the dataset), not an application bug.
+    The task must still report it to the client via update_job, but must not
+    re-raise it - that would surface it a second time as an unhandled Celery
+    task exception.
+    """
+    from superset.tasks.async_queries import load_chart_data_into_cache
+
+    job_metadata = {"user_id": 1}
+    form_data: dict[str, Any] = {}
+    err_message = "Columns missing in dataset: ['foo']"
+
+    mock_security_manager.get_user_by_id.return_value = mock.MagicMock()
+    mock_async_query_manager.STATUS_ERROR = "error"
+    mock_query_context_schema_cls.return_value.load.return_value = mock.MagicMock()
+    mock_command_cls.return_value.run.side_effect = ChartDataQueryFailedError(
+        _(err_message)
+    )
+
+    # Should not raise.
+    load_chart_data_into_cache(job_metadata, form_data)
+
+    mock_async_query_manager.update_job.assert_called_once_with(
+        job_metadata, "error", errors=[{"message": err_message}]
+    )
+
+
+@mock.patch("superset.tasks.async_queries.security_manager")
+@mock.patch("superset.tasks.async_queries.async_query_manager")
+@mock.patch("superset.commands.chart.data.get_data_command.ChartDataCommand")
+@mock.patch("superset.tasks.async_queries.ChartDataQueryContextSchema")
+def test_load_chart_data_into_cache_with_cache_load_error_does_not_reraise(
+    mock_query_context_schema_cls: mock.MagicMock,
+    mock_command_cls: mock.MagicMock,
+    mock_async_query_manager: mock.MagicMock,
+    mock_security_manager: mock.MagicMock,
+) -> None:
+    """Same as above, for the sibling 422-mapped ChartDataCacheLoadError."""
+    from superset.tasks.async_queries import load_chart_data_into_cache
+
+    job_metadata = {"user_id": 1}
+    form_data: dict[str, Any] = {}
+    err_message = "Cache load failed"
+
+    mock_security_manager.get_user_by_id.return_value = mock.MagicMock()
+    mock_async_query_manager.STATUS_ERROR = "error"
+    mock_query_context_schema_cls.return_value.load.return_value = mock.MagicMock()
+    mock_command_cls.return_value.run.side_effect = ChartDataCacheLoadError(err_message)
+
+    # Should not raise.
+    load_chart_data_into_cache(job_metadata, form_data)
+
+    mock_async_query_manager.update_job.assert_called_once_with(
+        job_metadata, "error", errors=[{"message": err_message}]
+    )
+
+
 @mock.patch("superset.tasks.async_queries.security_manager")
 @mock.patch("superset.tasks.async_queries.async_query_manager")
 @mock.patch("superset.tasks.async_queries.ChartDataQueryContextSchema")