You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by aa...@apache.org on 2022/08/18 14:32:35 UTC

[superset] branch master updated: feat: add header_data into emails (#20903)

This is an automated email from the ASF dual-hosted git repository.

aafghahi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git


The following commit(s) were added to refs/heads/master by this push:
     new dda1dcf8ee feat: add header_data into emails (#20903)
dda1dcf8ee is described below

commit dda1dcf8ee217438acb45f2ad016ff1869c16112
Author: AAfghahi <48...@users.noreply.github.com>
AuthorDate: Thu Aug 18 10:32:25 2022 -0400

    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                                 | 10 +++++
 superset/reports/commands/execute.py               | 38 +++++++++++++++++-
 superset/reports/models.py                         |  5 +++
 superset/reports/notifications/base.py             |  2 +
 superset/reports/notifications/email.py            | 11 +++++-
 superset/utils/core.py                             | 18 ++++++++-
 tests/integration_tests/email_tests.py             | 31 +++++++++++++++
 .../commands/execute_dashboard_report_tests.py     | 45 ++++++++++++++++++++++
 tests/unit_tests/notifications/email_tests.py      |  9 +++++
 9 files changed, 163 insertions(+), 6 deletions(-)

diff --git a/superset/config.py b/superset/config.py
index cc566fe6ee..1c8cf3c8b7 100644
--- a/superset/config.py
+++ b/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
diff --git a/superset/reports/commands/execute.py b/superset/reports/commands/execute.py
index 721a4db267..90c5040cc1 100644
--- a/superset/reports/commands/execute.py
+++ b/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 = [
diff --git a/superset/reports/models.py b/superset/reports/models.py
index d7be26eb32..24d4657b7d 100644
--- a/superset/reports/models.py
+++ b/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,
diff --git a/superset/reports/notifications/base.py b/superset/reports/notifications/base.py
index f5c9b79fd7..6eb2405d0f 100644
--- a/superset/reports/notifications/base.py
+++ b/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
diff --git a/superset/reports/notifications/email.py b/superset/reports/notifications/email.py
index 42ee9b4663..358294f1e1 100644
--- a/superset/reports/notifications/email.py
+++ b/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:
diff --git a/superset/utils/core.py b/superset/utils/core.py
index 8298f47832..46318dd50e 100644
--- a/superset/utils/core.py
+++ b/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)
-
-    send_mime_email(smtp_mail_from, recipients, msg, config, dryrun=dryrun)
+    msg_mutator = config["EMAIL_HEADER_MUTATOR"]
+    # 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(
diff --git a/tests/integration_tests/email_tests.py b/tests/integration_tests/email_tests.py
index d6c46a08d9..4b546fea0e 100644
--- a/tests/integration_tests/email_tests.py
+++ b/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(
diff --git a/tests/integration_tests/reports/commands/execute_dashboard_report_tests.py b/tests/integration_tests/reports/commands/execute_dashboard_report_tests.py
index 6621c1fce1..54de4cdf28 100644
--- a/tests/integration_tests/reports/commands/execute_dashboard_report_tests.py
+++ b/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
diff --git a/tests/unit_tests/notifications/email_tests.py b/tests/unit_tests/notifications/email_tests.py
index f9827580c6..1df40e22f9 100644
--- a/tests/unit_tests/notifications/email_tests.py
+++ b/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(