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 <[email protected]> Authored: Wed Oct 18 12:27:14 2017 -0700 Committer: Chris Riccomini <[email protected]> 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 = [email protected] + + [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 = ['[email protected]', '[email protected]'] + 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': '[email protected]'}, {'email': '[email protected]'}]}], + 'from': {'email': u'[email protected]'}, + '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 = '[email protected]' + send_email_sendgrid(self.to, self.subject, self.html_content) + mock_post.assert_called_with(self.expected_mail_data)
