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)
number_from = serializers.IntegerField()
number_to = serializers.IntegerField()
default_email = serializers.CharField(max_length=64)
password_length = serializers.IntegerField(required=False, max_value=16)
password_length = serializers.IntegerField(max_value=16, default=8)
class ImportUserSeralizer(serializers.Serializer):

View File

@ -1,4 +1,5 @@
import time
from unittest import mock
from datetime import timedelta
from copy import deepcopy
@ -610,6 +611,3 @@ class GenerateUserAPITest(APITestCase):
resp = self.client.post(self.url, data=self.data)
self.assertSuccess(resp)
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"]:
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)
filename = f"/tmp/{file_id}.xlsx"
workbook = xlsxwriter.Workbook(filename)
@ -184,26 +181,27 @@ class GenerateUserAPI(APIView):
worksheet.write("A1", "Username")
worksheet.write("B1", "Password")
i = 1
created_count = 0
get_count = 0
user_list = []
for number in range(data["number_from"], data["number_to"] + 1):
username = f"{data['prefix']}{number}{data['suffix']}"
password = rand_str(password_length)
user, created = User.objects.get_or_create(username=username)
user.email = default_email
user.set_password(password)
user.save()
if created:
UserProfile.objects.create(user=user)
created_count += 1
else:
get_count += 1
worksheet.write_string(i, 0, username)
worksheet.write_string(i, 1, password)
raw_password = rand_str(data["password_length"])
user = User(username=f"{data['prefix']}{number}{data['suffix']}", password=make_password(raw_password))
user.raw_password = raw_password
user_list.append(user)
try:
with transaction.atomic():
ret = User.objects.bulk_create(user_list)
UserProfile.objects.bulk_create([UserProfile(user=user) for user in ret])
for item in user_list:
worksheet.write_string(i, 0, item.username)
worksheet.write_string(i, 1, item.raw_password)
i += 1
workbook.close()
return self.success({
"file_id": file_id,
"created_count": created_count,
"get_count": get_count
})
return self.success({"file_id": file_id})
except IntegrityError as e:
# Extract detail from exception message
# 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': {
'ENGINE': 'django.db.backends.postgresql_psycopg2',
'HOST': '127.0.0.1',
'PORT': 5433,
'PORT': 5432,
'NAME': "onlinejudge",
'USER': "onlinejudge",
'PASSWORD': 'onlinejudge'