From 04e4c2b7a1e8b9b162eb5f8dc13a773542a07655 Mon Sep 17 00:00:00 2001
From: Varsha Menon <varshamenon04@gmail.com>
Date: Mon, 26 Aug 2024 15:03:35 -0400
Subject: [PATCH] temp: testing bulk create

---
 .../commands/bulk_add_course_staff.py         | 23 ++++++++-----------
 1 file changed, 9 insertions(+), 14 deletions(-)

diff --git a/edx_exams/apps/core/management/commands/bulk_add_course_staff.py b/edx_exams/apps/core/management/commands/bulk_add_course_staff.py
index b97226e5..cf1517ef 100644
--- a/edx_exams/apps/core/management/commands/bulk_add_course_staff.py
+++ b/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)