feat: add course staff mgmt command

edx_exams/apps/core/management/commands/bulk_add_course_staff.py

@@ -0,0 +1,98 @@
+"""Management command to populate the CourseStaffRole model and related User objects, from LMS, using CSV"""
+import logging
+import time
+
+import unicodecsv
+from django.core.management.base import BaseCommand
+from django.db import transaction
+
+from edx_exams.apps.core.models import CourseStaffRole, User
+
+logger = logging.getLogger(__name__)
+
+
+class Command(BaseCommand):
+    """
+    Management command to add Course Staff (and User) in batches from CSV
+    """
+    help = """
+    Add Course Staff in bulk from CSV.
+    Expects that the data will be provided in a csv file format with the first row
+    being the header and columns being: username, email, role, course_id.
+    Example:
+        $ ... bulk_add_course_staff --csv_path=foo.csv
+    """
+
+    def add_arguments(self, parser):
+        parser.add_argument(
+            '-p', '--csv_path',
+            metavar='csv_path',
+            dest='csv_path',
+            required=False,
+            help='Path to CSV file.')
+        parser.add_argument(
+            '--batch_size',
+            type=int,
+            default=200,
+            dest='batch_size',
+            help='Batch size')
+        parser.add_argument(
+            '--batch_delay',
+            type=float,
+            default=1.0,
+            dest='batch_delay',
+            help='Time delay in each iteration')
+
+    @transaction.atomic
+    def handle(self, *args, **options):
+        """
+        The main logic and entry point of the management command
+        """
+        csv_path = options['csv_path']
+        batch_size = options['batch_size']
+        batch_delay = options['batch_delay']
+
+        if csv_path:
+            with open(csv_path, 'rb') as csv_file:
+                self.add_course_staff_from_csv(csv_file, batch_size, batch_delay)
+
+        logger.info('Bulk add course staff complete!')
+
+    def add_course_staff_from_csv(self, csv_file, batch_size, batch_delay):
+        """
+        Add the given set of course staff provided in csv
+        """
+        reader = list(unicodecsv.DictReader(csv_file))
+        users_to_create = []
+        users_existing = {u.username for u in User.objects.filter(username__in=[r.get('username') for r in reader])}
+        for row in reader:
+            # if not User.objects.filter(username=row.get('username')).exists():
+            #     users_to_create.append(row)
+            if row.get('username') not in users_existing:
+                users_to_create.append(row)
+
+        # bulk create users
+        for i in range(0, len(users_to_create), batch_size):
+            User.objects.bulk_create(
+                User(
+                    username=user.get('username'),
+                    email=user.get('email'),
+                )
+                for user in users_to_create[i:i + batch_size]
+            )
+            time.sleep(batch_delay)
+
+        # bulk create course staff
+        for i in range(0, len(reader), batch_size):
+            CourseStaffRole.objects.bulk_create(
+                CourseStaffRole(
+                    user=User.objects.get(
+                        username=row.get('username'),
+                        email=row.get('email'),
+                    ),
+                    course_id=row.get('course_id'),
+                    role=row.get('role'),
+                )
+                for row in reader[i:i + batch_size]
+            )
+            time.sleep(batch_delay)

edx_exams/apps/core/management/commands/test/test_bulk_add_course_staff.py

