feat: global logs context (apache#26418)

eschutho · web-flow · commit aaa4a7b371a1 · 2024-01-16T14:44:30.000-08:00
diff --git a/superset/commands/report/execute.py b/superset/commands/report/execute.py
@@ -71,6 +71,7 @@
 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
@@ -22,6 +22,7 @@
 
 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 @@ def send(self) -> None:
         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 @@ def send(self) -> None:
                     )
             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
@@ -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 wrapped(*args: Any, **kwargs: Any) -> 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
@@ -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", '<a href="http://www.example.com">333</a>'],
+            }
+        ),
+        description='<p>This is <a href="#">a test</a> alert</p><br />',
+    )
+    with patch.object(
+        WebClient, "chat_postMessage", return_value=True
+    ) as chat_post_message_mock:
+        SlackNotification(
+            recipient=ReportRecipients(
+                type=ReportRecipientType.SLACK,
+                recipient_config_json='{"target": "some_channel"}',
+            ),
+            content=content,
+        ).send()
+        logger_mock.info.assert_called_with(
+            "Report sent to slack", extra={"execution_id": execution_id}
+        )
+        chat_post_message_mock.assert_called_with(
+            channel="some_channel",
+            text="""*test alert*
+
+<p>This is <a href="#">a test</a> alert</p><br />
+
+<None|Explore in Superset>
+
+```
+|    |   A |   B | C                                        |
+|---:|----:|----:|:-----------------------------------------|
+|  0 |   1 |   4 | 111                                      |
+|  1 |   2 |   5 | 222                                      |
+|  2 |   3 |   6 | <a href="http://www.example.com">333</a> |
+```
+""",
+        )
diff --git a/tests/unit_tests/utils/test_decorators.py b/tests/unit_tests/utils/test_decorators.py
