Skip to content

Commit

Permalink
[AIRFLOW-1867] Fix sendgrid py3k bug; add sandbox mode
Browse files Browse the repository at this point in the history 
- Fix sendgrid attachments bug in py3k
- Add sandbox mode option
- Minor style fixes
- Add test
  • Loading branch information
@thesquelched
thesquelched committed Oct 16, 2018
1 parent 74c613b commit 123cf66
Show file tree
Hide file tree
Showing 2 changed files with 75 additions and 43 deletions.
40 changes: 25 additions & 15 deletions airflow/contrib/utils/sendgrid.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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()
Expand Down Expand Up @@ -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())

Expand All @@ -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)
78 changes: 50 additions & 28 deletions tests/contrib/utils/test_sendgrid.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@

import copy
import unittest
import tempfile
import os

from airflow.contrib.utils.sendgrid import send_email

Expand All @@ -41,58 +43,78 @@ def setUp(self):
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'}],
'content': [{'type': u'text/html', 'value': self.html_content}],
'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'}
'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': '[email protected]'}
self.expected_mail_data_extras['from'] = {
'name': 'Foo',
'email': '[email protected]',
}
# sender
self.expected_mail_data_sender = copy.deepcopy(self.expected_mail_data)
self.expected_mail_data_sender['from'] = \
{'name': 'Foo Bar', 'email': '[email protected]'}
self.expected_mail_data_sender['from'] = {
'name': 'Foo Bar',
'email': '[email protected]',
}

# 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='[email protected]'))
@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': '[email protected]'}.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='[email protected]',
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': '[email protected]',
'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='[email protected]', from_name='Foo Bar')
mock_post.assert_called_with(self.expected_mail_data_sender)

0 comments on commit 123cf66

Please sign in to comment.