airflow-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From criccom...@apache.org
Subject incubator-airflow git commit: [AIRFLOW-1723] Support sendgrid in email backend
Date Wed, 18 Oct 2017 19:27:19 GMT
Repository: incubator-airflow
Updated Branches:
  refs/heads/master 6078e753a -> 7cb818bba


[AIRFLOW-1723] Support sendgrid in email backend

Closes #2695 from fenglu-g/master


Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/7cb818bb
Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/7cb818bb
Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/7cb818bb

Branch: refs/heads/master
Commit: 7cb818bbacb2a2695282471591a9e323d8efbf5c
Parents: 6078e75
Author: fenglu-g <fenglu@google.com>
Authored: Wed Oct 18 12:27:14 2017 -0700
Committer: Chris Riccomini <criccomini@apache.org>
Committed: Wed Oct 18 12:27:14 2017 -0700

----------------------------------------------------------------------
 airflow/config_templates/default_airflow.cfg |  6 +++
 airflow/utils/email.py                       | 43 +++++++++++++++++++
 scripts/ci/requirements.txt                  |  1 +
 setup.py                                     |  2 +
 tests/utils/test_email.py                    | 51 +++++++++++++++++++++++
 5 files changed, 103 insertions(+)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/7cb818bb/airflow/config_templates/default_airflow.cfg
----------------------------------------------------------------------
diff --git a/airflow/config_templates/default_airflow.cfg b/airflow/config_templates/default_airflow.cfg
index dee6dc7..fe20261 100644
--- a/airflow/config_templates/default_airflow.cfg
+++ b/airflow/config_templates/default_airflow.cfg
@@ -244,6 +244,12 @@ page_size = 100
 email_backend = airflow.utils.email.send_email_smtp
 
 
+[sendgrid]
+# Recommend an API key with Mail.send permission only.
+sendgrid_api_key = <your send grid api key>
+sendgrid_mail_from = airflow@example.com
+
+
 [smtp]
 # If you want airflow to send emails on retries, failure, and you want to use
 # the airflow.utils.email.send_email_smtp function, you have to configure an

http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/7cb818bb/airflow/utils/email.py
----------------------------------------------------------------------
diff --git a/airflow/utils/email.py b/airflow/utils/email.py
index fadd4d5..21ae707 100644
--- a/airflow/utils/email.py
+++ b/airflow/utils/email.py
@@ -21,13 +21,16 @@ from builtins import str
 from past.builtins import basestring
 
 import importlib
+import mimetypes
 import os
+import sendgrid
 import smtplib
 
 from email.mime.text import MIMEText
 from email.mime.multipart import MIMEMultipart
 from email.mime.application import MIMEApplication
 from email.utils import formatdate
+from sendgrid.helpers.mail import Attachment, Content, Email, Mail, Personalization
 
 from airflow import configuration
 from airflow.exceptions import AirflowConfigException
