Fokko closed pull request #2824: [AIRFLOW-1867] Fix sendgrid py3k bug; add
sandbox mode
URL: https://github.com/apache/incubator-airflow/pull/2824
This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:
As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):
diff --git a/airflow/contrib/utils/sendgrid.py
b/airflow/contrib/utils/sendgrid.py
index ceb3718ff5..1b932bcb9b 100644
--- a/airflow/contrib/utils/sendgrid.py
+++ b/airflow/contrib/utils/sendgrid.py
@@ -27,16 +27,16 @@
import os
import sendgrid
-from sendgrid.helpers.mail import Attachment, Content, Email, Mail, \
- Personalization, CustomArg, Category
+from sendgrid.helpers.mail import (
+ Attachment, Content, Email, Mail, Personalization, CustomArg, Category,
+ MailSettings, SandBoxMode)
from airflow.utils.email import get_email_address_list
from airflow.utils.log.logging_mixin import LoggingMixin
-def send_email(to, subject, html_content, files=None,
- dryrun=False, cc=None, bcc=None,
- mime_subtype='mixed', **kwargs):
+def send_email(to, subject, html_content, files=None, dryrun=False, cc=None,
+ bcc=None, mime_subtype='mixed', sandbox_mode=False, **kwargs):
"""
Send an email with html content using sendgrid.
@@ -50,11 +50,18 @@ def send_email(to, subject, html_content, files=None,
SENDGRID_MAIL_FROM={your-mail-from}
SENDGRID_API_KEY={your-sendgrid-api-key}.
"""
+ if files is None:
+ files = []
+
mail = Mail()
from_email = kwargs.get('from_email') or
os.environ.get('SENDGRID_MAIL_FROM')
from_name = kwargs.get('from_name') or
os.environ.get('SENDGRID_MAIL_SENDER')
mail.from_email = Email(from_email, from_name)
mail.subject = subject
+ mail.mail_settings = MailSettings()
+
+ if sandbox_mode:
+ mail.mail_settings.sandbox_mode = SandBoxMode(enable=True)
# Add the recipient list of to emails.
personalization = Personalization()
@@ -84,15 +91,18 @@ def send_email(to, subject, html_content, files=None,
mail.add_category(Category(cat))
# Add email attachment.
- for fname in files or []:
+ for fname in files:
basename = os.path.basename(fname)
+
attachment = Attachment()
+ attachment.type = mimetypes.guess_type(basename)[0]
+ attachment.filename = basename
+ attachment.disposition = "attachment"
+ attachment.content_id = '<{0}>'.format(basename)
+
with open(fname, "rb") as f:
- attachment.content = str(base64.b64encode(f.read()), 'utf-8')
- attachment.type = mimetypes.guess_type(basename)[0]
- attachment.filename = basename
- attachment.disposition = "attachment"
- attachment.content_id = '<%s>' % basename
+ attachment.content = base64.b64encode(f.read()).decode('utf-8')
+
mail.add_attachment(attachment)
_post_sendgrid_mail(mail.get())
@@ -103,8 +113,8 @@ def _post_sendgrid_mail(mail_data):
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']))
+ 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))
+ log.warning('Failed to send out email with subject %s, status code:
%s',
+ mail_data['subject'], response.status_code)
diff --git a/tests/contrib/utils/test_sendgrid.py
b/tests/contrib/utils/test_sendgrid.py
index 6710076c2d..d12c597d70 100644
--- a/tests/contrib/utils/test_sendgrid.py
+++ b/tests/contrib/utils/test_sendgrid.py
@@ -20,6 +20,8 @@
import copy
import unittest
+import tempfile
+import os
from airflow.contrib.utils.sendgrid import send_email
@@ -41,58 +43,78 @@ def setUp(self):
self.cc = ['foo...@foo.com', 'bar...@bar.com']
self.bcc = ['foo-...@foo.com', 'bar-...@bar.com']
self.expected_mail_data = {
- 'content': [{'type': u'text/html', 'value': '<b>Foo</b> bar'}],
+ 'content': [{'type': u'text/html', 'value': self.html_content}],
'personalizations': [
{'cc': [{'email': 'foo...@foo.com'}, {'email':
'bar...@bar.com'}],
'to': [{'email': 'f...@foo.com'}, {'email': 'b...@bar.com'}],
'bcc': [{'email': 'foo-...@foo.com'}, {'email':
'bar-...@bar.com'}]}],
'from': {'email': u'f...@bar.com'},
- 'subject': 'sendgrid-send-email unit test'}
+ 'subject': 'sendgrid-send-email unit test',
+ 'mail_settings': {},
+ }
self.personalization_custom_args = {'arg1': 'val1', 'arg2': 'val2'}
self.categories = ['cat1', 'cat2']
# extras
self.expected_mail_data_extras = copy.deepcopy(self.expected_mail_data)
- self.expected_mail_data_extras['personalizations'][0]['custom_args'] =
\
- self.personalization_custom_args
+ self.expected_mail_data_extras['personalizations'][0]['custom_args'] =
(
+ self.personalization_custom_args)
self.expected_mail_data_extras['categories'] = self.categories
- self.expected_mail_data_extras['from'] = \
- {'name': 'Foo', 'email': 'f...@bar.com'}
+ self.expected_mail_data_extras['from'] = {
+ 'name': 'Foo',
+ 'email': 'f...@bar.com',
+ }
# sender
self.expected_mail_data_sender = copy.deepcopy(self.expected_mail_data)
- self.expected_mail_data_sender['from'] = \
- {'name': 'Foo Bar', 'email': 'f...@foo.bar'}
+ self.expected_mail_data_sender['from'] = {
+ 'name': 'Foo Bar',
+ 'email': 'f...@foo.bar',
+ }
- # Test the right email is constructed.
-
- @mock.patch('os.environ.get')
+ # Test the right email is constructed.
+ @mock.patch('os.environ', dict(os.environ,
SENDGRID_MAIL_FROM='f...@bar.com'))
@mock.patch('airflow.contrib.utils.sendgrid._post_sendgrid_mail')
- def test_send_email_sendgrid_correct_email(self, mock_post, mock_get):
- def get_return(var):
- return {'SENDGRID_MAIL_FROM': 'f...@bar.com'}.get(var)
+ def test_send_email_sendgrid_correct_email(self, mock_post):
+ with tempfile.NamedTemporaryFile(mode='wt', suffix='.txt') as f:
+ f.write('this is some test data')
+ f.flush()
+
+ filename = os.path.basename(f.name)
+ expected_mail_data = dict(
+ self.expected_mail_data,
+ attachments=[{
+ 'content': 'dGhpcyBpcyBzb21lIHRlc3QgZGF0YQ==',
+ 'content_id': '<{0}>'.format(filename),
+ 'disposition': 'attachment',
+ 'filename': filename,
+ 'type': 'text/plain',
+ }],
+ )
- mock_get.side_effect = get_return
- send_email(self.to, self.subject, self.html_content, cc=self.cc,
bcc=self.bcc)
- mock_post.assert_called_with(self.expected_mail_data)
+ send_email(self.to,
+ self.subject,
+ self.html_content,
+ cc=self.cc,
+ bcc=self.bcc,
+ files=[f.name])
+ mock_post.assert_called_with(expected_mail_data)
# Test the right email is constructed.
- @mock.patch('os.environ.get')
+ @mock.patch(
+ 'os.environ',
+ dict(os.environ,
+ SENDGRID_MAIL_FROM='f...@bar.com',
+ SENDGRID_MAIL_SENDER='Foo')
+ )
@mock.patch('airflow.contrib.utils.sendgrid._post_sendgrid_mail')
- def test_send_email_sendgrid_correct_email_extras(self, mock_post,
mock_get):
- def get_return(var):
- return {'SENDGRID_MAIL_FROM': 'f...@bar.com',
- 'SENDGRID_MAIL_SENDER': 'Foo'}.get(var)
-
- mock_get.side_effect = get_return
+ def test_send_email_sendgrid_correct_email_extras(self, mock_post):
send_email(self.to, self.subject, self.html_content, cc=self.cc,
bcc=self.bcc,
personalization_custom_args=self.personalization_custom_args,
categories=self.categories)
mock_post.assert_called_with(self.expected_mail_data_extras)
- @mock.patch('os.environ.get')
+ @mock.patch('os.environ', {})
@mock.patch('airflow.contrib.utils.sendgrid._post_sendgrid_mail')
- def test_send_email_sendgrid_sender(self, mock_post, mock_get):
-
- mock_get.return_value = None
+ def test_send_email_sendgrid_sender(self, mock_post):
send_email(self.to, self.subject, self.html_content, cc=self.cc,
bcc=self.bcc,
from_email='f...@foo.bar', from_name='Foo Bar')
mock_post.assert_called_with(self.expected_mail_data_sender)
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
With regards,
Apache Git Services
