move execution_id to global context

eschutho · eschutho · commit dd1c7d2bb41f · 2023-11-16T17:05:30.000-08:00
diff --git a/superset/reports/commands/execute.py b/superset/reports/commands/execute.py
@@ -73,6 +73,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 context
 from superset.utils.screenshots import ChartScreenshot, DashboardScreenshot
 from superset.utils.urls import get_url_path
 
@@ -83,6 +84,7 @@ class BaseReportState:
     current_states: list[ReportState] = []
     initial: bool = False
 
+    @context()
     def __init__(
         self,
         session: Session,
@@ -234,7 +236,12 @@ def _get_screenshots(self) -> list[bytes]:
         try:
             image = screenshot.get_screenshot(user=user)
         except SoftTimeLimitExceeded as ex:
-            logger.warning("A timeout occurred while taking a screenshot.")
+            logger.warning(
+                "A timeout occurred while taking a screenshot.",
+                extra={
+                    "execution_id": self._execution_id,
+                },
+            )
             raise ReportScheduleScreenshotTimeout() from ex
         except Exception as ex:
             raise ReportScheduleScreenshotFailedError(
@@ -254,11 +261,23 @@ def _get_csv_data(self) -> bytes:
         auth_cookies = machine_auth_provider_factory.instance.get_auth_cookies(user)
 
         if self._report_schedule.chart.query_context is None:
-            logger.warning("No query context found, taking a screenshot to generate it")
+            logger.warning(
+                "No query context found, taking a screenshot to generate it",
+                extra={
+                    "execution_id": self._execution_id,
+                },
+            )
             self._update_query_context()
 
         try:
-            logger.info("Getting chart from %s as user %s", url, user.username)
+            logger.info(
+                "Getting chart from %s as user %s",
+                url,
+                user.username,
+                extra={
+                    "execution_id": self._execution_id,
+                },
+            )
             csv_data = get_chart_csv_data(chart_url=url, auth_cookies=auth_cookies)
         except SoftTimeLimitExceeded as ex:
             raise ReportScheduleCsvTimeout() from ex
@@ -283,11 +302,23 @@ def _get_embedded_data(self) -> pd.DataFrame:
         auth_cookies = machine_auth_provider_factory.instance.get_auth_cookies(user)
 
         if self._report_schedule.chart.query_context is None:
-            logger.warning("No query context found, taking a screenshot to generate it")
+            logger.warning(
+                "No query context found, taking a screenshot to generate it",
+                extra={
+                    "execution_id": self._execution_id,
+                },
+            )
             self._update_query_context()
 
         try:
-            logger.info("Getting chart from %s as user %s", url, user.username)
+            logger.info(
+                "Getting chart from %s as user %s",
+                url,
+                user.username,
+                extra={
+                    "execution_id": self._execution_id,
+                },
+            )
             dataframe = get_chart_dataframe(url, auth_cookies)
         except SoftTimeLimitExceeded as ex:
             raise ReportScheduleDataFrameTimeout() from ex
@@ -440,7 +471,12 @@ def _send(
         if notification_errors:
             # log all errors but raise based on the most severe
             for error in notification_errors:
-                logger.warning(str(error))
+                logger.warning(
+                    str(error),
+                    extra={
+                        "execution_id": self._execution_id,
+                    },
+                )
 
             if any(error.level == ErrorLevel.ERROR for error in notification_errors):
                 raise ReportScheduleSystemErrorsException(errors=notification_errors)
@@ -466,7 +502,9 @@ def send_error(self, name: str, message: str) -> None:
         logger.info(
             "header_data in notifications for alerts and reports %s, taskid, %s",
             header_data,
-            self._execution_id,
+            extra={
+                "execution_id": self._execution_id,
+            },
         )
         notification_content = NotificationContent(
             name=name, text=message, header_data=header_data
@@ -725,6 +763,9 @@ def run(self) -> None:
                         "Running report schedule %s as user %s",
                         self._execution_id,
                         username,
+                        extra={
+                            "execution_id": self._execution_id,
+                        },
                     )
                     ReportScheduleStateMachine(
                         session, self._execution_id, self._model, self._scheduled_dttm
@@ -740,6 +781,9 @@ def validate(self, session: Session = None) -> None:
             "session is validated: id %s, executionid: %s",
             self._model_id,
             self._execution_id,
+            extra={
+                "execution_id": self._execution_id,
+            },
         )
         self._model = (
             session.query(ReportSchedule).filter_by(id=self._model_id).one_or_none()
diff --git a/superset/reports/notifications/email.py b/superset/reports/notifications/email.py
@@ -22,6 +22,7 @@
 from typing import Any, Optional
 
 import nh3
+from flask import g
 from flask_babel import gettext as __
 
 from superset import app
@@ -190,6 +191,7 @@ def send(self) -> None:
         subject = self._get_subject()
         content = self._get_content()
         to = self._get_to()
+        global_context = getattr(g, "context", {}) or {}
         try:
             send_email_smtp(
                 to,
@@ -205,7 +207,11 @@ def send(self) -> None:
                 header_data=content.header_data,
             )
             logger.info(
-                "Report sent to email, notification content is %s", content.header_data
+                "Report sent to email, notification content is %s",
+                content.header_data,
+                extra={
+                    "execution_id": global_context.get("execution_id"),
+                },
             )
         except SupersetErrorsException as ex:
             raise NotificationError(
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 (
@@ -166,6 +167,8 @@ def send(self) -> None:
         channel = self._get_channel()
         body = self._get_body()
         file_type = "csv" if self._content.csv else "png"
+        global_context = getattr(g, "context", {}) or {}
+
         try:
             token = app.config["SLACK_API_TOKEN"]
             if callable(token):
@@ -183,7 +186,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_context.get("execution_id"),
+                },
+            )
         except (
             BotUserAccessError,
             SlackRequestError,
diff --git a/superset/utils/decorators.py b/superset/utils/decorators.py
@@ -20,8 +20,9 @@
 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
@@ -111,3 +112,38 @@ 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)})
+
+
+def context(
+    slice_id: int | None = None,
+    dashboard_id: int | None = None,
+    execution_id: str | UUID | None = None,
+) -> Callable[..., Any]:
+    """
+    Takes arguments and adds them to the global 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, "context"):
+                g.context = {}
+            available_context_values = ["slice_id", "dashboard_id", "execution_id"]
+            context_data = {
+                key: val
+                for key, val in kwargs.items()
+                if key in available_context_values
+            }
+
+            # if values are passed in to decorator directly, add them to context
+            # by overriding values from kwargs
+            for val in available_context_values:
+                if locals().get(val) is not None:
+                    context_data[val] = locals()[val]
+
+            g.context.update(context_data)
+            return f(*args, **kwargs)
+
+        return wrapped
+
+    return decorate
diff --git a/superset/utils/webdriver.py b/superset/utils/webdriver.py
@@ -23,7 +23,7 @@
 from time import sleep
 from typing import Any, TYPE_CHECKING
 
-from flask import current_app
+from flask import current_app, g
 from selenium.common.exceptions import (
     StaleElementReferenceException,
     TimeoutException,
@@ -227,6 +227,9 @@ def get_screenshot(self, url: str, element_name: str, user: User) -> bytes | Non
                     "Taking a PNG screenshot of url %s as user %s",
                     url,
                     user.username,
+                    extra={
+                        "execution_id": getattr(g, "context", {}).get("execution_id"),
+                    },
                 )
                 if current_app.config["SCREENSHOT_REPLACE_UNEXPECTED_ERRORS"]:
                     unexpected_errors = WebDriverPlaywright.find_unexpected_errors(page)
diff --git a/tests/unit_tests/notifications/email_tests.py b/tests/unit_tests/notifications/email_tests.py
@@ -14,7 +14,11 @@
 # 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, PropertyMock
+
 import pandas as pd
+from flask import g
 
 
 def test_render_description_with_html() -> None:
@@ -55,3 +59,82 @@ def test_render_description_with_html() -> None:
         in email_body
     )
     assert '<td>&lt;a href="http://www.example.com"&gt;333&lt;/a&gt;</td>' in email_body
+
+
+@patch("superset.reports.notifications.email.send_email_smtp")
+@patch("superset.reports.notifications.email.logger")
+def test_send_email(
+    logger_mock: MagicMock,
+    send_email_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.email import EmailNotification
+
+    execution_id = uuid.uuid4()
+    g.context = {"execution_id": execution_id}
+    content = NotificationContent(
+        name="test alert",
+        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 />',
+        header_data={
+            "notification_format": "PNG",
+            "notification_type": "Alert",
+            "owners": [1],
+            "notification_source": None,
+            "chart_id": None,
+            "dashboard_id": None,
+        },
+    )
+    EmailNotification(
+        recipient=ReportRecipients(
+            type=ReportRecipientType.EMAIL,
+            recipient_config_json='{"target": "foo@bar.com"}',
+        ),
+        content=content,
+    ).send()
+
+    logger_mock.info.assert_called_with(
+        "Report sent to email, notification content is %s",
+        {
+            "notification_format": "PNG",
+            "notification_type": "Alert",
+            "owners": [1],
+            "notification_source": None,
+            "chart_id": None,
+            "dashboard_id": None,
+        },
+        extra={"execution_id": execution_id},
+    )
+    # clear logger_mock and g.context
+    logger_mock.reset_mock()
+    g.context = None
+
+    # test with no execution_id
+    EmailNotification(
+        recipient=ReportRecipients(
+            type=ReportRecipientType.EMAIL,
+            recipient_config_json='{"target": "foo@bar.com"}',
+        ),
+        content=content,
+    ).send()
+    logger_mock.info.assert_called_with(
+        "Report sent to email, notification content is %s",
+        {
+            "notification_format": "PNG",
+            "notification_type": "Alert",
+            "owners": [1],
+            "notification_source": None,
+            "chart_id": None,
+            "dashboard_id": None,
+        },
+        extra={"execution_id": None},
+    )
diff --git a/tests/unit_tests/notifications/slack_tests.py b/tests/unit_tests/notifications/slack_tests.py
diff --git a/tests/unit_tests/utils/test_decorators.py b/tests/unit_tests/utils/test_decorators.py