@@ -44,6 +47,46 @@ def send_email(to, subject, html_content, files=None, dryrun=False, cc=None,
bcc
     return backend(to, subject, html_content, files=files, dryrun=dryrun, cc=cc, bcc=bcc,
mime_subtype=mime_subtype)
 
 
+def send_email_sendgrid(to, subject, html_content, files=None, dryrun=False, cc=None, bcc=None,
mime_subtype='mixed'):
+    """
+    Send an email with html content using sendgrid.
+    """
+    mail = Mail()
+    mail.from_email = Email(configuration.get('sendgrid', 'SENDGRID_MAIL_FROM'))
+    mail.subject = subject
+
+    # Add the list of to emails.
+    to = get_email_address_list(to)
+    personalization = Personalization()
+    for to_address in to:
+      personalization.add_to(Email(to_address))
+    mail.add_personalization(personalization)
+    mail.add_content(Content('text/html', html_content))
+
+    # Add email attachment.
+    for fname in files or []:
+        basename = os.path.basename(fname)
+        attachment = Attachment()
+        with open(fname, "rb") as f:
+          attachment.content = base64.b64encode(f.read())
+          attachment.type = mimetypes.guess_type(basename)[0]
+          attachment.filename = basename
+          attachment.disposition = "attachment"
+          attachment.content_id = '<%s>' % basename
+        mail.add_attachment(attachment)
+    _post_sendgrid_mail(mail.get())
+
+
+def _post_sendgrid_mail(mail_data):
+    log = LoggingMixin().log
+    sg = sendgrid.SendGridAPIClient(apikey=configuration.get('sendgrid', 'SENDGRID_API_KEY'))
+    response = sg.client.mail.send.post(request_body=mail_data)
+    # 2xx status code.
+    if response.status_code >= 200 and response.status_code < 300:
+        log.info('The following email with subject %s is successfully sent to sendgrid.'
% subject)
+    else:
+        log.warning('Failed to send out email with subject %s, status code: %s' % (subject,
response.status_code))
+
 def send_email_smtp(to, subject, html_content, files=None, dryrun=False, cc=None, bcc=None,
mime_subtype='mixed'):
     """
     Send an email with html content

http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/7cb818bb/scripts/ci/requirements.txt
----------------------------------------------------------------------
diff --git a/scripts/ci/requirements.txt b/scripts/ci/requirements.txt
index d612d6f..1ea7a0b 100644
--- a/scripts/ci/requirements.txt
+++ b/scripts/ci/requirements.txt
@@ -79,6 +79,7 @@ rednose
 requests
 requests-kerberos
 requests_mock
+sendgrid
 setproctitle
 slackclient
 sphinx

http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/7cb818bb/setup.py
----------------------------------------------------------------------
diff --git a/setup.py b/setup.py
index d52bd3b..d520445 100644
--- a/setup.py
+++ b/setup.py
@@ -105,6 +105,7 @@ async = [
     'gevent>=0.13'
 ]
 azure = ['azure-storage>=0.34.0']
+sendgrid = ['sendgrid>=5.2.0']
 celery = [
     'celery>=4.0.0',
     'flower>=0.7.3'
@@ -273,6 +274,7 @@ def do_setup():
             's3': s3,
             'salesforce': salesforce,
             'samba': samba,
+            'sendgrid' : sendgrid,
             'slack': slack,
             'ssh': ssh,
             'statsd': statsd,

http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/7cb818bb/tests/utils/test_email.py
----------------------------------------------------------------------
diff --git a/tests/utils/test_email.py b/tests/utils/test_email.py
new file mode 100644
index 0000000..568a5bd
--- /dev/null
+++ b/tests/utils/test_email.py
@@ -0,0 +1,51 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import logging
+import unittest
+
+from airflow.utils.email import send_email_sendgrid
+
+try:
+    from unittest import mock
+except ImportError:
+    try:
+        import mock
+    except ImportError:
+        mock = None
+
+from mock import Mock
+from mock import patch
+
+class SendEmailSendGridTest(unittest.TestCase):
+    # Unit test for send_email_sendgrid()
+    def setUp(self):
+        self.to = ['foo@foo.com', 'bar@bar.com']
+        self.subject = 'send-email-sendgrid unit test'
+        self.html_content = '<b>Foo</b> bar'
+        self.expected_mail_data = {
+            'content': [{'type': u'text/html', 'value': '<b>Foo</b> bar'}],
+            'personalizations': [
+                {'to': [{'email': 'foo@foo.com'}, {'email': 'bar@bar.com'}]}],
+            'from': {'email': u'foo@bar.com'},
+            'subject': 'send-email-sendgrid unit test'}
+
+    # Test the right email is constructed.
+    @mock.patch('airflow.configuration.get')
+    @mock.patch('airflow.utils.email._post_sendgrid_mail')
+    def test_send_email_sendgrid_correct_email(self, mock_post, mock_get):
+        mock_get.return_value = 'foo@bar.com'
+        send_email_sendgrid(self.to, self.subject, self.html_content)
+        mock_post.assert_called_with(self.expected_mail_data)


Mime
View raw message