feat: add header_data into emails (#20903)

* test sparkpost

* added logging info

* header function implementation

* added test

* daniel revisions

* daniel revision

* elizabeth review

superset/config.py

@@ -30,6 +30,7 @@ import re
 import sys
 from collections import OrderedDict
 from datetime import timedelta
+from email.mime.multipart import MIMEMultipart
 from typing import (
     Any,
     Callable,
@@ -1090,6 +1091,15 @@ def SQL_QUERY_MUTATOR(  # pylint: disable=invalid-name,unused-argument
     return sql
 
 
+# This allows for a user to add header data to any outgoing emails. For example,
+# if you need to include metadata in the header or you want to change the specifications
+# of the email title, header, or sender.
+def EMAIL_HEADER_MUTATOR(  # pylint: disable=invalid-name,unused-argument
+    msg: MIMEMultipart, **kwargs: Any
+) -> MIMEMultipart:
+    return msg
+
+
 # This auth provider is used by background (offline) tasks that need to access
 # protected resources. Can be overridden by end users in order to support
 # custom auth mechanisms

superset/reports/commands/execute.py

@@ -62,12 +62,14 @@ from superset.reports.models import (
     ReportRecipientType,
     ReportSchedule,
     ReportScheduleType,
+    ReportSourceFormat,
     ReportState,
 )
 from superset.reports.notifications import create_notification
 from superset.reports.notifications.base import NotificationContent
 from superset.reports.notifications.exceptions import NotificationError
 from superset.utils.celery import session_scope
+from superset.utils.core import HeaderDataType
 from superset.utils.csv import get_chart_csv_data, get_chart_dataframe
 from superset.utils.screenshots import ChartScreenshot, DashboardScreenshot
 from superset.utils.urls import get_url_path
@@ -305,6 +307,28 @@ class BaseReportState:
                 "Please try loading the chart and saving it again."
             ) from ex
 
+    def _get_log_data(self) -> HeaderDataType:
+        chart_id = None
+        dashboard_id = None
+        report_source = None
+        if self._report_schedule.chart:
+            report_source = ReportSourceFormat.CHART
+            chart_id = self._report_schedule.chart_id
+        else:
+            report_source = ReportSourceFormat.DASHBOARD
+            dashboard_id = self._report_schedule.dashboard_id
+
+        log_data: HeaderDataType = {
+            "notification_type": self._report_schedule.type,
+            "notification_source": report_source,
+            "notification_format": self._report_schedule.report_format,
+            "chart_id": chart_id,
+            "dashboard_id": dashboard_id,
+            "owners": self._report_schedule.owners,
+            "error_text": None,
+        }
+        return log_data
+
     def _get_notification_content(self) -> NotificationContent:
         """
         Gets a notification content, this is composed by a title and a screenshot
@@ -315,6 +339,7 @@ class BaseReportState:
         embedded_data = None
         error_text = None
         screenshot_data = []
+        header_data = self._get_log_data()
         url = self._get_url(user_friendly=True)
         if (
             feature_flag_manager.is_feature_enabled("ALERTS_ATTACH_REPORTS")
@@ -332,8 +357,11 @@ class BaseReportState:
                 if not csv_data:
                     error_text = "Unexpected missing csv file"
             if error_text:
+                header_data["error_text"] = error_text
                 return NotificationContent(
-                    name=self._report_schedule.name, text=error_text
+                    name=self._report_schedule.name,
+                    text=error_text,
+                    header_data=header_data,
                 )
 
         if (
@@ -352,6 +380,7 @@ class BaseReportState:
                 f"{self._report_schedule.name}: "
                 f"{self._report_schedule.dashboard.dashboard_title}"
             )
+
         return NotificationContent(
             name=name,
             url=url,
@@ -359,6 +388,7 @@ class BaseReportState:
             description=self._report_schedule.description,
             csv=csv_data,
             embedded_data=embedded_data,
+            header_data=header_data,
         )
 
     def _send(
@@ -404,7 +434,11 @@ class BaseReportState:
 
         :raises: ReportScheduleNotificationError
         """
-        notification_content = NotificationContent(name=name, text=message)
+        header_data = self._get_log_data()
+        header_data["error_text"] = message
+        notification_content = NotificationContent(
+            name=name, text=message, header_data=header_data
+        )
 
         # filter recipients to recipients who are also owners
         owner_recipients = [

superset/reports/models.py

@@ -82,6 +82,11 @@ class ReportCreationMethod(str, enum.Enum):
     ALERTS_REPORTS = "alerts_reports"
 
 
+class ReportSourceFormat(str, enum.Enum):
+    CHART = "chart"
+    DASHBOARD = "dashboard"
+
+
 report_schedule_user = Table(
     "report_schedule_user",
     metadata,

superset/reports/notifications/base.py

@@ -21,11 +21,13 @@ from typing import Any, List, Optional, Type
 import pandas as pd
 
 from superset.reports.models import ReportRecipients, ReportRecipientType
+from superset.utils.core import HeaderDataType
 
 
 @dataclass
 class NotificationContent:
     name: str
+    header_data: HeaderDataType  # this is optional to account for error states
     csv: Optional[bytes] = None  # bytes for csv file
     screenshots: Optional[List[bytes]] = None  # bytes for a list of screenshots
     text: Optional[str] = None

superset/reports/notifications/email.py

@@ -29,7 +29,7 @@ from superset import app
 from superset.reports.models import ReportRecipientType
 from superset.reports.notifications.base import BaseNotification
 from superset.reports.notifications.exceptions import NotificationError
-from superset.utils.core import send_email_smtp
+from superset.utils.core import HeaderDataType, send_email_smtp
 from superset.utils.decorators import statsd_gauge
 from superset.utils.urls import modify_url_query
 
@@ -67,6 +67,7 @@ ALLOWED_ATTRIBUTES = {
 @dataclass
 class EmailContent:
     body: str
+    header_data: Optional[HeaderDataType] = None
     data: Optional[Dict[str, Any]] = None
     images: Optional[Dict[str, bytes]] = None
 
@@ -170,7 +171,12 @@ class EmailNotification(BaseNotification):  # pylint: disable=too-few-public-met
 
         if self._content.csv:
             csv_data = {__("%(name)s.csv", name=self._content.name): self._content.csv}
-        return EmailContent(body=body, images=images, data=csv_data)
+        return EmailContent(
+            body=body,
+            images=images,
+            data=csv_data,
+            header_data=self._content.header_data,
+        )
 
     def _get_subject(self) -> str:
         return __(
@@ -199,6 +205,7 @@ class EmailNotification(BaseNotification):  # pylint: disable=too-few-public-met
                 bcc="",
                 mime_subtype="related",
                 dryrun=False,
+                header_data=content.header_data,
             )
             logger.info("Report sent to email")
         except Exception as ex:

superset/utils/core.py

@@ -184,6 +184,16 @@ class DatasourceType(str, Enum):
     VIEW = "view"
 
 
+class HeaderDataType(TypedDict):
+    notification_format: str
+    owners: List[int]
+    notification_type: str
+    notification_source: Optional[str]
+    chart_id: Optional[int]
+    dashboard_id: Optional[int]
+    error_text: Optional[str]
+
+
 class DatasourceDict(TypedDict):
     type: str  # todo(hugh): update this to be DatasourceType
     id: int
@@ -904,6 +914,7 @@ def send_email_smtp(  # pylint: disable=invalid-name,too-many-arguments,too-many
     cc: Optional[str] = None,
     bcc: Optional[str] = None,
     mime_subtype: str = "mixed",
+    header_data: Optional[HeaderDataType] = None,
 ) -> None:
     """
     Send an email with html content, eg:
@@ -917,6 +928,7 @@ def send_email_smtp(  # pylint: disable=invalid-name,too-many-arguments,too-many
     msg["Subject"] = subject
     msg["From"] = smtp_mail_from
     msg["To"] = ", ".join(smtp_mail_to)
+
     msg.preamble = "This is a multi-part message in MIME format."
 
     recipients = smtp_mail_to
@@ -963,8 +975,10 @@ def send_email_smtp(  # pylint: disable=invalid-name,too-many-arguments,too-many
         image.add_header("Content-ID", "<%s>" % msgid)
         image.add_header("Content-Disposition", "inline")
         msg.attach(image)
-
+    msg_mutator = config["EMAIL_HEADER_MUTATOR"]
-    send_mime_email(smtp_mail_from, recipients, msg, config, dryrun=dryrun)
+    # the base notification returns the message without any editing.
+    new_msg = msg_mutator(msg, **(header_data or {}))
+    send_mime_email(smtp_mail_from, recipients, new_msg, config, dryrun=dryrun)
 
 
 def send_mime_email(

tests/integration_tests/email_tests.py

@@ -58,6 +58,37 @@ class TestEmailSmtp(SupersetTestCase):
         mimeapp = MIMEApplication("attachment")
         assert msg.get_payload()[-1].get_payload() == mimeapp.get_payload()
 
+    @mock.patch("superset.utils.core.send_mime_email")
+    def test_send_smtp_with_email_mutator(self, mock_send_mime):
+        attachment = tempfile.NamedTemporaryFile()
+        attachment.write(b"attachment")
+        attachment.seek(0)
+
+        # putting this into a variable so that we can reset after the test
+        base_email_mutator = app.config["EMAIL_HEADER_MUTATOR"]
+
+        def mutator(msg, **kwargs):
+            msg["foo"] = "bar"
+            return msg
+
+        app.config["EMAIL_HEADER_MUTATOR"] = mutator
+        utils.send_email_smtp(
+            "to", "subject", "content", app.config, files=[attachment.name]
+        )
+        assert mock_send_mime.called
+        call_args = mock_send_mime.call_args[0]
+        logger.debug(call_args)
+        assert call_args[0] == app.config["SMTP_MAIL_FROM"]
+        assert call_args[1] == ["to"]
+        msg = call_args[2]
+        assert msg["Subject"] == "subject"
+        assert msg["From"] == app.config["SMTP_MAIL_FROM"]
+        assert msg["foo"] == "bar"
+        assert len(msg.get_payload()) == 2
+        mimeapp = MIMEApplication("attachment")
+        assert msg.get_payload()[-1].get_payload() == mimeapp.get_payload()
+        app.config["EMAIL_HEADER_MUTATOR"] = base_email_mutator
+
     @mock.patch("superset.utils.core.send_mime_email")
     def test_send_smtp_data(self, mock_send_mime):
         utils.send_email_smtp(

tests/integration_tests/reports/commands/execute_dashboard_report_tests.py

@@ -25,6 +25,7 @@ from superset.dashboards.permalink.commands.create import (
 )
 from superset.models.dashboard import Dashboard
 from superset.reports.commands.execute import AsyncExecuteReportScheduleCommand
+from superset.reports.models import ReportSourceFormat
 from tests.integration_tests.fixtures.tabbed_dashboard import tabbed_dashboard
 from tests.integration_tests.reports.utils import create_dashboard_report
 
@@ -66,3 +67,47 @@ def test_report_for_dashboard_with_tabs(
         assert digest == dashboard.digest
         assert send_email_smtp_mock.call_count == 1
         assert len(send_email_smtp_mock.call_args.kwargs["images"]) == 1
+
+
+@patch("superset.reports.notifications.email.send_email_smtp")
+@patch(
+    "superset.reports.commands.execute.DashboardScreenshot",
+)
+@patch(
+    "superset.dashboards.permalink.commands.create.CreateDashboardPermalinkCommand.run"
+)
+def test_report_with_header_data(
+    create_dashboard_permalink_mock: MagicMock,
+    dashboard_screenshot_mock: MagicMock,
+    send_email_smtp_mock: MagicMock,
+    tabbed_dashboard: Dashboard,
+) -> None:
+    create_dashboard_permalink_mock.return_value = "permalink"
+    dashboard_screenshot_mock.get_screenshot.return_value = b"test-image"
+    current_app.config["ALERT_REPORTS_NOTIFICATION_DRY_RUN"] = False
+
+    with create_dashboard_report(
+        dashboard=tabbed_dashboard,
+        extra={"active_tabs": ["TAB-L1B"]},
+        name="test report tabbed dashboard",
+    ) as report_schedule:
+        dashboard: Dashboard = report_schedule.dashboard
+        AsyncExecuteReportScheduleCommand(
+            str(uuid4()), report_schedule.id, datetime.utcnow()
+        ).run()
+        dashboard_state = report_schedule.extra.get("dashboard", {})
+        permalink_key = CreateDashboardPermalinkCommand(
+            dashboard.id, dashboard_state
+        ).run()
+
+        assert dashboard_screenshot_mock.call_count == 1
+        (url, digest) = dashboard_screenshot_mock.call_args.args
+        assert url.endswith(f"/superset/dashboard/p/{permalink_key}/")
+        assert digest == dashboard.digest
+        assert send_email_smtp_mock.call_count == 1
+        header_data = send_email_smtp_mock.call_args.kwargs["header_data"]
+        assert header_data.get("dashboard_id") == dashboard.id
+        assert header_data.get("notification_format") == report_schedule.report_format
+        assert header_data.get("notification_source") == ReportSourceFormat.DASHBOARD
+        assert header_data.get("notification_type") == report_schedule.type
+        assert len(send_email_smtp_mock.call_args.kwargs["header_data"]) == 7

tests/unit_tests/notifications/email_tests.py

@@ -34,6 +34,15 @@ def test_render_description_with_html() -> None:
             }
         ),
         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,
+            "error_text": None,
+        },
     )
     email_body = (
         EmailNotification(