diff --git a/airflow/bin/cli.py b/airflow/bin/cli.py index a24ed5ee426f6..9a752784888c0 100755 --- a/airflow/bin/cli.py +++ b/airflow/bin/cli.py @@ -935,7 +935,7 @@ def connections(args): if args.list: # Check that no other flags were passed to the command invalid_args = list() - for arg in ['conn_id', 'conn_uri', 'conn_extra']: + for arg in ['conn_id', 'conn_uri', 'conn_type', 'conn_extra']: if getattr(args, arg) is not None: invalid_args.append(arg) if invalid_args: @@ -960,7 +960,7 @@ def connections(args): if args.delete: # Check that only the `conn_id` arg was passed to the command invalid_args = list() - for arg in ['conn_uri', 'conn_extra']: + for arg in ['conn_uri', 'conn_type', 'conn_extra']: if getattr(args, arg) is not None: invalid_args.append(arg) if invalid_args: @@ -1013,7 +1013,8 @@ def connections(args): print(msg) return - new_conn = Connection(conn_id=args.conn_id, uri=args.conn_uri) + new_conn = Connection(conn_id=args.conn_id, uri=args.conn_uri, + conn_type=args.conn_type) if args.conn_extra is not None: new_conn.set_extra(args.conn_extra) @@ -1422,6 +1423,10 @@ class CLIFactory(object): ('--conn_uri',), help='Connection URI, required to add a connection', type=str), + 'conn_type': Arg( + ('--conn_type',), + help='Connection type, optional when overwriting conn_uri.scheme', + type=str), 'conn_extra': Arg( ('--conn_extra',), help='Connection `Extra` field, optional when adding a connection', @@ -1558,7 +1563,7 @@ class CLIFactory(object): 'func': connections, 'help': "List/Add/Delete connections", 'args': ('list_connections', 'add_connection', 'delete_connection', - 'conn_id', 'conn_uri', 'conn_extra'), + 'conn_id', 'conn_uri', 'conn_type', 'conn_extra'), }, ) subparsers_dict = {sp['func'].__name__: sp for sp in subparsers} diff --git a/airflow/config_templates/default_airflow.cfg b/airflow/config_templates/default_airflow.cfg index dee6dc7197222..fe2026194d6e2 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 = +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/models.py b/airflow/models.py index e3c52b5b72873..27f1df5e71296 100755 --- a/airflow/models.py +++ b/airflow/models.py @@ -564,7 +564,7 @@ def __init__( uri=None): self.conn_id = conn_id if uri: - self.parse_from_uri(uri) + self.parse_from_uri(uri, conn_type) else: self.conn_type = conn_type self.host = host @@ -574,12 +574,12 @@ def __init__( self.port = port self.extra = extra - def parse_from_uri(self, uri): + def parse_from_uri(self, uri, conn_type=None): temp_uri = urlparse(uri) hostname = temp_uri.hostname or '' if '%2f' in hostname: hostname = hostname.replace('%2f', '/').replace('%2F', '/') - conn_type = temp_uri.scheme + conn_type = conn_type or temp_uri.scheme if conn_type == 'postgresql': conn_type = 'postgres' self.conn_type = conn_type diff --git a/airflow/utils/email.py b/airflow/utils/email.py index fadd4d51ff9ae..472c0017b3252 100644 --- a/airflow/utils/email.py +++ b/airflow/utils/email.py @@ -21,13 +21,16 @@ 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 * 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 diff --git a/tests/core.py b/tests/core.py index 0c94137d15608..162b0c52032e3 100644 --- a/tests/core.py +++ b/tests/core.py @@ -1239,6 +1239,27 @@ def test_cli_connections_add_delete(self): self.assertEqual(result, (conn_id, 'postgres', 'host', 5432, extra[conn_id])) + # Check conn_type overwrite + session = settings.Session() + with mock.patch('sys.stdout', + new_callable=six.StringIO) as mock_stdout: + cli.connections(self.parser.parse_args( + ['connections', '--add', '--conn_id=overwrite-test', + '--conn_uri=scheme_with_underscore:', + '--conn_type=scheme_with_underscore'])) + stdout = mock_stdout.getvalue() + lines = [l for l in stdout.split('\n') if len(l) > 0] + self.assertListEqual(lines, [ + ("\tSuccessfully added `conn_id`=overwrite-test : " + + "scheme_with_underscore:") + ]) + result = (session + .query(models.Connection) + .filter(models.Connection.conn_id == 'overwrite-test') + .first()) + result = (result.conn_id, result.conn_type) + self.assertEqual(result, ('overwrite-test', 'scheme_with_underscore')) + # Delete connections with mock.patch('sys.stdout', new_callable=six.StringIO) as mock_stdout: diff --git a/tests/utils/test_email.py b/tests/utils/test_email.py new file mode 100644 index 0000000000000..568a5bd1b0ba1 --- /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 = ['foo@foo.com', 'bar@bar.com'] + self.subject = 'send-email-sendgrid unit test' + self.html_content = 'Foo bar' + self.expected_mail_data = { + 'content': [{'type': u'text/html', 'value': 'Foo bar'}], + 'personalizations': [ + {'to': [{'email': 'foo@foo.com'}, {'email': 'bar@bar.com'}]}], + 'from': {'email': u'foo@bar.com'}, + '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 = 'foo@bar.com' + send_email_sendgrid(self.to, self.subject, self.html_content) + mock_post.assert_called_with(self.expected_mail_data)