[ENG-12603] Split API keys between standard notifications and ENTER ones

osf/email/__init__.py

@@ -4,6 +4,7 @@
 import logging
 import importlib
 import sys
+from functools import cache
 from html import unescape
 from typing import Optional
 from mako.template import Template as MakoTemplate
@@ -285,7 +286,22 @@ def personalization(recipients, custom_args=None):
     ]
 
 
-def send_email_with_send_grid(to_addr, notification_type, context, email_context=None, *, is_multiple=False, rendered_html=None):
+@cache
+def _sendgrid_client(api_key):
+    return SendGridAPIClient(api_key)
+
+
+def get_sendgrid_client(use_webhook=False):
+    """Return a per-process cached SendGrid client.
+
+    When ``use_webhook`` is True, the client authenticates as the sub-user with Event Webhook enabled
+    (``SENDGRID_API_KEY_WITH_WEBHOOK``), otherwise as the default user (``SENDGRID_API_KEY``).
+    """
+    api_key = settings.SENDGRID_API_KEY_WITH_WEBHOOK if use_webhook else None
+    return _sendgrid_client(api_key or settings.SENDGRID_API_KEY)
+
+
+def send_email_with_send_grid(to_addr, notification_type, context, email_context=None, *, is_multiple=False, rendered_html=None, client=None):
 
     email_context = email_context or {}
     to_list = [to_addr] if isinstance(to_addr, str) else [a for a in (to_addr or []) if a]
