feat(api): add fast / slow lane depending on email type, closes #242

Adjust rate limits by running something like:
$ celery -A api control rate_limit email_slow_lane 10/m

api/api/celery.py

@@ -4,7 +4,7 @@ from celery import Celery
 
 os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'api.settings')
 
-app = Celery('api')
+app = Celery('api', include='desecapi.mail_backends')
 app.config_from_object('django.conf:settings', namespace='CELERY')
 app.autodiscover_tasks()
 

api/api/settings.py

@@ -42,7 +42,6 @@ INSTALLED_APPS = (
     'rest_framework',
     'desecapi.apps.AppConfig',
     'corsheaders',
-    'djcelery_email',
 )
 
 MIDDLEWARE = (
@@ -119,7 +118,7 @@ TEMPLATES = [
 ]
 
 # How and where to send mail
-EMAIL_BACKEND = 'djcelery_email.backends.CeleryEmailBackend'
+EMAIL_BACKEND = 'desecapi.mail_backends.MultiLaneEmailBackend'
 EMAIL_HOST = os.environ['DESECSTACK_API_EMAIL_HOST']
 EMAIL_HOST_USER = os.environ['DESECSTACK_API_EMAIL_HOST_USER']
 EMAIL_HOST_PASSWORD = os.environ['DESECSTACK_API_EMAIL_HOST_PASSWORD']
@@ -144,13 +143,12 @@ NSMASTER_PDNS_API_TOKEN = os.environ['DESECSTACK_NSMASTER_APIKEY']
 
 # Celery
 CELERY_BROKER_URL = 'amqp://rabbitmq'
-CELERY_EMAIL_CHUNK_SIZE = 1
-CELERY_EMAIL_TASK_CONFIG = {
-    'queue' : 'celery',
-    'rate_limit' : '3/m',
-}
-CELERY_TASK_DEFAULT_RATE_LIMIT = '3/m'
+CELERY_EMAIL_MESSAGE_EXTRA_ATTRIBUTES = []  # required because djcelery_email.utils accesses it
 CELERY_TASK_TIME_LIMIT = 30
+TASK_CONFIG = {
+    'email_fast_lane': {'rate_limit': '1/s'},
+    'email_slow_lane': {'rate_limit': '3/m'},
+}
 
 # pdns accepts request payloads of this size.
 # This will hopefully soon be configurable: https://github.com/PowerDNS/pdns/pull/7550

api/api/settings_quick_test.py

@@ -19,3 +19,6 @@ PASSWORD_HASHERS = [
 ]
 
 REST_FRAMEWORK['PAGE_SIZE'] = 20
+
+# Carry email backend connection over to test mail outbox
+CELERY_EMAIL_MESSAGE_EXTRA_ATTRIBUTES = ['connection']

api/desecapi/mail_backends.py

@@ -0,0 +1,47 @@
+import logging
+
+from celery import shared_task
+from django.conf import settings
+from django.core.mail import get_connection
+from django.core.mail.backends.base import BaseEmailBackend
+from djcelery_email.utils import dict_to_email, email_to_dict
+
+
+logger = logging.getLogger(__name__)
+
+
+class MultiLaneEmailBackend(BaseEmailBackend):
+    config = {'ignore_result': True, 'queue': 'celery'}
+    default_backend = 'django.core.mail.backends.smtp.EmailBackend'
+
+    def __init__(self, lane: str, fail_silently=False, **kwargs):
+        self.config.update(name=lane)
+        self.config.update(settings.TASK_CONFIG[lane])
+        self.task_kwargs = kwargs.copy()
+        # Make a copy tor ensure we don't modify input dict when we set the 'lane'
+        self.task_kwargs['debug'] = self.task_kwargs.pop('debug', {}).copy()
+        self.task_kwargs['debug']['lane'] = lane
+        super().__init__(fail_silently)
+
+    def send_messages(self, email_messages):
+        dict_messages = [email_to_dict(msg) for msg in email_messages]
+        TASKS[self.config['name']].delay(dict_messages, **self.task_kwargs)
+        return len(email_messages)
+
+    @staticmethod
+    def _run_task(messages, debug, **kwargs):
+        logger.warning('Sending queued email, details: %s', debug)
+        kwargs.setdefault('backend', kwargs.pop('backbackend', MultiLaneEmailBackend.default_backend))
+        with get_connection(**kwargs) as connection:
+            return connection.send_messages([dict_to_email(message) for message in messages])
+
+    @property
+    def task(self):
+        return shared_task(**self.config)(self._run_task)
+
+
+# Define tasks so that Celery can discovery them
+TASKS = {
+    name: MultiLaneEmailBackend(lane=name, fail_silently=True).task
+    for name in settings.TASK_CONFIG
+}

api/desecapi/models.py

@@ -16,7 +16,7 @@ import rest_framework.authtoken.models
 from django.conf import settings
 from django.contrib.auth.models import BaseUserManager, AbstractBaseUser, AnonymousUser
 from django.core.exceptions import ValidationError
-from django.core.mail import EmailMessage
+from django.core.mail import EmailMessage, get_connection
 from django.core.signing import Signer
 from django.core.validators import RegexValidator
 from django.db import models
@@ -140,30 +140,33 @@ class User(AbstractBaseUser):
         self.send_email('password-change-confirmation')
 
     def send_email(self, reason, context=None, recipient=None):
+        fast_lane = 'email_fast_lane'
+        slow_lane = 'email_slow_lane'
+        lanes = {
+            'activate': slow_lane,
+            'activate-with-domain': slow_lane,
+            'change-email': slow_lane,
+            'change-email-confirmation-old-email': fast_lane,
+            'password-change-confirmation': fast_lane,
+            'reset-password': fast_lane,
+            'delete-user': fast_lane,
+            'domain-dyndns': fast_lane,
+        }
+        if reason not in lanes:
+            raise ValueError(f'Cannot send email to user {self.pk} without a good reason: {reason}')
+
         context = context or {}
-        reasons = [
-            'activate',
-            'activate-with-domain',
-            'change-email',
-            'change-email-confirmation-old-email',
-            'password-change-confirmation',
-            'reset-password',
-            'delete-user',
-            'domain-dyndns',
-        ]
-        recipient = recipient or self.email
-        if reason not in reasons:
-            raise ValueError('Cannot send email to user {} without a good reason: {}'.format(self.pk, reason))
-        content_tmpl = get_template('emails/{}/content.txt'.format(reason))
-        subject_tmpl = get_template('emails/{}/subject.txt'.format(reason))
-        from_tmpl = get_template('emails/from.txt')
-        footer_tmpl = get_template('emails/footer.txt')
-        email = EmailMessage(subject_tmpl.render(context).strip(),
-                             content_tmpl.render(context) + footer_tmpl.render(),
-                             from_tmpl.render(),
-                             [recipient])
-        logger.warning('Sending email for user account %s (reason: %s)', str(self.pk), reason)
-        email.send()
+        content = get_template(f'emails/{reason}/content.txt').render(context)
+        footer = get_template('emails/footer.txt').render()
+
+        logger.warning(f'Queuing email for user account {self.pk} (reason: {reason})')
+        return EmailMessage(
+            subject=get_template(f'emails/{reason}/subject.txt').render(context).strip(),
+            body=content + footer,
+            from_email=get_template('emails/from.txt').render(),
+            to=[recipient or self.email],
+            connection=get_connection(lane=lanes[reason], debug={'user': self.pk, 'reason': reason})
+        ).send()
 
 
 class Token(rest_framework.authtoken.models.Token):

api/desecapi/tests/test_mail_backends.py

@@ -0,0 +1,30 @@
+from unittest import mock
+
+from django.conf import settings
+from django.core import mail
+from django.core.mail import EmailMessage, get_connection
+from django.test import TestCase
+
+from desecapi import mail_backends
+
+
+@mock.patch.dict(mail_backends.TASKS,
+                 {key: type('obj', (object,), {'delay': mail_backends.MultiLaneEmailBackend._run_task})
+                  for key in mail_backends.TASKS})
+class MultiLaneEmailBackendTestCase(TestCase):
+    test_backend = settings.EMAIL_BACKEND
+
+    def test_lanes(self):
+        debug_params = {'foo': 'bar'}
+        debug_params_orig = debug_params.copy()
+
+        with self.settings(EMAIL_BACKEND='desecapi.mail_backends.MultiLaneEmailBackend'):
+            for lane in ['email_slow_lane', 'email_fast_lane']:
+                subject = f'Test subject for lane {lane}'
+                connection = get_connection(lane=lane, backbackend=self.test_backend, debug=debug_params)
+                EmailMessage(subject=subject, to=['to@test.invalid'], connection=connection).send()
+                self.assertDictEqual(mail.outbox[-1].connection.task_kwargs['debug'], {'lane': lane, **debug_params})
+                self.assertEqual(mail.outbox[-1].subject, subject)
+
+        # Check that the backend hasn't modified the dict we passed
+        self.assertEqual(debug_params, debug_params_orig)