Merge pull request #1619 from shrmnk/enhance/email-task

Enhance EmailTask to accept smtp server settings and email_from

CHANGELOG.md

@@ -17,6 +17,7 @@ These changes are available in the [master branch](https://github.com/PrefectHQ/
 - Local Secrets set through environment variable now retain their casing - [#1601](https://github.com/PrefectHQ/prefect/issues/1601)
 - Agents can accept an optional `name` for logging and debugging - [#1612](https://github.com/PrefectHQ/prefect/pull/1612)
 - Added AWS configuration options for Fargate Agent (task_role_arn, execution_role_arn) - [#1614](https://github.com/PrefectHQ/prefect/pull/1614)
+- Change EmailTask to accept SMTP server settings as well as an email_from kwarg - [#1619](https://github.com/PrefectHQ/prefect/pull/1619)
 
 ### Task Library
 
@@ -40,6 +41,7 @@ These changes are available in the [master branch](https://github.com/PrefectHQ/
 ### Contributors
 
 - [Mark McDonald](https://github.com/mhmcdonal)
+- [Sherman K](https://github.com/shrmnk)
 
 ## 0.6.6 <Badge text="beta" type="success"/>
 

docs/core/task_library/notifications.md

@@ -2,6 +2,6 @@
 
 ## EmailTask <Badge text="task"/>
 
-Task for sending email from an authenticated Gmail address.  For this task to function properly, you must have the `"EMAIL_USERNAME"` and `"EMAIL_PASSWORD"` Prefect Secrets set.  It is recommended you use a [Google App Password](https://support.google.com/accounts/answer/185833) for this purpose.
+Task for sending email from an authenticated email service over SMTP. For this task to function properly, you must have the `"EMAIL_USERNAME"` and `"EMAIL_PASSWORD"` Prefect Secrets set.  It is recommended you use a [Google App Password](https://support.google.com/accounts/answer/185833) if you use Gmail. The default SMTP server is set to the Gmail SMTP server on port 465 (SMTP-over-SSL)
 
 [API Reference](/api/unreleased/tasks/notifications.html#prefect-tasks-notifications-email-task-emailtask)

src/prefect/tasks/notifications/email_task.py

@@ -11,28 +11,62 @@
 
 class EmailTask(Task):
     """
-    Task for sending email from an authenticated Gmail address.  For this task to function properly,
+    Task for sending email from an authenticated email service over SMTP. For this task to function properly,
     you must have the `"EMAIL_USERNAME"` and `"EMAIL_PASSWORD"` Prefect Secrets set.  It is recommended
-    you use a [Google App Password](https://support.google.com/accounts/answer/185833) for this purpose.
+    you use a [Google App Password](https://support.google.com/accounts/answer/185833) if you use Gmail.
+    The default SMTP server is set to the Gmail SMTP server on port 465 (SMTP-over-SSL)
 
     Args:
         - subject (str, optional): the subject of the email; can also be provided at runtime
         - msg (str, optional): the contents of the email; can also be provided at runtime
         - email_to (str, optional): the destination email address to send the message to; can also
             be provided at runtime
+        - email_from (str, optional): the email address to send from; defaults to notifications@prefect.io
+        - smtp_server (str, optional): the hostname of the SMTP server; defaults to smtp.gmail.com
+        - smtp_port (int, optional): the port number of the SMTP server; defaults to 465
+        - smtp_type (str, optional): either SSL or STARTTLS; defaults to SSL
         - **kwargs (Any, optional): additional keyword arguments to pass to the base Task initialization
     """
 
     def __init__(
-        self, subject: str = None, msg: str = None, email_to: str = None, **kwargs: Any
+        self,
+        subject: str = None,
+        msg: str = None,
+        email_to: str = None,
+        email_from: str = "notifications@prefect.io",
+        smtp_server: str = "smtp.gmail.com",
+        smtp_port: int = 465,
+        smtp_type: str = "SSL",
+        **kwargs: Any
     ):
         self.subject = subject
         self.msg = msg
         self.email_to = email_to
+        self.email_from = email_from
+        self.smtp_server = smtp_server
+        self.smtp_port = smtp_port
+        self.smtp_type = smtp_type
         super().__init__(**kwargs)
 
-    @defaults_from_attrs("subject", "msg", "email_to")
-    def run(self, subject: str = None, msg: str = None, email_to: str = None) -> None:
+    @defaults_from_attrs(
+        "subject",
+        "msg",
+        "email_to",
+        "email_from",
+        "smtp_server",
+        "smtp_port",
+        "smtp_type",
+    )
+    def run(
+        self,
+        subject: str = None,
+        msg: str = None,
+        email_to: str = None,
+        email_from: str = None,
+        smtp_server: str = None,
+        smtp_port: int = None,
+        smtp_type: str = None,
+    ) -> None:
         """
         Run method which sends an email.
 
@@ -43,6 +77,14 @@ def run(self, subject: str = None, msg: str = None, email_to: str = None) -> Non
                 at initialization
             - email_to (str, optional): the destination email address to send the message to;
                 defaults to the one provided at initialization
+            - email_from (str, optional): the email address to send from; defaults to the one
+                provided at initialization
+            - smtp_server (str, optional): the hostname of the SMTP server; defaults to the one
+                provided at initialization
+            - smtp_port (int, optional): the port number of the SMTP server; defaults to the one
+                provided at initialization
+            - smtp_type (str, optional): either SSL or STARTTLS; defaults to the one provided
+                at initialization
 
         Returns:
             - None
@@ -56,14 +98,19 @@ def run(self, subject: str = None, msg: str = None, email_to: str = None) -> Non
         contents.attach(MIMEText(cast(str, msg), "plain"))
 
         contents["Subject"] = Header(subject, "UTF-8")
-        contents["From"] = "notifications@prefect.io"
+        contents["From"] = email_from
         contents["To"] = email_to
 
         message = contents.as_string()
 
-        server = smtplib.SMTP_SSL("smtp.gmail.com", 465)
+        if smtp_type == "SSL":
+            server = smtplib.SMTP_SSL(smtp_server, smtp_port)
+        elif smtp_type == "STARTTLS":
+            server = smtplib.SMTP(smtp_server, smtp_port)
+            server.starttls()
+
         server.login(username, password)
         try:
-            server.sendmail("notifications@prefect.io", email_to, message)
+            server.sendmail(email_from, email_to, message)
         finally:
             server.quit()

tests/tasks/notifications/test_email_task.py

@@ -26,3 +26,18 @@ def test_username_password_pulled_from_secrets(self, monkeypatch):
             with context({"secrets": dict(EMAIL_USERNAME="foo", EMAIL_PASSWORD="bar")}):
                 res = t.run()
         assert smtp.SMTP_SSL.return_value.login.call_args[0] == ("foo", "bar")
+
+    def test_kwarg_for_email_from_get_passed_to_task_init(self, monkeypatch):
+        smtp = MagicMock()
+        monkeypatch.setattr("prefect.tasks.notifications.email_task.smtplib", smtp)
+        t = EmailTask(msg="", email_from="test@lvh.me")
+        with set_temporary_config({"cloud.use_local_secrets": True}):
+            with context({"secrets": dict(EMAIL_USERNAME="foo", EMAIL_PASSWORD="bar")}):
+                res = t.run()
+        assert smtp.SMTP_SSL.return_value.sendmail.call_args[0][0] == ("test@lvh.me")
+
+    def test_kwargs_for_smtp_server_get_passed_to_task_init(self):
+        t = EmailTask(smtp_server="mail.lvh.me", smtp_port=587, smtp_type="STARTTLS")
+        assert t.smtp_server == "mail.lvh.me"
+        assert t.smtp_port == 587
+        assert t.smtp_type == "STARTTLS"