[AIRFLOW-1700] Support conn_type overwrite when url scheme is invalid

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}

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 = <your send grid 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

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

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

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:

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 = '<b>Foo</b> bar'
+        self.expected_mail_data = {
+            'content': [{'type': u'text/html', 'value': '<b>Foo</b> 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)