Merge pull request #10469 from edx/efischer/hotfix_management_command

Post-migration management command

openedx/core/djangoapps/course_groups/management/commands/post_cohort_membership_fix.py

+"""
+Intended to fix any inconsistencies that may arise during the rollout of the CohortMembership model.
+Illustration: https://gist.github.com/efischer19/d62f8ee42b7fbfbc6c9a
+"""
+from django.core.management.base import BaseCommand
+from django.db import IntegrityError
+
+from openedx.core.djangoapps.course_groups.models import CourseUserGroup, CohortMembership
+
+
+class Command(BaseCommand):
+    """
+    Repair any inconsistencies between CourseUserGroup and CohortMembership. To be run after migration 0006.
+    """
+    help = '''
+    Repairs any potential inconsistencies made in the window between running migrations 0005 and 0006, and deploying
+    the code changes to enforce use of CohortMembership that go with said migrations.
+    |commit|: optional argument. If not provided, will dry-run and list number of operations that would be made.
+    '''
+
+    def handle(self, *args, **options):
+        """
+        Execute the command. Since this is designed to fix any issues cause by running pre-CohortMembership code
+        with the database already migrated to post-CohortMembership state, we will use the pre-CohortMembership
+        table CourseUserGroup as the canonical source of truth. This way, changes made in the window are persisted.
+        """
+        commit = False
+        if len(args) == 1:
+            commit = args[0] == 'commit'
+
+        memberships_to_delete = 0
+        memberships_to_add = 0
+
+        # Begin by removing any data in CohortMemberships that does not match CourseUserGroups data
+        for membership in CohortMembership.objects.all():
+            try:
+                CourseUserGroup.objects.get(
+                    group_type=CourseUserGroup.COHORT,
+                    users__id=membership.user.id,
+                    course_id=membership.course_id,
+                    id=membership.course_user_group.id
+                )
+            except CourseUserGroup.DoesNotExist:
+                memberships_to_delete += 1
+                if commit:
+                    membership.delete()
+
+        # Now we can add any CourseUserGroup data that is missing a backing CohortMembership
+        for course_group in CourseUserGroup.objects.filter(group_type=CourseUserGroup.COHORT):
+            for user in course_group.users.all():
+                try:
+                    CohortMembership.objects.get(
+                        user=user,
+                        course_id=course_group.course_id,
+                        course_user_group_id=course_group.id
+                    )
+                except CohortMembership.DoesNotExist:
+                    memberships_to_add += 1
+                    if commit:
+                        membership = CohortMembership(
+                            course_user_group=course_group,
+                            user=user,
+                            course_id=course_group.course_id
+                        )
+                        try:
+                            membership.save()
+                        except IntegrityError:  # If the user is in multiple cohorts, we arbitrarily choose between them
+                            # In this case, allow the pre-existing entry to be "correct"
+                            course_group.users.remove(user)
+                            user.course_groups.remove(course_group)
+
+        print '{} CohortMemberships did not match the CourseUserGroup table and will be deleted'.format(
+            memberships_to_delete
+        )
+        print '{} CourseUserGroup users do not have a CohortMembership; one will be added if it is valid'.format(
+            memberships_to_add
+        )
+        if commit:
+            print 'Changes have been made and saved.'
+        else:
+            print 'Dry run, changes have not been saved. Run again with "commit" argument to save changes'

openedx/core/djangoapps/course_groups/management/commands/remove_users_from_multiple_cohorts.py

