diff --git a/superset/utils/decorators.py b/superset/utils/decorators.py index 43b6909caf..7e34b98360 100644 --- a/superset/utils/decorators.py +++ b/superset/utils/decorators.py @@ -191,3 +191,22 @@ def debounce(duration: float | int = 0.1) -> Callable[..., Any]: def on_security_exception(self: Any, ex: Exception) -> Response: return self.response(403, **{"message": utils.error_msg_from_exception(ex)}) + + +@contextmanager +def suppress_logging( + logger_name: str | None = None, + new_level: int = logging.CRITICAL, +) -> Iterator[None]: + """ + Context manager to suppress logging during the execution of code block. + + Use with caution and make sure you have the least amount of code inside it. + """ + target_logger = logging.getLogger(logger_name) + original_level = target_logger.getEffectiveLevel() + target_logger.setLevel(new_level) + try: + yield + finally: + target_logger.setLevel(original_level) diff --git a/superset/utils/pandas_postprocessing/prophet.py b/superset/utils/pandas_postprocessing/prophet.py index 47d956fed5..a1c823ee2a 100644 --- a/superset/utils/pandas_postprocessing/prophet.py +++ b/superset/utils/pandas_postprocessing/prophet.py @@ -23,6 +23,7 @@ from pandas import DataFrame from superset.exceptions import InvalidPostProcessingError from superset.utils.core import DTTM_ALIAS +from superset.utils.decorators import suppress_logging from superset.utils.pandas_postprocessing.utils import PROPHET_TIME_GRAIN_MAP @@ -52,8 +53,10 @@ def _prophet_fit_and_predict( # pylint: disable=too-many-arguments Fit a prophet model and return a DataFrame with predicted results. """ try: - # pylint: disable=import-outside-toplevel - from prophet import Prophet + # `prophet` complains about `plotly` not being installed + with suppress_logging("prophet.plot"): + # pylint: disable=import-outside-toplevel + from prophet import Prophet prophet_logger = logging.getLogger("prophet.plot") prophet_logger.setLevel(logging.CRITICAL) diff --git a/tests/unit_tests/utils/test_decorators.py b/tests/unit_tests/utils/test_decorators.py index 6cfe55b641..0a622f4bce 100644 --- a/tests/unit_tests/utils/test_decorators.py +++ b/tests/unit_tests/utils/test_decorators.py @@ -16,6 +16,7 @@ # under the License. +import logging import uuid from contextlib import nullcontext from inspect import isclass @@ -249,3 +250,47 @@ def test_context_decorator(flask_g_mock) -> None: context_func_not_callable() assert flask_g_mock.logs_context == {} + + +class ListHandler(logging.Handler): + """ + Simple logging handler that stores records in a list. + """ + + def __init__(self, *args: Any, **kwargs: Any): + super().__init__(*args, **kwargs) + self.log_records: list[logging.LogRecord] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.log_records.append(record) + + def reset(self) -> None: + self.log_records = [] + + +def test_suppress_logging() -> None: + """ + Test the `suppress_logging` decorator. + """ + handler = ListHandler() + logger = logging.getLogger("test-logger") + logger.setLevel(logging.INFO) + logger.addHandler(handler) + + def func() -> None: + logger.error("error") + logger.critical("critical") + + func() + assert len(handler.log_records) == 2 + + handler.log_records = [] + decorated = decorators.suppress_logging("test-logger")(func) + decorated() + assert len(handler.log_records) == 1 + assert handler.log_records[0].levelname == "CRITICAL" + + handler.log_records = [] + decorated = decorators.suppress_logging("test-logger", logging.CRITICAL + 1)(func) + decorated() + assert len(handler.log_records) == 0