use bulk_create and transaction for user generator

This commit is contained in:
virusdefender 2017-11-24 23:30:17 +08:00
parent 9889ac5b4a
commit 2d038c7bcc
4 changed files with 26 additions and 31 deletions

View File

@ -40,8 +40,7 @@ class GenerateUserSerializer(serializers.Serializer):
suffix = serializers.CharField(max_length=16, allow_blank=True) suffix = serializers.CharField(max_length=16, allow_blank=True)
number_from = serializers.IntegerField() number_from = serializers.IntegerField()
number_to = serializers.IntegerField() number_to = serializers.IntegerField()
default_email = serializers.CharField(max_length=64) password_length = serializers.IntegerField(max_value=16, default=8)
password_length = serializers.IntegerField(required=False, max_value=16)
class ImportUserSeralizer(serializers.Serializer): class ImportUserSeralizer(serializers.Serializer):

View File

@ -1,4 +1,5 @@
import time import time
from unittest import mock from unittest import mock
from datetime import timedelta from datetime import timedelta
from copy import deepcopy from copy import deepcopy
@ -610,6 +611,3 @@ class GenerateUserAPITest(APITestCase):
resp = self.client.post(self.url, data=self.data) resp = self.client.post(self.url, data=self.data)
self.assertSuccess(resp) self.assertSuccess(resp)
mock_workbook.assert_called() mock_workbook.assert_called()
data = resp.data["data"]
self.assertEqual(data["created_count"], 6)
self.assertEqual(data["get_count"], 0)

View File

@ -173,9 +173,6 @@ class GenerateUserAPI(APIView):
if data["number_from"] > data["number_to"]: if data["number_from"] > data["number_to"]:
return self.error("Start number must be lower than end number") return self.error("Start number must be lower than end number")
password_length = data.get("password_length", 8)
default_email = data.get("default_email")
file_id = rand_str(8) file_id = rand_str(8)
filename = f"/tmp/{file_id}.xlsx" filename = f"/tmp/{file_id}.xlsx"
workbook = xlsxwriter.Workbook(filename) workbook = xlsxwriter.Workbook(filename)
@ -184,26 +181,27 @@ class GenerateUserAPI(APIView):
worksheet.write("A1", "Username") worksheet.write("A1", "Username")
worksheet.write("B1", "Password") worksheet.write("B1", "Password")
i = 1 i = 1
created_count = 0
get_count = 0 user_list = []
for number in range(data["number_from"], data["number_to"] + 1): for number in range(data["number_from"], data["number_to"] + 1):
username = f"{data['prefix']}{number}{data['suffix']}" raw_password = rand_str(data["password_length"])
password = rand_str(password_length) user = User(username=f"{data['prefix']}{number}{data['suffix']}", password=make_password(raw_password))
user, created = User.objects.get_or_create(username=username) user.raw_password = raw_password
user.email = default_email user_list.append(user)
user.set_password(password)
user.save() try:
if created: with transaction.atomic():
UserProfile.objects.create(user=user)
created_count += 1 ret = User.objects.bulk_create(user_list)
else: UserProfile.objects.bulk_create([UserProfile(user=user) for user in ret])
get_count += 1 for item in user_list:
worksheet.write_string(i, 0, username) worksheet.write_string(i, 0, item.username)
worksheet.write_string(i, 1, password) worksheet.write_string(i, 1, item.raw_password)
i += 1 i += 1
workbook.close() workbook.close()
return self.success({ return self.success({"file_id": file_id})
"file_id": file_id, except IntegrityError as e:
"created_count": created_count, # Extract detail from exception message
"get_count": get_count # duplicate key value violates unique constraint "user_username_key"
}) # DETAIL: Key (username)=(root11) already exists.
return self.error(str(e).split("\n")[1])

View File

@ -7,7 +7,7 @@ DATABASES = {
'default': { 'default': {
'ENGINE': 'django.db.backends.postgresql_psycopg2', 'ENGINE': 'django.db.backends.postgresql_psycopg2',
'HOST': '127.0.0.1', 'HOST': '127.0.0.1',
'PORT': 5433, 'PORT': 5432,
'NAME': "onlinejudge", 'NAME': "onlinejudge",
'USER': "onlinejudge", 'USER': "onlinejudge",
'PASSWORD': 'onlinejudge' 'PASSWORD': 'onlinejudge'