Repository: apache/airflow

Issue:
Allow custom email backends
The current `send_email` helper uses SMTP, but sometimes it's useful to be able to send mail without SMTP--for example when using API-based services like Sendgrid, Mailgun, etc. It would be useful to include a hook for overriding `send_email`--maybe a Django-style override like `email_backend = path.to.plugin:send_email`. If folks would be interested, I'm happy to send in a patch for this.



Discussion:
Seems reasonable. Good call on mimicking the Django approach. 


Resolving patch:
diff --git a/airflow/configuration.py b/airflow/configuration.py
--- a/airflow/configuration.py
+++ b/airflow/configuration.py
@@ -105,6 +105,9 @@ class AirflowConfigException(Exception):
         'default_queue': 'default',
         'flower_port': '5555'
     },
+    'email': {
+        'email_backend': 'airflow.utils.send_email_smtp',
+    },
     'smtp': {
         'smtp_starttls': True,
         'smtp_ssl': False,
@@ -218,6 +221,9 @@ class AirflowConfigException(Exception):
 # Filter the list of dags by owner name (requires authentication to be enabled)
 filter_by_owner = False
 
+[email]
+email_backend = airflow.utils.send_email_smtp
+
 [smtp]
 # If you want airflow to send emails on retries, failure, and you want to
 # the airflow.utils.send_email function, you have to configure an smtp
@@ -338,6 +344,9 @@ class AirflowConfigException(Exception):
 web_server_host = 0.0.0.0
 web_server_port = 8080
 
+[email]
+email_backend = airflow.utils.send_email_smtp
+
 [smtp]
 smtp_host = localhost
 smtp_user = airflow
diff --git a/airflow/utils.py b/airflow/utils.py
--- a/airflow/utils.py
+++ b/airflow/utils.py
@@ -15,6 +15,7 @@
 import errno
 from functools import wraps
 import imp
+import importlib
 import inspect
 import json
 import logging
@@ -494,6 +495,16 @@ def ask_yesno(question):
 
 
 def send_email(to, subject, html_content, files=None, dryrun=False):
+    """
+    Send email using backend specified in EMAIL_BACKEND.
+    """
+    path, attr = configuration.get('email', 'EMAIL_BACKEND').rsplit('.', 1)
+    module = importlib.import_module(path)
+    backend = getattr(module, attr)
+    return backend(to, subject, html_content, files=files, dryrun=dryrun)
+
+
+def send_email_smtp(to, subject, html_content, files=None, dryrun=False):
     """
     Send an email with html content
 


Test patch:
diff --git a/tests/core.py b/tests/core.py
--- a/tests/core.py
+++ b/tests/core.py
@@ -7,7 +7,11 @@
 import re
 import unittest
 import mock
+import tempfile
 from datetime import datetime, time, timedelta
+from email.mime.multipart import MIMEMultipart
+from email.mime.application import MIMEApplication
+import errno
 from time import sleep
 
 from dateutil.relativedelta import relativedelta
@@ -1287,6 +1291,90 @@ def test_tunnel(self):
             print("Closing tunnel")
 
 
+send_email_test = mock.Mock()
+
+
+class EmailTest(unittest.TestCase):
+
+    def setUp(self):
+        configuration.remove_option('email', 'EMAIL_BACKEND')
+
+    @mock.patch('airflow.utils.send_email_smtp')
+    def test_default_backend(self, mock_send_email):
+        res = utils.send_email('to', 'subject', 'content')
+        mock_send_email.assert_called_with('to', 'subject', 'content', files=None, dryrun=False)
+        assert res == mock_send_email.return_value
+
+    @mock.patch('airflow.utils.send_email_smtp')
+    def test_custom_backend(self, mock_send_email):
+        configuration.set('email', 'EMAIL_BACKEND', 'tests.core.send_email_test')
+        utils.send_email('to', 'subject', 'content')
+        send_email_test.assert_called_with('to', 'subject', 'content', files=None, dryrun=False)
+        assert not mock_send_email.called
+
+
+class EmailSmtpTest(unittest.TestCase):
+
+    def setUp(self):
+        configuration.set('smtp', 'SMTP_SSL', 'False')
+
+    @mock.patch('airflow.utils.send_MIME_email')
+    def test_send_smtp(self, mock_send_mime):
+        attachment = tempfile.NamedTemporaryFile()
+        attachment.write(b'attachment')
+        attachment.seek(0)
+        utils.send_email_smtp('to', 'subject', 'content', files=[attachment.name])
+        assert mock_send_mime.called
+        call_args = mock_send_mime.call_args[0]
+        assert call_args[0] == configuration.get('smtp', 'SMTP_MAIL_FROM')
+        assert call_args[1] == ['to']
+        msg = call_args[2]
+        assert msg['Subject'] == 'subject'
+        assert msg['From'] == configuration.get('smtp', 'SMTP_MAIL_FROM')
+        assert len(msg.get_payload()) == 2
+        mimeapp = MIMEApplication('attachment')
+        assert msg.get_payload()[-1].get_payload() == mimeapp.get_payload()
+
+    @mock.patch('smtplib.SMTP_SSL')
+    @mock.patch('smtplib.SMTP')
+    def test_send_mime(self, mock_smtp, mock_smtp_ssl):
+        mock_smtp.return_value = mock.Mock()
+        mock_smtp_ssl.return_value = mock.Mock()
+        msg = MIMEMultipart()
+        utils.send_MIME_email('from', 'to', msg, dryrun=False)
+        mock_smtp.assert_called_with(
+            configuration.get('smtp', 'SMTP_HOST'),
+            configuration.getint('smtp', 'SMTP_PORT'),
+        )
+        assert mock_smtp.return_value.starttls.called
+        mock_smtp.return_value.login.assert_called_with(
+            configuration.get('smtp', 'SMTP_USER'),
+            configuration.get('smtp', 'SMTP_PASSWORD'),
+        )
+        mock_smtp.return_value.sendmail.assert_called_with('from', 'to', msg.as_string())
+        assert mock_smtp.return_value.quit.called
+
+    @mock.patch('smtplib.SMTP_SSL')
+    @mock.patch('smtplib.SMTP')
+    def test_send_mime_ssl(self, mock_smtp, mock_smtp_ssl):
+        configuration.set('smtp', 'SMTP_SSL', 'True')
+        mock_smtp.return_value = mock.Mock()
+        mock_smtp_ssl.return_value = mock.Mock()
+        utils.send_MIME_email('from', 'to', MIMEMultipart(), dryrun=False)
+        assert not mock_smtp.called
+        mock_smtp_ssl.assert_called_with(
+            configuration.get('smtp', 'SMTP_HOST'),
+            configuration.getint('smtp', 'SMTP_PORT'),
+        )
+
+    @mock.patch('smtplib.SMTP_SSL')
+    @mock.patch('smtplib.SMTP')
+    def test_send_mime_dryrun(self, mock_smtp, mock_smtp_ssl):
+        utils.send_MIME_email('from', 'to', MIMEMultipart(), dryrun=True)
+        assert not mock_smtp.called
+        assert not mock_smtp_ssl.called
+
+
 if 'AIRFLOW_RUNALL_TESTS' in os.environ: