chore: prevent prophet from logging non-errors as errors (#27053)

This commit is contained in:
Beto Dealmeida 2024-02-08 17:33:07 -05:00 committed by GitHub
parent 81301ee54e
commit cf84f36a6c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 69 additions and 2 deletions

View File

@ -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)

View File

@ -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)

View File

@ -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