diff --git a/superset/charts/data/commands/get_data_command.py b/superset/charts/data/commands/get_data_command.py index a84870a1dd306..c791ace9de3ee 100644 --- a/superset/charts/data/commands/get_data_command.py +++ b/superset/charts/data/commands/get_data_command.py @@ -17,7 +17,7 @@ import logging from typing import Any -from flask_babel import lazy_gettext as _ +from flask_babel import gettext as _ from superset.charts.commands.exceptions import ( ChartDataCacheLoadError, diff --git a/superset/common/query_context_processor.py b/superset/common/query_context_processor.py index 58e9022736730..56f53cc16a983 100644 --- a/superset/common/query_context_processor.py +++ b/superset/common/query_context_processor.py @@ -23,7 +23,7 @@ import numpy as np import pandas as pd -from flask_babel import _ +from flask_babel import gettext as _ from pandas import DateOffset from typing_extensions import TypedDict diff --git a/superset/tasks/async_queries.py b/superset/tasks/async_queries.py index cfcb3e31c6ba8..8a8b52e7d5231 100644 --- a/superset/tasks/async_queries.py +++ b/superset/tasks/async_queries.py @@ -90,10 +90,10 @@ def load_chart_data_into_cache( raise ex except Exception as ex: # TODO: QueryContext should support SIP-40 style errors - error = ( + error = str( ex.message # pylint: disable=no-member if hasattr(ex, "message") - else str(ex) + else ex ) errors = [{"message": error}] async_query_manager.update_job( diff --git a/tests/unit_tests/tasks/test_async_queries.py b/tests/unit_tests/tasks/test_async_queries.py new file mode 100644 index 0000000000000..5787bbdc8bb6b --- /dev/null +++ b/tests/unit_tests/tasks/test_async_queries.py @@ -0,0 +1,39 @@ +from unittest import mock + +import pytest +from flask_babel import lazy_gettext as _ + +from superset.charts.commands.exceptions import ChartDataQueryFailedError + + +@mock.patch("superset.tasks.async_queries.security_manager") +@mock.patch("superset.tasks.async_queries.async_query_manager") +@mock.patch("superset.tasks.async_queries.ChartDataQueryContextSchema") +def test_load_chart_data_into_cache_with_error( + mock_query_context_schema_cls, mock_async_query_manager, mock_security_manager +): + """Test that the task is gracefully marked failed in event of error""" + from superset.tasks.async_queries import load_chart_data_into_cache + + job_metadata = {"user_id": 1} + form_data = {} + err_message = "Something went wrong" + err = ChartDataQueryFailedError(_(err_message)) + + mock_user = mock.MagicMock() + mock_query_context_schema = mock.MagicMock() + + mock_security_manager.get_user_by_id.return_value = mock_user + mock_async_query_manager.STATUS_ERROR = "error" + mock_query_context_schema_cls.return_value = mock_query_context_schema + + mock_query_context_schema.load.side_effect = err + + with pytest.raises(ChartDataQueryFailedError): + load_chart_data_into_cache(job_metadata, form_data) + + expected_errors = [{"message": err_message}] + + mock_async_query_manager.update_job.assert_called_once_with( + job_metadata, "error", errors=expected_errors + )