feat: add header_data into emails (#20903)

* test sparkpost

* added logging info

* header function implementation

* added test

* daniel revisions

* daniel revision

* elizabeth review
This commit is contained in:
AAfghahi 2022-08-18 10:32:25 -04:00 committed by GitHub
parent fa0be30d49
commit dda1dcf8ee
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 163 additions and 6 deletions

View File

@ -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

View File

@ -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 = [

View File

@ -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,

View File

@ -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

View File

@ -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:

View File

@ -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(

View File

@ -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(

View File

@ -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

View File

@ -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(