diff --git a/openedx/core/djangoapps/external_user_ids/models.py b/openedx/core/djangoapps/external_user_ids/models.py index 30af507e5b..9bb633fb5e 100644 --- a/openedx/core/djangoapps/external_user_ids/models.py +++ b/openedx/core/djangoapps/external_user_ids/models.py @@ -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 diff --git a/openedx/core/djangoapps/external_user_ids/tests/test_batch_generate_id.py b/openedx/core/djangoapps/external_user_ids/tests/test_batch_generate_id.py new file mode 100644 index 0000000000..1a21fd1840 --- /dev/null +++ b/openedx/core/djangoapps/external_user_ids/tests/test_batch_generate_id.py @@ -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