@@ -337,7 +353,7 @@ def send_email_with_send_grid(to_addr, notification_type, context, email_context
 
             payload['attachments'] = [item]
     try:
-        sg = SendGridAPIClient(settings.SENDGRID_API_KEY)
+        sg = client or SendGridAPIClient(settings.SENDGRID_API_KEY)
         resp = sg.client.mail.send.post(request_body=payload)
         if resp.status_code not in (200, 201, 202):
             logging.error(
@@ -380,9 +396,10 @@ def send_email_with_send_grid(to_addr, notification_type, context, email_context
             logging.error('SendGrid hit a blocked socket error: %r | payload=%s', exc, payload)
         raise
 
-def send_email(recipient_address, notification_type, event_context=None, email_context=None, rendered_html=None):
+def send_email(recipient_address, notification_type, event_context=None, email_context=None, rendered_html=None, client=None):
     """
     Send an email using either SMTP or SendGrid based on settings and feature flags.
+    ``client`` is an optional SendGrid client to reuse; a default one is created per call otherwise.
     """
     if waffle.switch_is_active(features.ENABLE_MAILHOG):
         send_email_over_smtp(
@@ -400,6 +417,7 @@ def send_email(recipient_address, notification_type, event_context=None, email_c
             event_context,
             email_context,
             rendered_html=rendered_html,
+            client=client,
         )
 
     if settings.LOCAL_MODE and not waffle.switch_is_active(features.ENABLE_MAILHOG):

osf/email/notification_campaign/execution.py

@@ -6,7 +6,7 @@
 
 from framework import sentry
 from framework.celery_tasks import app as celery_app
-from osf.email import _render_email_html, send_email, send_email_with_send_grid
+from osf.email import _render_email_html, get_sendgrid_client, send_email, send_email_with_send_grid
 from osf.models import NotificationType, NotificationTypeEnum
 from osf.models.notification_campaign import (
     NotificationCampaign,
@@ -23,7 +23,7 @@
 
 
 @celery_app.task(bind=True, base=NotificationCampaignTask, name='email.process_campaign_retry')
-def process_campaign_retry(self, campaign_id, run_id):
+def process_campaign_retry(self, campaign_id, run_id, use_sendgrid_webhook=False):
     campaign = self.get_campaign(campaign_id, run_id, abort_if_cancelled=False)
     if campaign is None:
         return
@@ -49,7 +49,11 @@ def process_campaign_retry(self, campaign_id, run_id):
             # Still waiting for in-flight sends / SendGrid delivery webhooks.
             self.sync_campaign_stats(campaign)
             process_campaign_retry.apply_async(
-                kwargs={'campaign_id': campaign_id, 'run_id': campaign.run_id},
+                kwargs={
+                    'campaign_id': campaign_id,
+                    'run_id': campaign.run_id,
+                    'use_sendgrid_webhook': use_sendgrid_webhook,
+                },
                 countdown=execution.get('dispatch_interval', settings.CAMPAIGN_DISPATCH_INTERVAL),
             )
             return
@@ -90,6 +94,7 @@ def process_campaign_retry(self, campaign_id, run_id):
                 args=[campaign_id, campaign.run_id],
                 kwargs={
                     'restart_failed': True,
+                    'use_sendgrid_webhook': use_sendgrid_webhook,
                 },
             )
             return
@@ -118,16 +123,22 @@ def start_notification_campaign(self, campaign_id, restart_failed=False, restart
             status=NotificationCampaignRecipientStatus.QUEUED,
         ).update(status=NotificationCampaignRecipientStatus.PENDING, batch_id=None)
 
+    # Decide once per run which SendGrid user sends the campaign: delivery tracked via
+    # webhook requires the sub-user with Event Webhook enabled.
+    use_sendgrid_webhook = bool(
+        campaign.metadata.get('sendgrid_bulk', False) or settings.CAMPAIGN_SINGLE_SEND_WEBHOOK
+    )
     dispatch_campaign.apply_async(
         args=[campaign_id, campaign.run_id],
         kwargs={
             'restart_failed': restart_failed,
+            'use_sendgrid_webhook': use_sendgrid_webhook,
         },
     )
 
 
 @celery_app.task(bind=True, base=NotificationCampaignTask, name='email.dispatch_campaign')
-def dispatch_campaign(self, campaign_id, run_id, restart_failed=False, restart_stuck=False):
+def dispatch_campaign(self, campaign_id, run_id, restart_failed=False, restart_stuck=False, use_sendgrid_webhook=False):
     campaign = self.get_campaign(campaign_id, run_id)
     if campaign is None:
         return
@@ -196,6 +207,7 @@ def dispatch_campaign(self, campaign_id, run_id, restart_failed=False, restart_s
                 notification_type_name=notification_type_name,
                 run_id=run_id,
                 developer_reminder=developer_reminder,
+                use_sendgrid_webhook=use_sendgrid_webhook,
             )
             new_queued_batches += 1
 
@@ -209,13 +221,14 @@ def dispatch_campaign(self, campaign_id, run_id, restart_failed=False, restart_s
 
     if no_more_recipients:
         process_campaign_retry.delay(
-            campaign_id=campaign_id, run_id=campaign.run_id
+            campaign_id=campaign_id, run_id=campaign.run_id, use_sendgrid_webhook=use_sendgrid_webhook
         )
     else:
         self.apply_async(
             args=[campaign_id, run_id],
             kwargs={
                 'restart_failed': restart_failed,
+                'use_sendgrid_webhook': use_sendgrid_webhook,
             },
             countdown=execution.get('dispatch_interval', settings.CAMPAIGN_DISPATCH_INTERVAL),
         )
@@ -230,6 +243,7 @@ def send_campaign_batch(
     campaign_id=None,
     run_id=None,
     developer_reminder=False,
+    use_sendgrid_webhook=False,
 ):
     campaign = self.get_campaign(campaign_id, run_id)
     if campaign is None:
@@ -286,6 +300,8 @@ def send_campaign_batch(
     invalid_emails_qs = recipients_qs_annotated.filter(recipient_address__isnull=True)
     invalid_emails_qs.update(status=NotificationCampaignRecipientStatus.SKIPPED, error_message='Invalid email address')
 
+    sendgrid_client = get_sendgrid_client(use_webhook=use_sendgrid_webhook)
+
     if campaign.metadata.get('sendgrid_bulk', False):
         recipients = list(valid_emails_qs)
         recipient_emails = []
@@ -307,6 +323,7 @@ def send_campaign_batch(
                 context=context,
                 email_context={'custom_args_list': custom_args_list},
                 is_multiple=True,
+                client=sendgrid_client,
             )
             valid_emails_qs.update(
                 status=NotificationCampaignRecipientStatus.AWAITING_DELIVERY,
@@ -336,8 +353,9 @@ def send_campaign_batch(
                         },
                     },
                     rendered_html=rendered_html,
+                    client=sendgrid_client,
                 )
-                if settings.CAMPAIGN_SINGLE_SEND_WEBHOOK:
+                if use_sendgrid_webhook:
                     recipient.status = NotificationCampaignRecipientStatus.AWAITING_DELIVERY
                 else:
                     recipient.status = NotificationCampaignRecipientStatus.SENT

osf_tests/test_notification_campaign.py

@@ -6,12 +6,13 @@
 from django.utils import timezone
 from django.db.models import Q
 
-from osf.email import _build_sendgrid_personalizations
+from osf.email import _build_sendgrid_personalizations, _sendgrid_client, get_sendgrid_client
 from osf.email.notification_campaign import (
     NotificationCampaignTask,
     assign_batch_id_to_recipients,
     build_campaign_filter_query,
     create_campaign_recipients,
+    dispatch_campaign,
     get_campaign_recipient_stats,
     process_campaign_retry,
     process_sendgrid_campaign_events,
@@ -730,6 +731,25 @@ def test_start_schedules_workflow(self, mock_dispatch_campaign, campaign):
         assert campaign.recipient_count == 3
         mock_dispatch_campaign.assert_called_once()
 
+    @pytest.mark.parametrize('sendgrid_bulk, single_send_webhook, expected', [
+        (False, False, False),
+        (True, False, True),
+        (False, True, True),
+    ])
+    @mock.patch('osf.email.notification_campaign.execution.dispatch_campaign.apply_async')
+    def test_start_decides_sendgrid_webhook_client(
+        self, mock_dispatch_campaign, campaign, sendgrid_bulk, single_send_webhook, expected
+    ):
+        campaign.metadata['sendgrid_bulk'] = sendgrid_bulk
+        campaign.run_id = uuid.uuid4()
+        campaign.save()
+
+        with mock.patch('website.settings.CAMPAIGN_SINGLE_SEND_WEBHOOK', single_send_webhook):
+            start_notification_campaign(campaign.id)
+
+        mock_dispatch_campaign.assert_called_once()
+        assert mock_dispatch_campaign.call_args.kwargs['kwargs']['use_sendgrid_webhook'] is expected
+
     @mock.patch('osf.email.notification_campaign.execution.dispatch_campaign')
     def test_start_excludes_unconfirmed_accounts_when_enabled(self, mock_dispatch_campaign, campaign):
         confirmed = UserFactory()
@@ -823,6 +843,28 @@ def test_start_restart_stuck_resets_only_queued(self, mock_dispatch_campaign, ca
         assert campaign.recipient_count == 2
 
 
+class TestDispatchCampaign:
+
+    @mock.patch('osf.email.notification_campaign.execution.process_campaign_retry.delay')
+    @mock.patch('osf.email.notification_campaign.execution.send_campaign_batch.delay')
+    def test_dispatch_passes_sendgrid_webhook_flag_down(self, mock_send_batch, mock_retry, campaign):
+        user = UserFactory()
+        create_campaign_recipients(Q(**{'id__in': [user.id]}), campaign_id=campaign.id)
+        campaign.run_id = uuid.uuid4()
+        campaign.status = NotificationCampaignStatus.RUNNING
+        campaign.save()
+
+        dispatch_campaign(campaign.id, campaign.run_id, use_sendgrid_webhook=True)
+
+        mock_send_batch.assert_called_once()
+        assert mock_send_batch.call_args.kwargs['use_sendgrid_webhook'] is True
+        mock_retry.assert_called_once_with(
+            campaign_id=campaign.id,
+            run_id=campaign.run_id,
+            use_sendgrid_webhook=True,
+        )
+
+
 class TestSendCampaignBatch:
 
     @pytest.fixture
@@ -934,6 +976,48 @@ def test_send_campaign_batch_sendgrid_bulk_success(self, mock_sendgrid, running_
         assert get_campaign_recipient_stats(running_campaign.id)['awaiting_count'] == 1
         assert get_campaign_recipient_stats(running_campaign.id)['queued_count'] == 0
 
+    @mock.patch('osf.email.notification_campaign.execution.get_sendgrid_client')
+    @mock.patch('osf.email.notification_campaign.execution.send_email_with_send_grid')
+    def test_send_campaign_batch_sendgrid_bulk_uses_webhook_client(self, mock_sendgrid, mock_get_client, running_campaign):
+        user = UserFactory()
+        running_campaign.metadata['sendgrid_bulk'] = True
+        running_campaign.save(update_fields=['metadata'])
+        create_campaign_recipients(Q(**{'id__in': [user.id]}), campaign_id=running_campaign.id)
+        batch_id = assign_batch_id_to_recipients(campaign_id=running_campaign.id, batch_size=1)
+
+        send_campaign_batch(
+            context={},
+            batch_id=batch_id,
+            notification_type_name='blank',
+            campaign_id=running_campaign.id,
+            run_id=running_campaign.run_id,
+            use_sendgrid_webhook=True,
+        )
+
+        mock_get_client.assert_called_once_with(use_webhook=True)
+        assert mock_sendgrid.call_args.kwargs['client'] is mock_get_client.return_value
+
+    @mock.patch('osf.email.notification_campaign.execution.get_sendgrid_client')
+    @mock.patch('osf.email.notification_campaign.execution.send_email')
+    def test_send_campaign_batch_single_send_reuses_client_for_batch(self, mock_send_email, mock_get_client, running_campaign):
+        users = [UserFactory(), UserFactory()]
+        create_campaign_recipients(Q(**{'id__in': [u.id for u in users]}), campaign_id=running_campaign.id)
+        batch_id = assign_batch_id_to_recipients(campaign_id=running_campaign.id, batch_size=2)
+
+        send_campaign_batch(
+            context={},
+            batch_id=batch_id,
+            notification_type_name='blank',
+            campaign_id=running_campaign.id,
+            run_id=running_campaign.run_id,
+            use_sendgrid_webhook=True,
+        )
+
+        mock_get_client.assert_called_once_with(use_webhook=True)
+        assert mock_send_email.call_count == 2
+        assert all(c.kwargs['client'] is mock_get_client.return_value for c in mock_send_email.call_args_list)
+        assert get_campaign_recipient_stats(running_campaign.id)['awaiting_count'] == 2
+
     def test_build_sendgrid_personalizations_shared_to_list(self):
         personalizations = _build_sendgrid_personalizations(
             ['a@example.com', 'b@example.com'],
@@ -1173,6 +1257,36 @@ def test_send_campaign_batch_uses_fallback_email_when_username_has_no_at(self, m
         assert get_campaign_recipient_stats(running_campaign.id)['failed_count'] == 0
         mock_send_email.assert_called_once()
 
+class TestGetSendgridClient:
+
+    @pytest.fixture(autouse=True)
+    def clear_client_cache(self):
+        _sendgrid_client.cache_clear()
+        yield
+        _sendgrid_client.cache_clear()
+
+    @mock.patch('website.settings.SENDGRID_API_KEY_WITH_WEBHOOK', 'webhook-key')
+    @mock.patch('website.settings.SENDGRID_API_KEY', 'default-key')
+    def test_selects_api_key_by_webhook_flag(self):
+        with mock.patch('osf.email.SendGridAPIClient') as mock_client_cls:
+            get_sendgrid_client(use_webhook=True)
+            get_sendgrid_client(use_webhook=False)
+
+        assert mock_client_cls.call_args_list == [mock.call('webhook-key'), mock.call('default-key')]
+
+    @mock.patch('website.settings.SENDGRID_API_KEY_WITH_WEBHOOK', None)
+    @mock.patch('website.settings.SENDGRID_API_KEY', 'default-key')
+    def test_webhook_falls_back_to_default_api_key(self):
+        with mock.patch('osf.email.SendGridAPIClient') as mock_client_cls:
+            get_sendgrid_client(use_webhook=True)
+
+        mock_client_cls.assert_called_once_with('default-key')
+
+    @mock.patch('website.settings.SENDGRID_API_KEY', 'default-key')
+    def test_client_is_cached_per_api_key(self):
+        assert get_sendgrid_client() is get_sendgrid_client()
+
+
 class TestNotificationCampaignCancel:
 
     def test_cancel_sets_cancelled_status_and_completed_at(self, campaign):
@@ -1329,7 +1443,7 @@ def test_process_campaign_retry_waits_for_awaiting_delivery_recipients(self, moc
         assert campaign.sent_count == 0
         assert campaign.failed_count == 0
         mock_apply_async.assert_called_once_with(
-            kwargs={'campaign_id': campaign.id, 'run_id': campaign.run_id},
+            kwargs={'campaign_id': campaign.id, 'run_id': campaign.run_id, 'use_sendgrid_webhook': False},
             countdown=60,
         )
 

tests/utils.py

@@ -331,7 +331,7 @@ def _fake_send_over_smtp(to_email, notification_type, context=None, email_contex
             if passthrough:
                 return _real_send_over_smtp(to_email, notification_type, context, email_context, rendered_html)
 
-        def _fake_send_with_sendgrid(user, notification_type, context=None, email_context=None, rendered_html=None):
+        def _fake_send_with_sendgrid(user, notification_type, context=None, email_context=None, rendered_html=None, client=None):
             captured['emails'].append({
                 'protocol': 'sendgrid',
                 'to': user,
@@ -340,7 +340,9 @@ def _fake_send_with_sendgrid(user, notification_type, context=None, email_contex
                 'email_context': email_context.copy() if isinstance(email_context, dict) else email_context,
             })
             if passthrough:
-                return _real_send_with_sendgrid(user, notification_type, context, email_context, rendered_html)
+                return _real_send_with_sendgrid(
+                    user, notification_type, context, email_context, rendered_html=rendered_html, client=client
+                )
 
         patches.extend([
             mock.patch('osf.email.send_email_over_smtp', new=_fake_send_over_smtp),

website/settings/defaults.py

@@ -186,7 +186,10 @@ def parent_dir(path):
 
 # OR, if using Sendgrid's API
 # WARNING: If `SENDGRID_WHITELIST_MODE` is True,
+# API key of either the primary/default SendGrid user or a different sub-suer, both without Event Webhook enabled
 SENDGRID_API_KEY = None
+# API key of the SendGrid sub-user with Event Webhook enabled
+SENDGRID_API_KEY_WITH_WEBHOOK = None
 # Public verification key from SendGrid Event Webhook (Mail Settings -> Event Webhook -> Signed Event Webhook)
 SENDGRID_EVENT_WEBHOOK_PUBLIC_KEY = None