Optimize group operations

azure_auth/handlers.py

@@ -21,6 +21,8 @@ class AuthHandler:
     Class to interface with `msal` package and execute authentication process.
     """
 
+    GROUPS_UPDATED = False
+
     def __init__(self, request: HttpRequest):
         """
 
@@ -119,22 +121,31 @@ def authenticate(self, token: dict) -> AbstractBaseUser:
                 **self._map_attributes_to_user(**attributes),
             )
 
+        if not AuthHandler.GROUPS_UPDATED:
+            role_mappings = settings.AZURE_AUTH.get("ROLES")
+            for group_name in role_mappings.values():
+                Group.objects.get_or_create(name=group_name)
+            AuthHandler.GROUPS_UPDATED = True
+
         # Syncing azure token claim roles with django user groups
         # A role mapping in the AZURE_AUTH settings is expected.
         role_mappings = settings.AZURE_AUTH.get("ROLES")
         azure_token_roles = token.get("id_token_claims", {}).get("roles", None)
         if role_mappings:  # pragma: no branch
-            for role, group_name in role_mappings.items():
-                # all groups are created by default if they not exist
-                django_group = Group.objects.get_or_create(name=group_name)[0]
-
-                if azure_token_roles and role in azure_token_roles:
-                    # Add user with permissions to the corresponding django group
-                    user.groups.add(django_group)
-                else:
-                    # No permission so check if user is in group and remove
-                    if user.groups.filter(name=group_name).exists():
-                        user.groups.remove(django_group)
+            token_groups = [
+                name for (id, name) in role_mappings.items() if id in azure_token_roles
+            ]
+            current_groups = list(user.groups.values_list("name", flat=True))
+
+            if to_add := [item for item in token_groups if item not in current_groups]:
+                user.groups.add(*Group.objects.filter(name__in=to_add))
+
+            if to_remove := [
+                item
+                for item in current_groups
+                if item in role_mappings.values() and item not in token_groups
+            ]:
+                user.groups.remove(*Group.objects.filter(name__in=to_remove))
 
         return user
 

azure_auth/tests/conftest.py

@@ -3,6 +3,7 @@
 import pytest
 from django.contrib.auth import get_user_model
 from mixer.backend.django import mixer
+from azure_auth.handlers import AuthHandler
 
 UserModel = get_user_model()
 
@@ -19,6 +20,11 @@ def user(request):
     return _user
 
 
+@pytest.fixture(scope="function", autouse=True)
+def clean_groups():
+    AuthHandler.GROUPS_UPDATED = False
+
+
 @pytest.fixture(scope="function")
 def auth_flow(request):
     _auth_flow = {