feat(alert/report): add ALERTS_ATTACH_REPORTS feature flags + feature (#13894)

* Add a feature flag ALERTS_ATTACH_REPORTS

* update test

* update feature flag

* add comment for feature flag

* add unit tests for alerts with attachments disabled

* fix lint

Co-authored-by: samtfm <sam@preset.io>

superset/config.py

@@ -359,6 +359,13 @@ DEFAULT_FEATURE_FLAGS: Dict[str, bool] = {
     "OMNIBAR": False,
     "DASHBOARD_RBAC": False,
     "ENABLE_EXPLORE_DRAG_AND_DROP": False,
+    # Enabling ALERTS_ATTACH_REPORTS, the system sends email and slack message
+    # with screenshot and link
+    # Disables ALERTS_ATTACH_REPORTS, the system DOES NOT generate screenshot
+    # for report with type 'alert' and sends email and slack message with only link;
+    # for report with type 'report' still send with email and slack message with
+    # screenshot and link
+    "ALERTS_ATTACH_REPORTS": True,
 }
 
 # Set the default view to card/grid view if thumbnail support is enabled.

superset/reports/commands/execute.py

@@ -26,6 +26,7 @@ from sqlalchemy.orm import Session
 from superset import app
 from superset.commands.base import BaseCommand
 from superset.commands.exceptions import CommandException
+from superset.extensions import feature_flag_manager
 from superset.models.reports import (
     ReportExecutionLog,
     ReportSchedule,
@@ -52,7 +53,7 @@ from superset.reports.dao import (
     ReportScheduleDAO,
 )
 from superset.reports.notifications import create_notification
-from superset.reports.notifications.base import NotificationContent, ScreenshotData
+from superset.reports.notifications.base import NotificationContent
 from superset.reports.notifications.exceptions import NotificationError
 from superset.utils.celery import session_scope
 from superset.utils.screenshots import (
@@ -153,7 +154,7 @@ class BaseReportState:
             raise ReportScheduleSelleniumUserNotFoundError()
         return user
 
-    def _get_screenshot(self) -> ScreenshotData:
+    def _get_screenshot(self) -> bytes:
         """
         Get a chart or dashboard screenshot
 
@@ -176,7 +177,6 @@ class BaseReportState:
                 window_size=app.config["WEBDRIVER_WINDOW"]["dashboard"],
                 thumb_size=app.config["WEBDRIVER_WINDOW"]["dashboard"],
             )
-        image_url = self._get_url(user_friendly=True)
         user = self._get_screenshot_user()
         try:
             image_data = screenshot.get_screenshot(user=user)
@@ -188,7 +188,7 @@ class BaseReportState:
             )
         if not image_data:
             raise ReportScheduleScreenshotFailedError()
-        return ScreenshotData(url=image_url, image=image_data)
+        return image_data
 
     def _get_notification_content(self) -> NotificationContent:
         """
@@ -196,7 +196,19 @@ class BaseReportState:
 
         :raises: ReportScheduleScreenshotFailedError
         """
-        screenshot_data = self._get_screenshot()
+        screenshot_data = None
+        url = self._get_url(user_friendly=True)
+        if (
+            feature_flag_manager.is_feature_enabled("ALERTS_ATTACH_REPORTS")
+            or self._report_schedule.type == ReportScheduleType.REPORT
+        ):
+            screenshot_data = self._get_screenshot()
+            if not screenshot_data:
+                return NotificationContent(
+                    name=self._report_schedule.name,
+                    text="Unexpected missing screenshot",
+                )
+
         if self._report_schedule.chart:
             name = (
                 f"{self._report_schedule.name}: "
@@ -207,7 +219,7 @@ class BaseReportState:
                 f"{self._report_schedule.name}: "
                 f"{self._report_schedule.dashboard.dashboard_title}"
             )
-        return NotificationContent(name=name, screenshot=screenshot_data)
+        return NotificationContent(name=name, url=url, screenshot=screenshot_data)
 
     def _send(self, notification_content: NotificationContent) -> None:
         """

superset/reports/notifications/base.py

@@ -21,16 +21,11 @@ from typing import Any, List, Optional, Type
 from superset.models.reports import ReportRecipients, ReportRecipientType
 
 
-@dataclass
-class ScreenshotData:
-    url: str  # url to chart/dashboard for this screenshot
-    image: bytes  # bytes for the screenshot
-
-
 @dataclass
 class NotificationContent:
     name: str
-    screenshot: Optional[ScreenshotData] = None
+    url: Optional[str] = None  # url to chart/dashboard for this screenshot
+    screenshot: Optional[bytes] = None  # bytes for the screenshot
     text: Optional[str] = None
 
 

superset/reports/notifications/email.py

@@ -63,21 +63,21 @@ class EmailNotification(BaseNotification):  # pylint: disable=too-few-public-met
             return EmailContent(body=self._error_template(self._content.text))
         # Get the domain from the 'From' address ..
         # and make a message id without the < > in the end
+        image = None
+        domain = self._get_smtp_domain()
+        msgid = make_msgid(domain)[1:-1]
+        body = __(
+            """
+            <b><a href="%(url)s">Explore in Superset</a></b><p></p>
+            <img src="cid:%(msgid)s">
+            """,
+            url=self._content.url,
+            msgid=msgid,
+        )
         if self._content.screenshot:
-            domain = self._get_smtp_domain()
+            image = {msgid: self._content.screenshot}
-            msgid = make_msgid(domain)[1:-1]
 
-            image = {msgid: self._content.screenshot.image}
+        return EmailContent(body=body, images=image)
-            body = __(
-                """
-                <b><a href="%(url)s">Explore in Superset</a></b><p></p>
-                <img src="cid:%(msgid)s">
-                """,
-                url=self._content.screenshot.url,
-                msgid=msgid,
-            )
-            return EmailContent(body=body, images=image)
-        return EmailContent(body=self._error_template("Unexpected missing screenshot"))
 
     def _get_subject(self) -> str:
         return __(

superset/reports/notifications/slack.py

@@ -57,20 +57,18 @@ class SlackNotification(BaseNotification):  # pylint: disable=too-few-public-met
     def _get_body(self) -> str:
         if self._content.text:
             return self._error_template(self._content.name, self._content.text)
-        if self._content.screenshot:
+        return __(
-            return __(
+            """
-                """
+            *%(name)s*\n
-                *%(name)s*\n
+            <%(url)s|Explore in Superset>
-                <%(url)s|Explore in Superset>
+            """,
-                """,
+            name=self._content.name,
-                name=self._content.name,
+            url=self._content.url,
-                url=self._content.screenshot.url,
+        )
-            )
-        return self._error_template(self._content.name, "Unexpected missing screenshot")
 
     def _get_inline_screenshot(self) -> Optional[Union[str, IOBase, bytes]]:
         if self._content.screenshot:
-            return self._content.screenshot.image
+            return self._content.screenshot
         return None
 
     @retry(SlackApiError, delay=10, backoff=2, tries=5)

tests/reports/commands_tests.py

@@ -750,6 +750,10 @@ def test_email_dashboard_report_fails(
 )
 @patch("superset.reports.notifications.email.send_email_smtp")
 @patch("superset.utils.screenshots.ChartScreenshot.get_screenshot")
+@patch.dict(
+    "superset.extensions.feature_flag_manager._feature_flags",
+    ALERTS_ATTACH_REPORTS=True,
+)
 def test_slack_chart_alert(screenshot_mock, email_mock, create_alert_email_chart):
     """
     ExecuteReport Command: Test chart slack alert
@@ -773,6 +777,34 @@ def test_slack_chart_alert(screenshot_mock, email_mock, create_alert_email_chart
         assert_log(ReportState.SUCCESS)
 
 
+@pytest.mark.usefixtures(
+    "load_birth_names_dashboard_with_slices", "create_alert_email_chart"
+)
+@patch("superset.reports.notifications.email.send_email_smtp")
+@patch.dict(
+    "superset.extensions.feature_flag_manager._feature_flags",
+    ALERTS_ATTACH_REPORTS=False,
+)
+def test_slack_chart_alert_no_attachment(email_mock, create_alert_email_chart):
+    """
+    ExecuteReport Command: Test chart slack alert
+    """
+    # setup screenshot mock
+
+    with freeze_time("2020-01-01T00:00:00Z"):
+        AsyncExecuteReportScheduleCommand(
+            test_id, create_alert_email_chart.id, datetime.utcnow()
+        ).run()
+
+        notification_targets = get_target_from_report_schedule(create_alert_email_chart)
+        # Assert the email smtp address
+        assert email_mock.call_args[0][0] == notification_targets[0]
+        # Assert the there is no attached image
+        assert email_mock.call_args[1]["images"] is None
+        # Assert logs are correct
+        assert_log(ReportState.SUCCESS)
+
+
 @pytest.mark.usefixtures(
     "load_birth_names_dashboard_with_slices", "create_report_slack_chart"
 )
@@ -859,6 +891,10 @@ def test_soft_timeout_alert(email_mock, create_alert_email_chart):
 )
 @patch("superset.reports.notifications.email.send_email_smtp")
 @patch("superset.utils.screenshots.ChartScreenshot.get_screenshot")
+@patch.dict(
+    "superset.extensions.feature_flag_manager._feature_flags",
+    ALERTS_ATTACH_REPORTS=True,
+)
 def test_soft_timeout_screenshot(screenshot_mock, email_mock, create_alert_email_chart):
     """
     ExecuteReport Command: Test soft timeout on screenshot
@@ -882,11 +918,11 @@ def test_soft_timeout_screenshot(screenshot_mock, email_mock, create_alert_email
 
 
 @pytest.mark.usefixtures(
-    "load_birth_names_dashboard_with_slices", "create_alert_email_chart"
+    "load_birth_names_dashboard_with_slices", "create_report_email_chart"
 )
 @patch("superset.reports.notifications.email.send_email_smtp")
 @patch("superset.utils.screenshots.ChartScreenshot.get_screenshot")
-def test_fail_screenshot(screenshot_mock, email_mock, create_alert_email_chart):
+def test_fail_screenshot(screenshot_mock, email_mock, create_report_email_chart):
     """
     ExecuteReport Command: Test soft timeout on screenshot
     """
@@ -896,10 +932,10 @@ def test_fail_screenshot(screenshot_mock, email_mock, create_alert_email_chart):
     screenshot_mock.side_effect = Exception("Unexpected error")
     with pytest.raises(ReportScheduleScreenshotFailedError):
         AsyncExecuteReportScheduleCommand(
-            test_id, create_alert_email_chart.id, datetime.utcnow()
+            test_id, create_report_email_chart.id, datetime.utcnow()
         ).run()
 
-    notification_targets = get_target_from_report_schedule(create_alert_email_chart)
+    notification_targets = get_target_from_report_schedule(create_report_email_chart)
     # Assert the email smtp address, asserts a notification was sent with the error
     assert email_mock.call_args[0][0] == notification_targets[0]
 
@@ -908,6 +944,32 @@ def test_fail_screenshot(screenshot_mock, email_mock, create_alert_email_chart):
     )
 
 
+@pytest.mark.usefixtures(
+    "load_birth_names_dashboard_with_slices", "create_alert_email_chart"
+)
+@patch("superset.reports.notifications.email.send_email_smtp")
+@patch.dict(
+    "superset.extensions.feature_flag_manager._feature_flags",
+    ALERTS_ATTACH_REPORTS=False,
+)
+def test_email_disable_screenshot(email_mock, create_alert_email_chart):
+    """
+    ExecuteReport Command: Test soft timeout on screenshot
+    """
+
+    AsyncExecuteReportScheduleCommand(
+        test_id, create_alert_email_chart.id, datetime.utcnow()
+    ).run()
+
+    notification_targets = get_target_from_report_schedule(create_alert_email_chart)
+    # Assert the email smtp address, asserts a notification was sent with the error
+    assert email_mock.call_args[0][0] == notification_targets[0]
+    # Assert the there is no attached image
+    assert email_mock.call_args[1]["images"] is None
+
+    assert_log(ReportState.SUCCESS)
+
+
 @pytest.mark.usefixtures("create_invalid_sql_alert_email_chart")
 @patch("superset.reports.notifications.email.send_email_smtp")
 def test_invalid_sql_alert(email_mock, create_invalid_sql_alert_email_chart):