(superset) branch master updated: chore: prevent prophet from logging non-errors as errors (#27053)

[be...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

beto pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git


The following commit(s) were added to refs/heads/master by this push:
     new cf84f36a6c chore: prevent prophet from logging non-errors as errors (#27053)
cf84f36a6c is described below

commit cf84f36a6cf3341db139fbf3457b91685bbe8486
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Thu Feb 8 17:33:07 2024 -0500

    chore: prevent prophet from logging non-errors as errors (#27053)
---
 superset/utils/decorators.py                    | 19 +++++++++++
 superset/utils/pandas_postprocessing/prophet.py |  7 ++--
 tests/unit_tests/utils/test_decorators.py       | 45 +++++++++++++++++++++++++
 3 files changed, 69 insertions(+), 2 deletions(-)

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