-"""
-Script for removing users with multiple cohorts of a course from cohorts
-to ensure user's uniqueness for a course cohorts
-"""
-from django.contrib.auth.models import User
-from django.core.management.base import BaseCommand
-from django.db.models import Count
-
-from openedx.core.djangoapps.course_groups.models import CourseUserGroup
-
-
-class Command(BaseCommand):
-    """
-    Remove users with multiple cohorts of a course from all cohorts
-    """
-    help = 'Remove all users from multiple cohorts (except one) of each course'
-
-    def handle(self, *args, **options):
-        """
-        Execute the command
-        """
-        # Get entries of cohorts which have same user added multiple times for a single course
-        multiple_objects_cohorts = CourseUserGroup.objects.filter(group_type=CourseUserGroup.COHORT).\
-            values_list('users', 'course_id').annotate(user_count=Count('users')).filter(user_count__gt=1).\
-            order_by('users')
-        multiple_objects_cohorts_count = multiple_objects_cohorts.count()
-        multiple_course_cohorts_users = set(multiple_objects_cohorts.values_list('users', flat=True))
-        users_failed_to_cleanup = []
-
-        for user in User.objects.filter(id__in=multiple_course_cohorts_users):
-            print u"Removing user with id '{0}' from cohort groups".format(user.id)
-            try:
-                # remove user from only cohorts
-                user.course_groups.remove(*user.course_groups.filter(group_type=CourseUserGroup.COHORT))
-            except AttributeError as err:
-                users_failed_to_cleanup.append(user.email)
-                print u"Failed to remove user with id {0} from cohort groups, error: {1}".format(user.id, err)
-
-        print "=" * 80
-        print u"=" * 30 + u"> Cohorts summary"
-        print(
-            u"Total number of CourseUserGroup of type '{0}' with multiple users: {1}".format(
-                CourseUserGroup.COHORT, multiple_objects_cohorts_count
-            )
-        )
-        print(
-            u"Total number of unique users with multiple course cohorts: {0}".format(
-                len(multiple_course_cohorts_users)
-            )
-        )
-        print(
-            u"Users which failed on cohorts cleanup [{0}]: [{1}]".format(
-                len(users_failed_to_cleanup), (', '.join(users_failed_to_cleanup))
-            )
-        )
-        print "=" * 80

openedx/core/djangoapps/course_groups/management/commands/tests/test_remove_users_from_multiple_cohorts.py → openedx/core/djangoapps/course_groups/management/commands/tests/test_post_cohort_membership_fix.py

 """
-Tests for cleanup of users which are added in multiple cohorts of a course
+Test for the post-migration fix commands that are included with this djangoapp
 """
-from django.core.exceptions import MultipleObjectsReturned
 from django.core.management import call_command
 from django.test.client import RequestFactory
 
 from openedx.core.djangoapps.course_groups.views import cohort_handler
-from openedx.core.djangoapps.course_groups.cohorts import get_cohort, get_cohort_by_name
+from openedx.core.djangoapps.course_groups.cohorts import get_cohort_by_name
 from openedx.core.djangoapps.course_groups.tests.helpers import config_course_cohorts
+from openedx.core.djangoapps.course_groups.models import CohortMembership
 from student.tests.factories import UserFactory
 from xmodule.modulestore.tests.django_utils import ModuleStoreTestCase
 from xmodule.modulestore.tests.factories import CourseFactory
 
 
-class TestMultipleCohortUsers(ModuleStoreTestCase):
+class TestPostMigrationFix(ModuleStoreTestCase):
     """
-    Base class for testing users with multiple cohorts
+    Base class for testing post-migration fix commands
     """
     def setUp(self):
         """
         setup course, user and request for tests
         """
-        super(TestMultipleCohortUsers, self).setUp()
+        super(TestPostMigrationFix, self).setUp()
         self.course1 = CourseFactory.create()
         self.course2 = CourseFactory.create()
         self.user1 = UserFactory(is_staff=True)
@@ -29,54 +29,56 @@ class TestMultipleCohortUsers(ModuleStoreTestCase):
         self.request = RequestFactory().get("dummy_url")
         self.request.user = self.user1
 
-    def test_users_with_multiple_cohorts_cleanup(self):
+    def test_post_cohortmembership_fix(self):
         """
-        Test that user which have been added in multiple cohorts of a course,
-        can get cohorts without error after running cohorts cleanup command
+        Test that changes made *after* migration, but *before* turning on new code are handled properly
         """
-        # set two auto_cohort_groups for both courses
-        config_course_cohorts(
-            self.course1, is_cohorted=True, auto_cohorts=["Course1AutoGroup1", "Course1AutoGroup2"]
-        )
-        config_course_cohorts(
-            self.course2, is_cohorted=True, auto_cohorts=["Course2AutoGroup1", "Course2AutoGroup2"]
-        )
+        # First, we're going to simulate some problem states that can arise during this window
+        config_course_cohorts(self.course1, is_cohorted=True, auto_cohorts=["Course1AutoGroup1", "Course1AutoGroup2"])
 
-        # get the cohorts from the courses, which will cause auto cohorts to be created
+        # Get the cohorts from the courses, which will cause auto cohorts to be created
         cohort_handler(self.request, unicode(self.course1.id))
-        cohort_handler(self.request, unicode(self.course2.id))
         course_1_auto_cohort_1 = get_cohort_by_name(self.course1.id, "Course1AutoGroup1")
         course_1_auto_cohort_2 = get_cohort_by_name(self.course1.id, "Course1AutoGroup2")
