reset grades management command choose database

lms/djangoapps/grades/management/commands/reset_grades.py

@@ -65,6 +65,12 @@ class Command(BaseCommand):
             dest='modified_end',
             help='Ending range for modified date (inclusive): e.g. "2016-12-23 16:43"',
         )
+        parser.add_argument(
+            '--db_table',
+            dest='db_table',
+            help='Specify "subsection" to reset subsection grades or "course" to reset course grades. If absent, both '
+                 'are reset.',
+        )
 
     def handle(self, *args, **options):
         course_keys = None
@@ -73,6 +79,9 @@ class Command(BaseCommand):
 
         run_mode = get_mutually_exclusive_required_option(options, 'delete', 'dry_run')
         courses_mode = get_mutually_exclusive_required_option(options, 'courses', 'all_courses')
+        db_table = options.get('db_table')
+        if db_table not in {'subsection', 'course', None}:
+            raise CommandError('Invalid value for db_table. Valid options are "subsection" or "course" only.')
 
         if options.get('modified_start'):
             modified_start = datetime.strptime(options['modified_start'], DATE_FORMAT)
@@ -89,8 +98,11 @@ class Command(BaseCommand):
 
         operation = self._query_grades if run_mode == 'dry_run' else self._delete_grades
 
-        operation(PersistentSubsectionGrade, course_keys, modified_start, modified_end)
-        operation(PersistentCourseGrade, course_keys, modified_start, modified_end)
+        if db_table == 'subsection' or db_table is None:
+            operation(PersistentSubsectionGrade, course_keys, modified_start, modified_end)
+
+        if db_table == 'course' or db_table is None:
+            operation(PersistentCourseGrade, course_keys, modified_start, modified_end)
 
         log.info("reset_grade: Finished in %s mode!", run_mode)
 

lms/djangoapps/grades/management/commands/tests/test_reset_grades.py

@@ -94,26 +94,30 @@ class TestResetGrades(TestCase):
                     subsection_grade_params['usage_key'] = subsection_key
                     PersistentSubsectionGrade.update_or_create_grade(**subsection_grade_params)
 
-    def _assert_grades_exist_for_courses(self, course_keys):
+    def _assert_grades_exist_for_courses(self, course_keys, db_table=None):
         """
         Assert grades for given courses exist.
         """
         for course_key in course_keys:
-            self.assertIsNotNone(PersistentCourseGrade.read_course_grade(self.user_ids[0], course_key))
-            for subsection_key in self.subsection_keys_by_course[course_key]:
-                self.assertIsNotNone(PersistentSubsectionGrade.read_grade(self.user_ids[0], subsection_key))
+            if db_table == "course" or db_table is None:
+                self.assertIsNotNone(PersistentCourseGrade.read_course_grade(self.user_ids[0], course_key))
+            if db_table == "subsection" or db_table is None:
+                for subsection_key in self.subsection_keys_by_course[course_key]:
+                    self.assertIsNotNone(PersistentSubsectionGrade.read_grade(self.user_ids[0], subsection_key))
 
-    def _assert_grades_absent_for_courses(self, course_keys):
+    def _assert_grades_absent_for_courses(self, course_keys, db_table=None):
         """
         Assert grades for given courses do not exist.
         """
         for course_key in course_keys:
-            with self.assertRaises(PersistentCourseGrade.DoesNotExist):
-                PersistentCourseGrade.read_course_grade(self.user_ids[0], course_key)
+            if db_table == "course" or db_table is None:
+                with self.assertRaises(PersistentCourseGrade.DoesNotExist):
+                    PersistentCourseGrade.read_course_grade(self.user_ids[0], course_key)
 
-            for subsection_key in self.subsection_keys_by_course[course_key]:
-                with self.assertRaises(PersistentSubsectionGrade.DoesNotExist):
-                    PersistentSubsectionGrade.read_grade(self.user_ids[0], subsection_key)
+            if db_table == "subsection" or db_table is None:
+                for subsection_key in self.subsection_keys_by_course[course_key]:
+                    with self.assertRaises(PersistentSubsectionGrade.DoesNotExist):
+                        PersistentSubsectionGrade.read_grade(self.user_ids[0], subsection_key)
 
     def _assert_stat_logged(self, mock_log, num_rows, grade_model_class, message_substring, log_offset):
         self.assertIn('reset_grade: ' + message_substring, mock_log.info.call_args_list[log_offset][0][0])
@@ -222,6 +226,17 @@ class TestResetGrades(TestCase):
         self._assert_grades_absent_for_courses(self.course_keys[:2])
         self._assert_grades_exist_for_courses(self.course_keys[2:])
 
+    @ddt.data('subsection', 'course')
+    def test_specify_db_table(self, db_table):
+        self._update_or_create_grades()
+        self._assert_grades_exist_for_courses(self.course_keys)
+        self.command.handle(delete=True, all_courses=True, db_table=db_table)
+        self._assert_grades_absent_for_courses(self.course_keys, db_table=db_table)
+        if db_table == "subsection":
+            self._assert_grades_exist_for_courses(self.course_keys, db_table='course')
+        else:
+            self._assert_grades_exist_for_courses(self.course_keys, db_table='subsection')
+
     @patch('lms.djangoapps.grades.management.commands.reset_grades.log')
     def test_dry_run_all_courses(self, mock_log):
         self._update_or_create_grades()
@@ -279,6 +294,13 @@ class TestResetGrades(TestCase):
         with self.assertRaisesRegexp(CommandError, 'Invalid key specified.*invalid/key'):
             self.command.handle(dry_run=True, courses=['invalid/key'])
 
+    def test_invalid_db_table(self):
+        with self.assertRaisesMessage(
+                CommandError,
+                'Invalid value for db_table. Valid options are "subsection" or "course" only.'
+        ):
+            self.command.handle(delete=True, all_courses=True, db_table="not course or subsection")
+
     def test_no_run_mode(self):
         with self.assertRaisesMessage(CommandError, 'Either --delete or --dry_run must be specified.'):
             self.command.handle(all_courses=True)