[BD-24] [BB-2726] [TNL-7330] Added batch_get_or_create class method for ExternalId

openedx/core/djangoapps/external_user_ids/models.py

@@ -103,3 +103,58 @@ def add_new_user_id(cls, user, type_name):
                 )
             )
         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: Returns 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}'.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

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(result)
+
+        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