-        course_2_auto_cohort_1 = get_cohort_by_name(self.course2.id, "Course2AutoGroup1")
 
-        # forcefully add user1 in two auto cohorts
-        course_1_auto_cohort_1.users.add(self.user1)
-        course_1_auto_cohort_2.users.add(self.user1)
-        # forcefully add user2 in auto cohorts of both courses
-        course_1_auto_cohort_1.users.add(self.user2)
-        course_2_auto_cohort_1.users.add(self.user2)
+        # When migrations were first run, the users were assigned to CohortMemberships correctly
+        membership1 = CohortMembership(
+            course_id=course_1_auto_cohort_1.course_id,
+            user=self.user1,
+            course_user_group=course_1_auto_cohort_1
+        )
+        membership1.save()
+        membership2 = CohortMembership(
+            course_id=course_1_auto_cohort_1.course_id,
+            user=self.user2,
+            course_user_group=course_1_auto_cohort_1
+        )
+        membership2.save()
+
+        # But before CohortMembership code was turned on, some changes were made:
+        course_1_auto_cohort_2.users.add(self.user1)  # user1 is now in 2 cohorts in the same course!
+        course_1_auto_cohort_2.users.add(self.user2)
+        course_1_auto_cohort_1.users.remove(self.user2)  # and user2 was moved, but no one told CohortMembership!
 
-        # now check that when user1 goes on discussion page and tries to get
-        # cohorts 'MultipleObjectsReturned' exception is returned
-        with self.assertRaises(MultipleObjectsReturned):
-            get_cohort(self.user1, self.course1.id)
-        # also check that user 2 can go on discussion page of both courses
-        # without any exception
-        get_cohort(self.user2, self.course1.id)
-        get_cohort(self.user2, self.course2.id)
+        # run the post-CohortMembership command, dry-run
+        call_command('post_cohort_membership_fix')
 
-        # call command to remove users added in multiple cohorts of a course
-        # are removed from all cohort groups
-        call_command('remove_users_from_multiple_cohorts')
+        # Verify nothing was changed in dry-run mode.
+        self.assertEqual(self.user1.course_groups.count(), 2)  # CourseUserGroup has 2 entries for user1
 
-        # check that only user1 (with multiple cohorts) is removed from cohorts
-        # and user2 is still in auto cohorts of both course after running
-        # 'remove_users_from_multiple_cohorts' management command
-        self.assertEqual(self.user1.course_groups.count(), 0)
-        self.assertEqual(self.user2.course_groups.count(), 2)
+        self.assertEqual(CohortMembership.objects.get(user=self.user2).course_user_group.name, 'Course1AutoGroup1')
         user2_cohorts = list(self.user2.course_groups.values_list('name', flat=True))
-        self.assertEqual(user2_cohorts, ['Course1AutoGroup1', 'Course2AutoGroup1'])
+        self.assertEqual(user2_cohorts, ['Course1AutoGroup2'])  # CourseUserGroup and CohortMembership disagree
+
+        # run the post-CohortMembership command, and commit it
+        call_command('post_cohort_membership_fix', 'commit')
 
-        # now check that user1 can get cohorts in which he is added
-        response = cohort_handler(self.request, unicode(self.course1.id))
-        self.assertEqual(response.status_code, 200)
+        # verify that both databases agree about the (corrected) state of the memberships
+        self.assertEqual(self.user1.course_groups.count(), 1)
+        self.assertEqual(CohortMembership.objects.filter(user=self.user1).count(), 1)
+
+        self.assertEqual(self.user2.course_groups.count(), 1)
+        self.assertEqual(CohortMembership.objects.filter(user=self.user2).count(), 1)
+        self.assertEqual(CohortMembership.objects.get(user=self.user2).course_user_group.name, 'Course1AutoGroup2')
+        user2_cohorts = list(self.user2.course_groups.values_list('name', flat=True))
+        self.assertEqual(user2_cohorts, ['Course1AutoGroup2'])

openedx/core/djangoapps/course_groups/migrations/rerun_0006.sh

-#!/bin/bash
-if [ $# -eq 0 ]; then
-    echo "$0: usage: rerun_0006.sh <arguments>. At minimum, '--settings=<environment>' is expected."
-    exit 1
-fi
-
-./manage.py lms migrate course_groups 0005 --fake "$@"
-./manage.py lms migrate course_groups 0006 "$@"