feat: added batch_get_or_create class method for ExternalId (#25844)

* Added batch_get_or_create_user_ids method for ExternalId Model

Update doc string

* Update docstring & fix bug on test

[BD-24] [BB-2726] [TNL-7330]
This commit is contained in:
Shimul Chowdhury
2021-03-25 22:09:32 +06:00
committed by GitHub
parent dabad9fdfd
commit eba710ccb5
2 changed files with 122 additions and 0 deletions

View File

@@ -103,3 +103,58 @@ class ExternalId(TimeStampedModel):
)
)
return external_id, created
@classmethod
def batch_get_or_create_user_ids(cls, users, type_name):
"""
Creates ExternalIds in batch.
Given a list of users and a type_name, this method creates new external ids for
users that do not already have one. Then returns a dictionary mapping user id with
corresponding external id.
Arguments:
users: List of User to create the IDs for
type_name (str): Name of the type of ExternalId
Returns:
dict: None if fails, otherwise ExternalIds mapped by User.id
{
user_id: ExternalId
}
"""
try:
type_obj = ExternalIdType.objects.get(name=type_name)
except ExternalIdType.DoesNotExist:
LOGGER.info(
'Batch ID Creation failed, no external id type of {type!r}'.format(
type=type_name
)
)
return None
# get user ids in a set
user_ids = {user.id for user in users}
# find users for those external ids needs to be created
externalid_count = models.Count('externalid', filter=models.Q(externalid__external_id_type=type_obj))
users_wo_externalid = User.objects.annotate(externalid_count=externalid_count).filter(
externalid_count=0,
id__in=user_ids,
)
# get external ids that already exists
existing_externalids = cls.objects.filter(user__in=users, external_id_type=type_obj)
# prepare result dict with existing externalids.
result = {eid.user_id: eid for eid in existing_externalids}
# if there are users with no external id, create external ids for them
if len(users_wo_externalid) > 0:
new_externalids = cls.objects.bulk_create([
cls(user=user, external_id_type=type_obj) for user in users_wo_externalid
])
# append newly created externalids to result dict
result.update({eid.user_id: eid for eid in new_externalids})
return result

View File

@@ -0,0 +1,67 @@
"""
Test batch_get_or_create in ExternalId model
"""
from django.test import TransactionTestCase
from common.djangoapps.student.tests.factories import UserFactory
from openedx.core.djangoapps.external_user_ids.models import ExternalId
from openedx.core.djangoapps.external_user_ids.tests.factories import ExternalIDTypeFactory
class TestBatchGenerateExternalIds(TransactionTestCase):
"""
Test ExternalId.batch_get_or_create_user_ids
"""
# Following are the queries
# 1 - Get ExternalIdType
# 2 - Find users for those external ids needs to be created
# 3 - Get external ids that already exists
# 4 - BEGIN (from bulk_create)
# 5 - Create new external ids
EXPECTED_NUM_OF_QUERIES = 5
def test_batch_get_or_create_user_ids(self):
"""
Test if batch_get_or_create creates ExternalIds in batch
"""
id_type = ExternalIDTypeFactory.create(name='test')
users = [UserFactory() for _ in range(10)]
with self.assertNumQueries(self.EXPECTED_NUM_OF_QUERIES):
result = ExternalId.batch_get_or_create_user_ids(users, id_type)
assert len(result) == len(users)
for user in users:
assert result[user.id].external_id_type.name == 'test'
assert result[user.id].user == user
def test_batch_get_or_create_user_ids_existing_ids(self):
"""
Test batch creation output when there are existing ids for some user
"""
id_type = ExternalIDTypeFactory.create(name='test')
# first let's create some user and externalids for them
users = [UserFactory() for _ in range(10)]
with self.assertNumQueries(self.EXPECTED_NUM_OF_QUERIES):
result = ExternalId.batch_get_or_create_user_ids(users, id_type)
# now create some new user and try to create externalids for all user
new_users = [UserFactory() for _ in range(5)]
all_users = users + new_users
with self.assertNumQueries(self.EXPECTED_NUM_OF_QUERIES):
result = ExternalId.batch_get_or_create_user_ids(all_users, id_type)
assert len(result) == len(all_users)
def test_batch_get_or_create_user_ids_wrong_type(self):
"""
Test if batch_get_or_create returns None if wrong type given
"""
users = [UserFactory() for _ in range(2)]
external_ids = ExternalId.batch_get_or_create_user_ids(users, 'invalid')
assert external_ids is None