From 83a084ac1f008366df68743dff626b4b6fb7fcab Mon Sep 17 00:00:00 2001 From: fenglu-g Date: Thu, 26 Oct 2017 02:21:29 -0400 Subject: [PATCH] [AIRFLOW-1723] make sendgrid a plugin --- 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 ++++ .../utils/test_sendgrid.py} | 20 +++-- 7 files changed, 130 insertions(+), 57 deletions(-) create mode 100644 airflow/contrib/utils/__init__.py create mode 100644 airflow/contrib/utils/sendgrid.py create mode 100644 tests/contrib/utils/__init__.py rename tests/{utils/test_email.py => contrib/utils/test_sendgrid.py} (64%) diff --git a/airflow/config_templates/default_airflow.cfg b/airflow/config_templates/default_airflow.cfg index 91669792f70a5..fd78253f18df6 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 = -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 diff --git a/airflow/contrib/utils/__init__.py b/airflow/contrib/utils/__init__.py new file mode 100644 index 0000000000000..c82f5790fe906 --- /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. + diff --git a/airflow/contrib/utils/sendgrid.py b/airflow/contrib/utils/sendgrid.py new file mode 100644 index 0000000000000..7e83df12a34f8 --- /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)) diff --git a/airflow/utils/email.py b/airflow/utils/email.py index 21ae707db51e5..fadd4d51ff9ae 100644 --- a/airflow/utils/email.py +++ b/airflow/utils/email.py @@ -21,16 +21,13 @@ 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 diff --git a/tests/contrib/__init__.py b/tests/contrib/__init__.py index ff6f9e2529d08..58a73d1ec3b56 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 * diff --git a/tests/contrib/utils/__init__.py b/tests/contrib/utils/__init__.py new file mode 100644 index 0000000000000..cdd21472ecf95 --- /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. +# + diff --git a/tests/utils/test_email.py b/tests/contrib/utils/test_sendgrid.py similarity index 64% rename from tests/utils/test_email.py rename to tests/contrib/utils/test_sendgrid.py index 568a5bd1b0ba1..2459e5d37cbca 100644 --- a/tests/utils/test_email.py +++ b/tests/contrib/utils/test_sendgrid.py @@ -16,7 +16,7 @@ import logging import unittest -from airflow.utils.email import send_email_sendgrid +from airflow.contrib.utils.sendgrid import send_email try: from unittest import mock @@ -30,22 +30,26 @@ from mock import patch class SendEmailSendGridTest(unittest.TestCase): - # Unit test for send_email_sendgrid() + # Unit test for sendgrid.send_email() def setUp(self): self.to = ['foo@foo.com', 'bar@bar.com'] - self.subject = 'send-email-sendgrid unit test' + self.subject = 'sendgrid-send-email unit test' self.html_content = 'Foo bar' + self.cc = ['foo-cc@foo.com', 'bar-cc@bar.com'] + self.bcc = ['foo-bcc@foo.com', 'bar-bcc@bar.com'] self.expected_mail_data = { 'content': [{'type': u'text/html', 'value': 'Foo bar'}], 'personalizations': [ - {'to': [{'email': 'foo@foo.com'}, {'email': 'bar@bar.com'}]}], + {'cc': [{'email': 'foo-cc@foo.com'}, {'email': 'bar-cc@bar.com'}], + 'to': [{'email': 'foo@foo.com'}, {'email': 'bar@bar.com'}], + 'bcc': [{'email': 'foo-bcc@foo.com'}, {'email': 'bar-bcc@bar.com'}]}], 'from': {'email': u'foo@bar.com'}, - 'subject': 'send-email-sendgrid unit test'} + 'subject': 'sendgrid-send-email unit test'} # Test the right email is constructed. - @mock.patch('airflow.configuration.get') - @mock.patch('airflow.utils.email._post_sendgrid_mail') + @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 = 'foo@bar.com' - send_email_sendgrid(self.to, self.subject, self.html_content) + send_email(self.to, self.subject, self.html_content, cc=self.cc, bcc=self.bcc) mock_post.assert_called_with(self.expected_mail_data)