Changed indexing structure and logic to use a user's program enrollments

courses/factories.py

@@ -59,3 +59,12 @@ class CourseRunFactory(DjangoModelFactory):
 
     class Meta:  # pylint: disable=missing-docstring
         model = CourseRun
+
+    @classmethod
+    def create(cls, **kwargs):
+        # If 'program' is provided, build a Course from it first
+        if 'program' in kwargs:
+            kwargs['course'] = CourseFactory.create(program=kwargs.pop('program'))
+        # Create the CourseRun as normal
+        attrs = cls.attributes(create=True, extra=kwargs)
+        return cls._generate(True, attrs)

dashboard/factories.py

@@ -4,9 +4,7 @@
 from datetime import datetime, timedelta
 import pytz
 
-from factory import (
-    SubFactory,
-)
+from factory import SubFactory
 from factory.django import DjangoModelFactory
 from factory.fuzzy import (
     FuzzyAttribute,
@@ -17,8 +15,13 @@
 from dashboard.models import (
     CachedCertificate,
     CachedEnrollment,
+    ProgramEnrollment
+)
+from courses.factories import (
+    CourseRunFactory,
+    CourseFactory,
+    ProgramFactory,
 )
-from courses.factories import CourseRunFactory
 from profiles.factories import UserFactory
 
 
@@ -47,3 +50,32 @@ class CachedEnrollmentFactory(DjangoModelFactory):
 
     class Meta:  # pylint: disable=missing-docstring,no-init,too-few-public-methods,old-style-class
         model = CachedEnrollment
+
+
+class ProgramEnrollmentFactory(DjangoModelFactory):
+    """Factory for ProgramEnrollment"""
+    user = SubFactory(UserFactory)
+    program = SubFactory(ProgramFactory)
+
+    class Meta:  # pylint: disable=missing-docstring,no-init,too-few-public-methods,old-style-class
+        model = ProgramEnrollment
+
+    @classmethod
+    def create(cls, **kwargs):
+        """
+        Overrides default ProgramEnrollment object creation for the factory.
+
+        ProgramEnrollments should only exist if there is a CachedEnrollment associated with
+        the given User and Program. Instead of creating a new record with the factory, we
+        will create the necessary objects to trigger its creation.
+        """
+        user = kwargs.get('user', UserFactory.create())
+        program = kwargs.get('program', ProgramFactory.create())
+        course = CourseFactory.create(program=program)
+        course_run = CourseRunFactory.create(course=course)
+        CachedEnrollmentFactory.create(user=user, course_run=course_run)
+        # CachedCertificate isn't strictly necessary to create a ProgramEnrollment. This is here for test convenience.
+        CachedCertificateFactory.create(user=user, course_run=course_run)
+        # Signal from the creation of a CachedEnrollment should have created a ProgramEnrollment
+        program_enrollment = ProgramEnrollment.objects.get(user=user, program=program)
+        return program_enrollment

dashboard/models.py

@@ -25,6 +25,16 @@ class CachedEnrollment(Model):
     class Meta:
         unique_together = (('user', 'course_run'), )
 
+    @classmethod
+    def active_count(cls, user, program):
+        """
+        Returns the number of active CachedEnrollments for a User/Program pair
+        """
+        return cls.objects.filter(
+            user=user,
+            course_run__course__program=program
+        ).exclude(data__isnull=True).count()
+
     def __str__(self):
         """
         String representation of the model object

dashboard/signals.py

@@ -1,10 +1,27 @@
 """
 Signals for user profiles
 """
-from django.db.models.signals import pre_save, post_save
+from django.db.models.signals import pre_save, post_save, post_delete
 from django.dispatch import receiver
 
-from dashboard.models import CachedEnrollment, ProgramEnrollment
+from dashboard.models import CachedEnrollment, CachedCertificate, ProgramEnrollment
+from search.tasks import index_program_enrolled_users, index_users, remove_program_enrolled_user
+
+
+@receiver(post_save, sender=ProgramEnrollment, dispatch_uid="programenrollment_post_save")
+def handle_create_programenrollment(sender, instance, created, **kwargs):  # pylint: disable=unused-argument
+    """
+    When a ProgramEnrollment model is created/updated, update index.
+    """
+    index_program_enrolled_users.delay([instance])
+
+
+@receiver(post_delete, sender=ProgramEnrollment, dispatch_uid="programenrollment_post_delete")
+def handle_delete_programenrollment(sender, instance, **kwargs):  # pylint: disable=unused-argument
+    """
+    When a ProgramEnrollment model is deleted, update index.
+    """
+    remove_program_enrolled_user.delay(instance)
 
 
 @receiver(pre_save, sender=CachedEnrollment, dispatch_uid="preupdate_programenrollment")
@@ -25,27 +42,65 @@ def precreate_programenrollment(sender, instance, **kwargs):  # pylint: disable=
         instance_in_db = CachedEnrollment.objects.filter(id=instance.id).exclude(data__isnull=True).count()
         # if the count is 1, it means the student unenrolled from the course run
         if instance_in_db == 1:
-            active_enrollment_count = CachedEnrollment.objects.filter(
-                user=user,
-                course_run__course__program=program
-            ).exclude(data__isnull=True).count()
             # if there is only one enrollment with data non None, it means that it is the
             # current instance is the only one for the program, so the program enrollment
             # needs to be deleted
-            if active_enrollment_count <= 1:  # theoretically this cannot be <1, but just in case
+            if CachedEnrollment.active_count(user, program) <= 1:  # theoretically this cannot be <1, but just in case
                 ProgramEnrollment.objects.filter(
                     user=user,
                     program=program
                 ).delete()
 
 
-@receiver(post_save, sender=CachedEnrollment, dispatch_uid="update_programenrollment")
-def create_programenrollment(sender, instance, **kwargs):  # pylint: disable=unused-argument
+@receiver(post_save, sender=CachedEnrollment, dispatch_uid="cachedenrollment_post_save")
+def handle_update_enrollment(sender, instance, **kwargs):  # pylint: disable=unused-argument
+    """
+    Create ProgramEnrollment when a CachedEnrollment is created/updated, and update the index.
+    """
+    if instance.data is not None:
+        program_enrollment, _ = ProgramEnrollment.objects.get_or_create(
+            user=instance.user,
+            program=instance.course_run.course.program
+        )
+        index_program_enrolled_users.delay([program_enrollment])
+
+
+@receiver(post_save, sender=CachedCertificate, dispatch_uid="cachedcertificate_post_save")
+def handle_update_certificate(sender, instance, **kwargs):  # pylint: disable=unused-argument
     """
-    Signal handler to create Program enrollment when the CachedEnrollment table is updated
+    When a CachedCertificate model is updated, update index.
     """
     if instance.data is not None:
-        ProgramEnrollment.objects.get_or_create(
+        program_enrollment, _ = ProgramEnrollment.objects.get_or_create(
             user=instance.user,
             program=instance.course_run.course.program
         )
+        index_program_enrolled_users.delay([program_enrollment])
+
+
+@receiver(post_delete, sender=CachedEnrollment, dispatch_uid="cachedenrollment_post_delete")
+def handle_delete_enrollment(sender, instance, **kwargs):  # pylint: disable=unused-argument
+    """
+    Update index when CachedEnrollment model instance is deleted.
+    """
+    user = instance.user
+    program = instance.course_run.course.program
+    program_enrollment = ProgramEnrollment.objects.filter(user=user, program=program).first()
+    if program_enrollment is not None:
+        if CachedEnrollment.active_count(user, program) == 0:
+            program_enrollment.delete()
+            index_users.delay([user])
+        else:
+            index_program_enrolled_users.delay([program_enrollment])
+
+
+@receiver(post_delete, sender=CachedCertificate, dispatch_uid="cachedcertificate_post_delete")
+def handle_delete_certificate(sender, instance, **kwargs):  # pylint: disable=unused-argument
+    """
+    Update index when CachedCertificate model instance is deleted.
+    """
+    user = instance.user
+    program = instance.course_run.course.program
+    program_enrollment = ProgramEnrollment.objects.filter(user=user, program=program).first()
+    if program_enrollment is not None:
+        index_program_enrolled_users.delay([program_enrollment])

micromasters/utils.py

@@ -19,6 +19,15 @@ def webpack_dev_server_url(request):
     return 'http://{}:{}'.format(webpack_dev_server_host(request), settings.WEBPACK_DEV_SERVER_PORT)
 
 
+def dict_without_key(dictionary, key):
+    """
+    Helper method to remove a key from a dict and return the dict. This can be used in cases like a list comprehension
+    where the actual dictionary is needed once the key is deleted ('del' does not return anything)
+    """
+    del dictionary[key]
+    return dictionary
+
+
 def load_json_from_file(project_rel_filepath):
     """
     Loads JSON data from a file