feat: add ability to override middlewares for recurring nudges (#790)

* feat: add ability to override middlewares for recurring nudges

* feat: add ability to run command for all sites

Author: kaustavb12

openedx/core/djangoapps/schedules/management/commands/__init__.py

@@ -29,12 +29,28 @@ def add_arguments(self, parser):
             '--override-recipient-email',
             help='Send all emails to this address instead of the actual recipient'
         )
-        parser.add_argument('site_domain_name')
+        parser.add_argument(
+            'site_domain_name',
+            nargs='?',
+            default=None,
+            help=(
+                'Domain name for the site to use. '
+                'Do not provide a domain if you wish to run this for all sites'
+            )
+        )
         parser.add_argument(
             '--weeks',
             type=int,
             help='Number of weekly emails to be sent',
         )
+        parser.add_argument(
+            '--override-middlewares',
+            action='append',
+            help=(
+                'Use this middleware when emulating http requests. '
+                'To use multiple middlewares, provide this argument multiple times'
+            )
+        )
 
     def handle(self, *args, **options):
         self.log_debug('Args = %r', options)
@@ -49,19 +65,26 @@ def handle(self, *args, **options):
             tzinfo=pytz.UTC
         )
         self.log_debug('Current date = %s', current_date.isoformat())
+        override_recipient_email = options.get('override_recipient_email')
+        override_middlewares = options.get('override_middlewares')
 
-        site = Site.objects.get(domain__iexact=options['site_domain_name'])
-        self.log_debug('Running for site %s', site.domain)
+        site_domain_name = options['site_domain_name']
+        sites = Site.objects.filter(domain__iexact=site_domain_name) if site_domain_name else Site.objects.all()
 
-        override_recipient_email = options.get('override_recipient_email')
-        self.send_emails(site, current_date, override_recipient_email)
+        if sites:
+            for site in sites:
+                self.log_debug('Running for site %s', site.domain)
+                self.send_emails(site, current_date, override_recipient_email, override_middlewares)
+        else:
+            self.log_info("No matching site found")
 
-    def enqueue(self, day_offset, site, current_date, override_recipient_email=None):
+    def enqueue(self, day_offset, site, current_date, override_recipient_email=None, override_middlewares=None):
         self.async_send_task.enqueue(
             site,
             current_date,
             day_offset,
             override_recipient_email,
+            override_middlewares,
         )
 
     def send_emails(self, *args, **kwargs):

openedx/core/djangoapps/schedules/management/commands/tests/test_send_email_base_command.py

@@ -10,6 +10,7 @@
 import ddt
 import pytz
 from django.conf import settings
+from django.contrib.sites.models import Site
 
 from openedx.core.djangoapps.schedules.management.commands import SendEmailBaseCommand
 from openedx.core.djangoapps.site_configuration.tests.factories import SiteConfigurationFactory, SiteFactory
@@ -33,9 +34,23 @@ def test_handle(self):
             send_emails.assert_called_once_with(
                 self.site,
                 datetime.datetime(2017, 9, 29, tzinfo=pytz.UTC),
+                None,
                 None
             )
 
+    def test_handle_all_sites(self):
+        with patch.object(self.command, 'send_emails') as send_emails:
+            self.command.handle(site_domain_name=None, date='2017-09-29')
+            expected_sites = Site.objects.all()
+            for expected_site in expected_sites:
+                send_emails.assert_any_call(
+                    expected_site,
+                    datetime.datetime(2017, 9, 29, tzinfo=pytz.UTC),
+                    None,
+                    None
+                )
+            assert send_emails.call_count == len(expected_sites)
+
     def test_weeks_option(self):
         with patch.object(self.command, 'enqueue') as enqueue:
             self.command.handle(site_domain_name=self.site.domain, date='2017-09-29', weeks=12)

openedx/core/djangoapps/schedules/tasks.py

@@ -20,6 +20,7 @@
     set_custom_attribute
 )
 from eventtracking import tracker
+from importlib import import_module
 from opaque_keys.edx.keys import CourseKey
 
 from openedx.core.djangoapps.content.course_overviews.models import CourseOverview
@@ -103,7 +104,7 @@ class BinnedScheduleMessageBaseTask(ScheduleMessageBaseTask):
     task_instance = None
 
     @classmethod
-    def enqueue(cls, site, current_date, day_offset, override_recipient_email=None):  # lint-amnesty, pylint: disable=missing-function-docstring
+    def enqueue(cls, site, current_date, day_offset, override_recipient_email=None, override_middlewares=None):  # lint-amnesty, pylint: disable=missing-function-docstring
         set_code_owner_attribute_from_module(__name__)
         current_date = resolvers._get_datetime_beginning_of_day(current_date)  # lint-amnesty, pylint: disable=protected-access
 
@@ -120,6 +121,7 @@ def enqueue(cls, site, current_date, day_offset, override_recipient_email=None):
                 day_offset,
                 bin,
                 override_recipient_email,
+                override_middlewares,
             )
             cls.log_info('Launching task with args = %r', task_args)
             cls.task_instance.apply_async(
@@ -128,16 +130,17 @@ def enqueue(cls, site, current_date, day_offset, override_recipient_email=None):
             )
 
     def run(  # lint-amnesty, pylint: disable=arguments-differ
-        self, site_id, target_day_str, day_offset, bin_num, override_recipient_email=None,
+        self, site_id, target_day_str, day_offset, bin_num, override_recipient_email=None, override_middlewares=None,
     ):
         set_code_owner_attribute_from_module(__name__)
         site = Site.objects.select_related('configuration').get(id=site_id)
-        with emulate_http_request(site=site):
+        middlewares = [self.class_from_classpath(cls) for cls in override_middlewares] if override_middlewares else None
+        with emulate_http_request(site=site, middleware_classes=middlewares) as request:
             msg_type = self.make_message_type(day_offset)
-            _annotate_for_monitoring(msg_type, site, bin_num, target_day_str, day_offset)
+            _annotate_for_monitoring(msg_type, request.site, bin_num, target_day_str, day_offset)
             return self.resolver(  # lint-amnesty, pylint: disable=not-callable
                 self.async_send_task,
-                site,
+                request.site,
                 deserialize(target_day_str),
                 day_offset,
                 bin_num,
@@ -147,6 +150,11 @@ def run(  # lint-amnesty, pylint: disable=arguments-differ
     def make_message_type(self, day_offset):
         raise NotImplementedError
 
+    def class_from_classpath(self, class_path):
+        module_name, klass = class_path.rsplit('.', 1)
+        module = import_module(module_name)
+        return getattr(module, klass)
+
 
 @shared_task(base=LoggedTask, ignore_result=True)
 @set_code_owner_attribute

openedx/core/lib/celery/task_utils.py

@@ -45,7 +45,7 @@ def emulate_http_request(site=None, user=None, middleware_classes=None):
         _run_method_if_implemented(middleware, 'process_request', request)
 
     try:
-        yield
+        yield request
     except Exception as exc:
         for middleware in reversed(middleware_instances):
             _run_method_if_implemented(middleware, 'process_exception', request, exc)