From 241fe94504bd81e08fdd83f9b3e54dcf416bc86b Mon Sep 17 00:00:00 2001 From: Varsha Menon Date: Tue, 6 Aug 2024 10:12:36 -0400 Subject: [PATCH] fix: bulk create course staff mgmt command --- .../commands/bulk_add_course_staff.py | 24 ++++++++++--------- .../test/test_bulk_add_course_staff.py | 10 ++++---- 2 files changed, 18 insertions(+), 16 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 8edc1778..b97226e5 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,7 +4,6 @@ import time from django.core.management.base import BaseCommand -from django.db import transaction from edx_exams.apps.core.models import CourseStaffRole, User @@ -44,7 +43,6 @@ def add_arguments(self, parser): dest='batch_delay', help='Time delay in seconds for each batch') - @transaction.atomic def handle(self, *args, **options): """ The main logic and entry point of the management command @@ -73,6 +71,10 @@ def add_course_staff_from_csv(self, csv_file, batch_size, batch_delay): ) 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')), @@ -85,12 +87,12 @@ def add_course_staff_from_csv(self, csv_file, batch_size, 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) + # 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) diff --git a/edx_exams/apps/core/management/commands/test/test_bulk_add_course_staff.py b/edx_exams/apps/core/management/commands/test/test_bulk_add_course_staff.py index 023974df..d8151674 100644 --- a/edx_exams/apps/core/management/commands/test/test_bulk_add_course_staff.py +++ b/edx_exams/apps/core/management/commands/test/test_bulk_add_course_staff.py @@ -90,7 +90,7 @@ def test_add_course_staff_with_not_default_batch_size(self): 'sam,sam@pond.com,staff,course-v1:edx+test+f20\n'] with NamedTemporaryFile() as csv: csv = self._write_test_csv(csv, lines) - with self.assertNumQueries(8): + with self.assertNumQueries(6): call_command(self.command, f'--csv_path={csv.name}', '--batch_size=1') def test_add_course_staff_with_batch_size_larger_than_list(self): @@ -99,7 +99,7 @@ def test_add_course_staff_with_batch_size_larger_than_list(self): 'sam,sam@pond.com,staff,course-v1:edx+test+f20\n'] with NamedTemporaryFile() as csv: csv = self._write_test_csv(csv, lines) - with self.assertNumQueries(6): + with self.assertNumQueries(4): call_command(self.command, f'--csv_path={csv.name}', '--batch_size=3') def test_add_course_staff_with_batch_size_smaller_than_list(self): @@ -109,7 +109,7 @@ def test_add_course_staff_with_batch_size_smaller_than_list(self): 'tam,tam@pond.com,staff,course-v1:edx+test+f20\n'] with NamedTemporaryFile() as csv: csv = self._write_test_csv(csv, lines) - with self.assertNumQueries(9): + with self.assertNumQueries(7): call_command(self.command, f'--csv_path={csv.name}', '--batch_size=2') def test_add_course_staff_with_not_default_batch_delay(self): @@ -125,7 +125,7 @@ def test_add_course_staff_with_not_default_batch_delay(self): def test_num_queries_correct(self): """ - Assert the number of queries to be 4 + 1 * number of lines: + Assert the number of queries to be 2 + 1 * number of lines: 2 for savepoint/release savepoint 1 to bulk create users, 1 to bulk create course role 1 for each user (to get user) @@ -134,7 +134,7 @@ def test_num_queries_correct(self): 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(4 + num_lines): + with self.assertNumQueries(2 + num_lines): call_command(self.command, f'--csv_path={csv.name}') def test_dupe_user_csv(self):