Repository: incubator-airflow Updated Branches: refs/heads/master b3c247d3b -> 574e1c63d
[AIRFLOW-1723] Make sendgrid a plugin Closes #2727 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/574e1c63 Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/574e1c63 Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/574e1c63 Branch: refs/heads/master Commit: 574e1c63d9f818be07978e5deda413bf50b6c667 Parents: b3c247d Author: fenglu-g <[email protected]> Authored: Mon Oct 30 10:55:25 2017 -0700 Committer: Chris Riccomini <[email protected]> Committed: Mon Oct 30 10:55:30 2017 -0700 ---------------------------------------------------------------------- airflow/config_templates/default_airflow.cfg | 6 -- airflow/contrib/utils/__init__.py | 14 ++++ airflow/contrib/utils/sendgrid.py | 88 +++++++++++++++++++++++ airflow/utils/email.py | 43 ----------- tests/contrib/__init__.py | 1 + tests/contrib/utils/__init__.py | 15 ++++ tests/contrib/utils/test_sendgrid.py | 55 ++++++++++++++ tests/utils/test_email.py | 51 ------------- 8 files changed, 173 insertions(+), 100 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/574e1c63/airflow/config_templates/default_airflow.cfg ---------------------------------------------------------------------- diff --git a/airflow/config_templates/default_airflow.cfg b/airflow/config_templates/default_airflow.cfg index 9166979..fd78253 100644 --- a/airflow/config_templates/default_airflow.cfg +++ b/airflow/config_templates/default_airflow.cfg @@ -244,12 +244,6 @@ 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/574e1c63/airflow/contrib/utils/__init__.py ---------------------------------------------------------------------- diff --git a/airflow/contrib/utils/__init__.py b/airflow/contrib/utils/__init__.py new file mode 100644 index 0000000..c82f579 --- /dev/null +++ b/airflow/contrib/utils/__init__.py @@ -0,0 +1,14 @@ +# -*- 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. + http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/574e1c63/airflow/contrib/utils/sendgrid.py ---------------------------------------------------------------------- diff --git a/airflow/contrib/utils/sendgrid.py b/airflow/contrib/utils/sendgrid.py new file mode 100644 index 0000000..7e83df1 --- /dev/null +++ b/airflow/contrib/utils/sendgrid.py @@ -0,0 +1,88 @@ +# -*- 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. + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function +from __future__ import unicode_literals + +import base64 +import mimetypes +import os +import sendgrid + +from airflow.utils.email import get_email_address_list +from airflow.utils.log.logging_mixin import LoggingMixin +from sendgrid.helpers.mail import Attachment, Content, Email, Mail, Personalization + + +def send_email(to, subject, html_content, files=None, dryrun=False, cc=None, bcc=None, mime_subtype='mixed'): + """ + Send an email with html content using sendgrid. + + To use this plugin: + 0. include sendgrid subpackage as part of your Airflow installation, e.g., + pip install airflow[sendgrid] + 1. update [email] backend in airflow.cfg, i.e., + [email] + email_backend = airflow.contrib.utils.sendgrid.send_email + 2. configure Sendgrid specific environment variables at all Airflow instances: + SENDGRID_MAIL_FROM={your-mail-from} + SENDGRID_API_KEY={your-sendgrid-api-key}. + """ + mail = Mail() + mail.from_email = Email(os.environ.get('SENDGRID_MAIL_FROM')) + mail.subject = subject + + # Add the recipient list of to emails. + personalization = Personalization() + to = get_email_address_list(to) + for to_address in to: + personalization.add_to(Email(to_address)) + if cc: + cc = get_email_address_list(cc) + for cc_address in cc: + personalization.add_cc(Email(cc_address)) + if bcc: + bcc = get_email_address_list(bcc) + for bcc_address in bcc: + personalization.add_bcc(Email(bcc_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=os.environ.get('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('Email with subject %s is successfully sent to recipients: %s' % + (mail_data['subject'], mail_data['personalizations'])) + else: + log.warning('Failed to send out email with subject %s, status code: %s' % + (mail_data['subject'], response.status_code)) http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/574e1c63/airflow/utils/email.py ---------------------------------------------------------------------- diff --git a/airflow/utils/email.py b/airflow/utils/email.py index 21ae707..fadd4d5 100644 --- a/airflow/utils/email.py +++ b/airflow/utils/email.py @@ -21,16 +21,13 @@ 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 @@ -47,46 +44,6 @@ 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/574e1c63/tests/contrib/__init__.py ---------------------------------------------------------------------- diff --git a/tests/contrib/__init__.py b/tests/contrib/__init__.py index ff6f9e2..58a73d1 100644 --- a/tests/contrib/__init__.py +++ b/tests/contrib/__init__.py @@ -15,3 +15,4 @@ from __future__ import absolute_import from .operators import * from .sensors import * +from .utils import * http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/574e1c63/tests/contrib/utils/__init__.py ---------------------------------------------------------------------- diff --git a/tests/contrib/utils/__init__.py b/tests/contrib/utils/__init__.py new file mode 100644 index 0000000..cdd2147 --- /dev/null +++ b/tests/contrib/utils/__init__.py @@ -0,0 +1,15 @@ +# -*- 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. +# + http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/574e1c63/tests/contrib/utils/test_sendgrid.py ---------------------------------------------------------------------- diff --git a/tests/contrib/utils/test_sendgrid.py b/tests/contrib/utils/test_sendgrid.py new file mode 100644 index 0000000..2459e5d --- /dev/null +++ b/tests/contrib/utils/test_sendgrid.py @@ -0,0 +1,55 @@ +# -*- 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.contrib.utils.sendgrid import send_email + +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 sendgrid.send_email() + def setUp(self): + self.to = ['[email protected]', '[email protected]'] + self.subject = 'sendgrid-send-email unit test' + self.html_content = '<b>Foo</b> bar' + self.cc = ['[email protected]', '[email protected]'] + self.bcc = ['[email protected]', '[email protected]'] + self.expected_mail_data = { + 'content': [{'type': u'text/html', 'value': '<b>Foo</b> bar'}], + 'personalizations': [ + {'cc': [{'email': '[email protected]'}, {'email': '[email protected]'}], + 'to': [{'email': '[email protected]'}, {'email': '[email protected]'}], + 'bcc': [{'email': '[email protected]'}, {'email': '[email protected]'}]}], + 'from': {'email': u'[email protected]'}, + 'subject': 'sendgrid-send-email unit test'} + + # Test the right email is constructed. + @mock.patch('os.environ.get') + @mock.patch('airflow.contrib.utils.sendgrid._post_sendgrid_mail') + def test_send_email_sendgrid_correct_email(self, mock_post, mock_get): + mock_get.return_value = '[email protected]' + send_email(self.to, self.subject, self.html_content, cc=self.cc, bcc=self.bcc) + mock_post.assert_called_with(self.expected_mail_data) http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/574e1c63/tests/utils/test_email.py ---------------------------------------------------------------------- diff --git a/tests/utils/test_email.py b/tests/utils/test_email.py deleted file mode 100644 index 568a5bd..0000000 --- a/tests/utils/test_email.py +++ /dev/null @@ -1,51 +0,0 @@ -# -*- 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)
