diff --git a/superset/commands/report/execute.py b/superset/commands/report/execute.py index d4b53e30d..349ef0c85 100644 --- a/superset/commands/report/execute.py +++ b/superset/commands/report/execute.py @@ -71,6 +71,7 @@ from superset.tasks.utils import get_executor from superset.utils.celery import session_scope from superset.utils.core import HeaderDataType, override_user from superset.utils.csv import get_chart_csv_data, get_chart_dataframe +from superset.utils.decorators import logs_context from superset.utils.screenshots import ChartScreenshot, DashboardScreenshot from superset.utils.urls import get_url_path @@ -81,6 +82,7 @@ class BaseReportState: current_states: list[ReportState] = [] initial: bool = False + @logs_context() def __init__( self, session: Session, diff --git a/superset/reports/notifications/slack.py b/superset/reports/notifications/slack.py index fbae398bc..41ed2ceeb 100644 --- a/superset/reports/notifications/slack.py +++ b/superset/reports/notifications/slack.py @@ -22,6 +22,7 @@ from typing import Union import backoff import pandas as pd +from flask import g from flask_babel import gettext as __ from slack_sdk import WebClient from slack_sdk.errors import ( @@ -175,6 +176,7 @@ Error: %(text)s channel = self._get_channel() body = self._get_body() file_type = "csv" if self._content.csv else "png" + global_logs_context = getattr(g, "logs_context", {}) or {} try: token = app.config["SLACK_API_TOKEN"] if callable(token): @@ -192,7 +194,12 @@ Error: %(text)s ) else: client.chat_postMessage(channel=channel, text=body) - logger.info("Report sent to slack") + logger.info( + "Report sent to slack", + extra={ + "execution_id": global_logs_context.get("execution_id"), + }, + ) except ( BotUserAccessError, SlackRequestError, diff --git a/superset/utils/decorators.py b/superset/utils/decorators.py index 9c21e3b5e..43b6909ca 100644 --- a/superset/utils/decorators.py +++ b/superset/utils/decorators.py @@ -16,16 +16,20 @@ # under the License. from __future__ import annotations +import logging import time from collections.abc import Iterator from contextlib import contextmanager from typing import Any, Callable, TYPE_CHECKING +from uuid import UUID -from flask import current_app, Response +from flask import current_app, g, Response from superset.utils import core as utils from superset.utils.dates import now_as_float +logger = logging.getLogger(__name__) + if TYPE_CHECKING: from superset.stats_logger import BaseStatsLogger @@ -61,6 +65,82 @@ def statsd_gauge(metric_prefix: str | None = None) -> Callable[..., Any]: return decorate +def logs_context( + context_func: Callable[..., dict[Any, Any]] | None = None, + **ctx_kwargs: int | str | UUID | None, +) -> Callable[..., Any]: + """ + Takes arguments and adds them to the global logs_context. + This is for logging purposes only and values should not be relied on or mutated + """ + + def decorate(f: Callable[..., Any]) -> Callable[..., Any]: + def wrapped(*args: Any, **kwargs: Any) -> Any: + if not hasattr(g, "logs_context"): + g.logs_context = {} + + # limit data that can be saved to logs_context + # in order to prevent antipatterns + available_logs_context_keys = [ + "slice_id", + "dashboard_id", + "dataset_id", + "execution_id", + "report_schedule_id", + ] + # set value from kwargs from + # wrapper function if it exists + # e.g. @logs_context() + # def my_func(slice_id=None, **kwargs) + # + # my_func(slice_id=2) + logs_context_data = { + key: val + for key, val in kwargs.items() + if key in available_logs_context_keys + if val is not None + } + + try: + # if keys are passed in to decorator directly, add them to logs_context + # by overriding values from kwargs + # e.g. @logs_context(slice_id=1, dashboard_id=1) + logs_context_data.update( + { + key: ctx_kwargs.get(key) + for key in available_logs_context_keys + if ctx_kwargs.get(key) is not None + } + ) + + if context_func is not None: + # if a context function is passed in, call it and add the + # returned values to logs_context + # context_func=lambda *args, **kwargs: { + # "slice_id": 1, "dashboard_id": 1 + # } + logs_context_data.update( + { + key: value + for key, value in context_func(*args, **kwargs).items() + if key in available_logs_context_keys + if value is not None + } + ) + + except (TypeError, KeyError, AttributeError): + # do nothing if the key doesn't exist + # or context is not callable + logger.warning("Invalid data was passed to the logs context decorator") + + g.logs_context.update(logs_context_data) + return f(*args, **kwargs) + + return wrapped + + return decorate + + @contextmanager def stats_timing(stats_key: str, stats_logger: BaseStatsLogger) -> Iterator[float]: """Provide a transactional scope around a series of operations.""" diff --git a/tests/unit_tests/notifications/slack_tests.py b/tests/unit_tests/notifications/slack_tests.py new file mode 100644 index 000000000..7e6cc3afc --- /dev/null +++ b/tests/unit_tests/notifications/slack_tests.py @@ -0,0 +1,86 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +import uuid +from unittest.mock import MagicMock, patch + +import pandas as pd +from flask import g + + +@patch("superset.reports.notifications.slack.g") +@patch("superset.reports.notifications.slack.logger") +def test_send_slack( + logger_mock: MagicMock, + flask_global_mock: MagicMock, +) -> None: + # `superset.models.helpers`, a dependency of following imports, + # requires app context + from superset.reports.models import ReportRecipients, ReportRecipientType + from superset.reports.notifications.base import NotificationContent + from superset.reports.notifications.slack import SlackNotification, WebClient + + execution_id = uuid.uuid4() + flask_global_mock.logs_context = {"execution_id": execution_id} + content = NotificationContent( + name="test alert", + header_data={ + "notification_format": "PNG", + "notification_type": "Alert", + "owners": [1], + "notification_source": None, + "chart_id": None, + "dashboard_id": None, + }, + embedded_data=pd.DataFrame( + { + "A": [1, 2, 3], + "B": [4, 5, 6], + "C": ["111", "222", '333'], + } + ), + description='
This is a test alert
This is a test alert