@@ -0,0 +1,86 @@
+""" Tests for bulk_add_course_staff management command """
+from tempfile import NamedTemporaryFile
+
+from django.core.management import call_command
+from django.test import TestCase
+
+from edx_exams.apps.core.models import CourseStaffRole
+from edx_exams.apps.core.test_utils.factories import UserFactory
+
+
+class TestBulkAddCourseStaff(TestCase):
+    """ Test bulk_add_course_staff management command """
+
+    def setUp(self):
+        super().setUp()
+        self.command = 'bulk_add_course_staff'
+        self.success_log_message = 'Bulk add course staff complete!'
+
+        # create existing user
+        self.user = UserFactory.create(
+            username='amy',
+            is_active=True,
+            is_staff=True,
+        )
+        self.user.email = '[email protected]'
+        self.user.save()
+
+        self.course_id = 'course-v1:edx+test+f19'
+
+    def _write_test_csv(self, csv, lines):
+        """ Write a test csv file with the lines provided """
+        csv.write(b'username,email,role,course_id\n')
+        for line in lines:
+            csv.write(line.encode())
+        csv.seek(0)
+        return csv
+
+    def test_empty_csv(self):
+        lines = []
+        with NamedTemporaryFile() as csv:
+            csv = self._write_test_csv(csv, lines)
+            call_command(self.command, f'--csv_path={csv.name}')
+
+    def test_add_course_staff_with_existing_user(self):
+        lines = ['amy,[email protected],staff,course-v1:edx+test+f19\n']
+        with NamedTemporaryFile() as csv:
+            csv = self._write_test_csv(csv, lines)
+            call_command(self.command, f'--csv_path={csv.name}')
+            assert CourseStaffRole.objects.filter(user=self.user.id).exists()
+
+    def test_add_course_staff_with_new_user(self):
+        lines = ['pam,[email protected],staff,course-v1:edx+test+f20\n']
+        with NamedTemporaryFile() as csv:
+            csv = self._write_test_csv(csv, lines)
+            call_command(self.command, f'--csv_path={csv.name}')
+            assert CourseStaffRole.objects.filter(course_id='course-v1:edx+test+f20').count() == 1
+
+    def test_add_course_staff_with_not_default_batch_size(self):
+        lines = ['pam,[email protected],staff,course-v1:edx+test+f20\n',
+                 'sam,[email protected],staff,course-v1:edx+test+f20\n']
+        with NamedTemporaryFile() as csv:
+            csv = self._write_test_csv(csv, lines)
+            call_command(self.command, f'--csv_path={csv.name}', '--batch_size=1')
+            assert CourseStaffRole.objects.filter(course_id='course-v1:edx+test+f20').count() == 2
+
+    def test_add_course_staff_with_not_default_batch_delay(self):
+        lines = ['pam,[email protected],staff,course-v1:edx+test+f20\n',
+                 'sam,[email protected],staff,course-v1:edx+test+f20\n']
+        with NamedTemporaryFile() as csv:
+            csv = self._write_test_csv(csv, lines)
+            call_command(self.command, f'--csv_path={csv.name}', '--batch_size=1', '--batch_delay=2')
+            assert CourseStaffRole.objects.filter(course_id='course-v1:edx+test+f20').count() == 2
+
+    def test_num_queries_correct(self):
+        """
+        Expect the number of queries to be 5 + 1 * number of lines
+        - 2 for savepoint/release savepoint, 1 to get existing usernames,
+        - 1 to bulk create users, 1 to bulk create course role
+        - 1 for each user (to get user)
+        """
+        num_lines = 20
+        lines = [f'pam{i},pam{i}@pond.com,staff,course-v1:edx+test+f20\n' for i in range(num_lines)]
+        with NamedTemporaryFile() as csv:
+            csv = self._write_test_csv(csv, lines)
+            with self.assertNumQueries(5 + num_lines):
+                call_command(self.command, f'--csv_path={csv.name}')

requirements/base.in

@@ -25,3 +25,4 @@ mysqlclient
 openedx-events
 pytz
 pymemcache
+unicodecsv

requirements/base.txt

@@ -333,6 +333,8 @@ typing-extensions==4.9.0
     # via
     #   asgiref
     #   edx-opaque-keys
+unicodecsv==0.14.1
+    # via -r requirements/base.in
 uritemplate==4.1.1
     # via
     #   coreapi

requirements/common_constraints.txt

@@ -1,4 +1,6 @@
 
+
+
 # A central location for most common version constraints
 # (across edx repos) for pip-installation.
 #

requirements/dev.txt

@@ -704,6 +704,8 @@ typing-extensions==4.9.0
     #   faker
     #   pylint
     #   rich
+unicodecsv==0.14.1
+    # via -r requirements/validation.txt
 uritemplate==4.1.1
     # via
     #   -r requirements/validation.txt

requirements/doc.txt

@@ -683,6 +683,8 @@ typing-extensions==4.9.0
     #   faker
     #   pylint
     #   rich
+unicodecsv==0.14.1
+    # via -r requirements/test.txt
 uritemplate==4.1.1
     # via
     #   -r requirements/test.txt

requirements/production.txt

@@ -462,6 +462,8 @@ typing-extensions==4.9.0
     #   -r requirements/base.txt
     #   asgiref
     #   edx-opaque-keys
+unicodecsv==0.14.1
+    # via -r requirements/base.txt
 uritemplate==4.1.1
     # via
     #   -r requirements/base.txt

requirements/quality.txt

@@ -640,6 +640,8 @@ typing-extensions==4.9.0
     #   faker
     #   pylint
     #   rich
+unicodecsv==0.14.1
+    # via -r requirements/test.txt
 uritemplate==4.1.1
     # via
     #   -r requirements/test.txt

requirements/test.txt

@@ -554,6 +554,8 @@ typing-extensions==4.9.0
     #   edx-opaque-keys
     #   faker
     #   pylint
+unicodecsv==0.14.1
+    # via -r requirements/base.txt
 uritemplate==4.1.1
     # via
     #   -r requirements/base.txt

requirements/validation.txt

@@ -819,6 +819,10 @@ typing-extensions==4.9.0
     #   faker
     #   pylint
     #   rich
+unicodecsv==0.14.1
+    # via
+    #   -r requirements/quality.txt
+    #   -r requirements/test.txt
 uritemplate==4.1.1
     # via
     #   -r requirements/quality.txt