From dda1dcf8ee217438acb45f2ad016ff1869c16112 Mon Sep 17 00:00:00 2001
From: AAfghahi <48933336+AAfghahi@users.noreply.github.com>
Date: Thu, 18 Aug 2022 10:32:25 -0400
Subject: [PATCH] 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 +++++++++++++
.../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 cc566fe6e..1c8cf3c8b 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 721a4db26..90c5040cc 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 d7be26eb3..24d4657b7 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 f5c9b79fd..6eb2405d0 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 42ee9b466..358294f1e 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 8298f4783..46318dd50 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 d6c46a08d..4b546fea0 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 6621c1fce..54de4cdf2 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 f9827580c..1df40e22f 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='This is a test alert
',
+ 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(