Fix cohort Integrity Errors

TNL-5725

lms/djangoapps/verified_track_content/tasks.py

@@ -2,7 +2,6 @@
 Celery task for Automatic Verifed Track Cohorting MVP feature.
 """
 from django.contrib.auth.models import User
-from django.db.utils import IntegrityError
 
 from celery.task import task
 from celery.utils.log import get_task_logger
@@ -31,18 +30,7 @@ def sync_cohort_with_mode(self, course_id, user_id, verified_cohort_name, defaul
         enrollment = CourseEnrollment.get_enrollment(user, course_key)
         # Note that this will enroll the user in the default cohort on initial enrollment.
         # That's good because it will force creation of the default cohort if necessary.
-        try:
-            current_cohort = get_cohort(user, course_key)
-        except IntegrityError as integrity_error:
-            # It is quite common that get_cohort will throw an IntegrityError. This happens
-            # when 2 celery workers are both handling enrollment change events for the same user
-            # (for example, if the enrollment mode goes from None -> Audit -> Honor); if the user
-            # was not previously in a cohort, calling get_cohort will result in a cohort assignment.
-            LOGGER.info(
-                "HANDLING_INTEGRITY_ERROR: IntegrityError encountered for course '%s' and user '%s': %s",
-                course_id, user.id, unicode(integrity_error)
-            )
-            current_cohort = get_cohort(user, course_key)
+        current_cohort = get_cohort(user, course_key)
 
         verified_cohort = get_cohort_by_name(course_key, verified_cohort_name)
 

openedx/core/djangoapps/course_groups/cohorts.py

@@ -6,6 +6,7 @@ forums, and to the cohort admin views.
 import logging
 import random
 
+from django.db import IntegrityError, transaction
 from django.db.models.signals import post_save, m2m_changed
 from django.dispatch import receiver
 from django.http import Http404
@@ -194,11 +195,22 @@ def get_cohort(user, course_key, assign=True, use_cached=False):
             return None
 
     # Otherwise assign the user a cohort.
-    membership = CohortMembership.objects.create(
-        user=user,
-        course_user_group=get_random_cohort(course_key)
-    )
-    return request_cache.data.setdefault(cache_key, membership.course_user_group)
+    try:
+        with transaction.atomic():
+            membership = CohortMembership.objects.create(
+                user=user,
+                course_user_group=get_random_cohort(course_key)
+            )
+            return request_cache.data.setdefault(cache_key, membership.course_user_group)
+    except IntegrityError as integrity_error:
+        # An IntegrityError is raised when multiple workers attempt to
+        # create the same row in one of the cohort model entries:
+        # CourseCohort, CohortMembership.
+        log.info(
+            "HANDLING_INTEGRITY_ERROR: IntegrityError encountered for course '%s' and user '%s': %s",
+            course_key, user.id, unicode(integrity_error)
+        )
+        return get_cohort(user, course_key, assign, use_cached)
 
 
 def get_random_cohort(course_key):