diff --git a/superset/tasks/schedules.py b/superset/tasks/schedules.py index 3e8d6fc207d5a..20ce2d66482d1 100644 --- a/superset/tasks/schedules.py +++ b/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]: diff --git a/superset/utils/core.py b/superset/utils/core.py index c6d895c43824c..9c9fc8f7f3898 100644 --- a/superset/utils/core.py +++ b/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] diff --git a/superset/views/alerts.py b/superset/views/alerts.py index a06cf8a9b0dd8..f034903cc8d15 100644 --- a/superset/views/alerts.py +++ b/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) diff --git a/tests/alerts_tests.py b/tests/alerts_tests.py index 5c08acc8bf1dc..2f9491b830650 100644 --- a/tests/alerts_tests.py +++ b/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="testing@email.com", + is_test_alert=True, + ) + assert mock_run_alert.call_count == 1 + assert mock_deliver_alert.call_count == 1