feat: add test email functionality to SQL-based email alerts (apache#…

…10476)

* added test email functionality

* formatting changes

* more formatting

* applied requested changes

* mypy

Co-authored-by: Jason Davis <@dropbox.com>

superset/tasks/schedules.py

@@ -540,6 +540,7 @@ def schedule_alert_query(  # pylint: disable=unused-argument
     report_type: ScheduleType,
     schedule_id: int,
     recipients: Optional[str] = None,
+    is_test_alert: Optional[bool] = False,
 ) -> None:
     model_cls = get_scheduler_model(report_type)
     dbsession = db.create_scoped_session()
@@ -551,6 +552,10 @@ def schedule_alert_query(  # pylint: disable=unused-argument
         return
 
     if report_type == ScheduleType.alert:
+        if is_test_alert and recipients:
+            deliver_alert(schedule, recipients)
+            return
+
         if run_alert_query(schedule, dbsession):
             # deliver_dashboard OR deliver_slice
             return
@@ -564,10 +569,12 @@ class AlertState:
     PASS = "pass"
 
 
-def deliver_alert(alert: Alert) -> None:
+def deliver_alert(alert: Alert, recipients: Optional[str] = None) -> None:
     logging.info("Triggering alert: %s", alert)
     img_data = None
     images = {}
+    recipients = recipients or alert.recipients
+
     if alert.slice:
 
         chart_url = get_url_path(
@@ -604,7 +611,7 @@ def deliver_alert(alert: Alert) -> None:
         image_url=image_url,
     )
 
-    _deliver_email(alert.recipients, deliver_as_group, subject, body, data, images)
+    _deliver_email(recipients, deliver_as_group, subject, body, data, images)
 
 
 def run_alert_query(alert: Alert, dbsession: Session) -> Optional[bool]:

superset/utils/core.py

@@ -832,6 +832,13 @@ def get_email_address_list(address_string: str) -> List[str]:
     return [x.strip() for x in address_string_list if x.strip()]
 
 
+def get_email_address_str(address_string: str) -> str:
+    address_list = get_email_address_list(address_string)
+    address_list_str = ", ".join(address_list)
+
+    return address_list_str
+
+
 def choicify(values: Iterable[Any]) -> List[Tuple[Any, Any]]:
     """Takes an iterable and makes an iterable of tuples with it"""
     return [(v, v) for v in values]

superset/views/alerts.py

@@ -14,14 +14,21 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+from typing import Dict, Optional, Union
+
+from croniter import croniter
 from flask_appbuilder import CompactCRUDMixin
 from flask_appbuilder.models.sqla.interface import SQLAInterface
 from flask_babel import lazy_gettext as _
+from wtforms import BooleanField, Form, StringField
 
 from superset.constants import RouteMethod
 from superset.models.alerts import Alert, AlertLog
-from superset.utils.core import markdown
+from superset.models.schedules import ScheduleType
+from superset.tasks.schedules import schedule_alert_query
+from superset.utils.core import get_email_address_str, markdown
 
+from ..exceptions import SupersetException
 from .base import SupersetModelView
 
 # TODO: access control rules for this module
@@ -44,6 +51,10 @@ class AlertModelView(SupersetModelView):  # pylint: disable=too-many-ancestors
     datamodel = SQLAInterface(Alert)
     route_base = "/alert"
     include_route_methods = RouteMethod.CRUD_SET
+    _extra_data: Dict[str, Union[bool, Optional[str]]] = {
+        "test_alert": False,
+        "test_email_recipients": None,
+    }
 
     list_columns = (
         "label",
@@ -68,6 +79,8 @@ class AlertModelView(SupersetModelView):  # pylint: disable=too-many-ancestors
         # "dashboard",
         "log_retention",
         "grace_period",
+        "test_alert",
+        "test_email_recipients",
     )
     label_columns = {
         "sql": "SQL",
@@ -95,5 +108,46 @@ class AlertModelView(SupersetModelView):  # pylint: disable=too-many-ancestors
             "Superset nags you again."
         ),
     }
+
+    add_form_extra_fields = {
+        "test_alert": BooleanField(
+            "Send Test Alert",
+            default=False,
+            description="If enabled, a test alert will be sent on the creation / update"
+            " of an active alert. All alerts after will be sent only if the SQL "
+            "statement defined above returns True.",
+        ),
+        "test_email_recipients": StringField(
+            "Test Email Recipients",
+            default=None,
+            description="List of recipients to send test email to. "
+            "If empty, an email will be sent to the original recipients.",
+        ),
+    }
+    edit_form_extra_fields = add_form_extra_fields
     edit_columns = add_columns
     related_views = [AlertLogModelView]
+
+    def process_form(self, form: Form, is_created: bool) -> None:
+        email_recipients = None
+        if form.test_email_recipients.data:
+            email_recipients = get_email_address_str(form.test_email_recipients.data)
+
+        self._extra_data["test_alert"] = form.test_alert.data
+        self._extra_data["test_email_recipients"] = email_recipients
+
+    def pre_add(self, item: "AlertModelView") -> None:
+        item.recipients = get_email_address_str(item.recipients)
+
+        if not croniter.is_valid(item.crontab):
+            raise SupersetException("Invalid crontab format")
+
+    def post_add(self, item: "AlertModelView") -> None:
+        if self._extra_data["test_alert"]:
+            recipients = self._extra_data["test_email_recipients"] or item.recipients
+            args = (ScheduleType.alert, item.id)
+            kwargs = dict(recipients=recipients, is_test_alert=True)
+            schedule_alert_query.apply_async(args=args, kwargs=kwargs)
+
+    def post_update(self, item: "AlertModelView") -> None:
+        self.post_add(item)

tests/alerts_tests.py

@@ -22,8 +22,9 @@
 
 from superset import db
 from superset.models.alerts import Alert, AlertLog
+from superset.models.schedules import ScheduleType
 from superset.models.slice import Slice
-from superset.tasks.schedules import run_alert_query
+from superset.tasks.schedules import run_alert_query, schedule_alert_query
 from superset.utils import core as utils
 from tests.test_app import app
 
@@ -112,3 +113,31 @@ def test_run_alert_query(mock_error, mock_deliver, setup_database):
     run_alert_query(database.query(Alert).filter_by(id=5).one(), database)
     assert mock_deliver.call_count == 1
     assert mock_error.call_count == 4
+
+
+@patch("superset.tasks.schedules.deliver_alert")
+@patch("superset.tasks.schedules.run_alert_query")
+def test_schedule_alert_query(mock_run_alert, mock_deliver_alert, setup_database):
+    database = setup_database
+    active_alert = database.query(Alert).filter_by(id=1).one()
+    inactive_alert = database.query(Alert).filter_by(id=3).one()
+
+    # Test that inactive alerts are no processed
+    schedule_alert_query(report_type=ScheduleType.alert, schedule_id=inactive_alert.id)
+    assert mock_run_alert.call_count == 0
+    assert mock_deliver_alert.call_count == 0
+
+    # Test that active alerts with no recipients passed in are processed regularly
+    schedule_alert_query(report_type=ScheduleType.alert, schedule_id=active_alert.id)
+    assert mock_run_alert.call_count == 1
+    assert mock_deliver_alert.call_count == 0
+
+    # Test that active alerts sent as a test are delivered immediately
+    schedule_alert_query(
+        report_type=ScheduleType.alert,
+        schedule_id=active_alert.id,
+        recipients="[email protected]",
+        is_test_alert=True,
+    )
+    assert mock_run_alert.call_count == 1
+    assert mock_deliver_alert.call_count == 1