temp: testing bulk create

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

@@ -4,6 +4,7 @@
 import time
 
 from django.core.management.base import BaseCommand
+from django.db import transaction
 
 from edx_exams.apps.core.models import CourseStaffRole, User
 
@@ -56,43 +57,37 @@ def handle(self, *args, **options):
 
         logger.info('Bulk add course staff complete!')
 
+    @transaction.atomic
     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(csv.DictReader(csv_file))
+        users = {}
 
         # bulk create users
         for i in range(0, len(reader), batch_size):
-            User.objects.bulk_create(
+            bulk_users = User.objects.bulk_create(
                 (User(
                     username=row.get('username'),
                     email=row.get('email'),
                 ) for row in reader[i:i + batch_size]),
                 ignore_conflicts=True,
             )
+            users_to_add = {k: v for k, v in [(u.username, u) for u in bulk_users]}
+            logger.info('*** existing users ***', users_to_add)
+            users.update(users_to_add)
+
             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')),
+                    user=users.get(row.get('username')),
                     course_id=row.get('course_id'),
                     role=row.get('role'),
                 ) for row in reader[i:i + batch_size]),
                 ignore_conflicts=True,
             )
             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')),
-            #         course_id=row.get('course_id'),
-            #         role=row.get('role'),
-            #     ) for row in reader[i:i + batch_size]),
-            #     ignore_conflicts=True,
-            # )
-            # time.sleep(batch_delay)