Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 2 commits into from
Feb 8, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions superset/utils/decorators.py
Original file line number Diff line number Diff line change
Expand Up @@ -191,3 +191,22 @@ def wrapped(*args: Any, **kwargs: Any) -> 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(

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)
7 changes: 5 additions & 2 deletions superset/utils/pandas_postprocessing/prophet.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@

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


Expand Down Expand Up @@ -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)
Expand Down
45 changes: 45 additions & 0 deletions tests/unit_tests/utils/test_decorators.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
# under the License.


import logging
import uuid
from contextlib import nullcontext
from inspect import isclass
Expand Down Expand Up @@ -249,3 +250,47 @@ def context_func_not_callable() -> str:

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
Loading