Add a user_roles_editable_org_unit_type_ids field to the Profile API

iaso/api/profiles/profiles.py

@@ -220,7 +220,7 @@ class ProfilesViewSet(viewsets.ViewSet):
 
     def get_queryset(self):
         account = self.request.user.iaso_profile.account
-        return Profile.objects.filter(account=account).prefetch_related("editable_org_unit_types")
+        return Profile.objects.filter(account=account).with_editable_org_unit_types()
 
     def list(self, request):
         limit = request.GET.get("limit", None)
@@ -260,7 +260,7 @@ def list(self, request):
             teams=teams,
             managed_users_only=managed_users_only,
             ids=ids,
-        )
+        ).order_by("id")
 
         queryset = queryset.prefetch_related(
             "user",

iaso/models/base.py

@@ -1405,6 +1405,21 @@ def as_dict(self):
         }
 
 
+class ProfileQuerySet(models.QuerySet):
+    def with_editable_org_unit_types(self):
+        qs = self
+        return qs.annotate(
+            annotated_editable_org_unit_types_ids=ArrayAgg(
+                "editable_org_unit_types__id", distinct=True, filter=Q(editable_org_unit_types__isnull=False)
+            ),
+            annotated_user_roles_editable_org_unit_type_ids=ArrayAgg(
+                "user_roles__editable_org_unit_types__id",
+                distinct=True,
+                filter=Q(user_roles__editable_org_unit_types__isnull=False),
+            ),
+        )
+
+
 class Profile(models.Model):
     account = models.ForeignKey(Account, on_delete=models.CASCADE)
     user = models.OneToOneField(User, on_delete=models.CASCADE, related_name="iaso_profile")
@@ -1424,6 +1439,8 @@ class Profile(models.Model):
         "OrgUnitType", related_name="editable_by_iaso_profile_set", blank=True
     )
 
+    objects = models.Manager.from_queryset(ProfileQuerySet)()
+
     class Meta:
         constraints = [models.UniqueConstraint(fields=["dhis2_id", "account"], name="dhis2_id_constraint")]
 
@@ -1440,6 +1457,16 @@ def as_dict(self, small=False):
         )
         all_permissions = user_group_permissions + user_permissions
         permissions = list(set(all_permissions))
+        try:
+            editable_org_unit_type_ids = self.annotated_editable_org_unit_types_ids
+        except AttributeError:
+            editable_org_unit_type_ids = [out.pk for out in self.editable_org_unit_types.all()]
+        try:
+            user_roles_editable_org_unit_type_ids = self.annotated_user_roles_editable_org_unit_type_ids
+        except AttributeError:
+            user_roles_editable_org_unit_type_ids = (
+                list(self.user_roles.values_list("editable_org_unit_types", flat=True)),
+            )
         if not small:
             return {
                 "id": self.id,
@@ -1462,7 +1489,8 @@ def as_dict(self, small=False):
                 "phone_number": self.phone_number.as_e164 if self.phone_number else None,
                 "country_code": region_code_for_number(self.phone_number).lower() if self.phone_number else None,
                 "projects": [p.as_dict() for p in self.projects.all().order_by("name")],
-                "editable_org_unit_type_ids": [out.pk for out in self.editable_org_unit_types.all()],
+                "editable_org_unit_type_ids": editable_org_unit_type_ids,
+                "user_roles_editable_org_unit_type_ids": user_roles_editable_org_unit_type_ids,
             }
         else:
             return {
@@ -1485,10 +1513,21 @@ def as_dict(self, small=False):
                 "phone_number": self.phone_number.as_e164 if self.phone_number else None,
                 "country_code": region_code_for_number(self.phone_number).lower() if self.phone_number else None,
                 "projects": [p.as_dict() for p in self.projects.all()],
-                "editable_org_unit_type_ids": [out.pk for out in self.editable_org_unit_types.all()],
+                "editable_org_unit_type_ids": editable_org_unit_type_ids,
+                "user_roles_editable_org_unit_type_ids": user_roles_editable_org_unit_type_ids,
             }
 
     def as_short_dict(self):
+        try:
+            editable_org_unit_type_ids = self.annotated_editable_org_unit_types_ids
+        except AttributeError:
+            editable_org_unit_type_ids = [out.pk for out in self.editable_org_unit_types.all()]
+        try:
+            user_roles_editable_org_unit_type_ids = self.annotated_user_roles_editable_org_unit_type_ids
+        except AttributeError:
+            user_roles_editable_org_unit_type_ids = (
+                list(self.user_roles.values_list("editable_org_unit_types", flat=True)),
+            )
         return {
             "id": self.id,
             "first_name": self.user.first_name,
@@ -1499,7 +1538,8 @@ def as_short_dict(self):
             "user_id": self.user.id,
             "phone_number": self.phone_number.as_e164 if self.phone_number else None,
             "country_code": region_code_for_number(self.phone_number).lower() if self.phone_number else None,
-            "editable_org_unit_type_ids": [out.pk for out in self.editable_org_unit_types.all()],
+            "editable_org_unit_type_ids": editable_org_unit_type_ids,
+            "user_roles_editable_org_unit_type_ids": user_roles_editable_org_unit_type_ids,
         }
 
     def has_a_team(self):

iaso/tests/models/test_profile.py

@@ -70,3 +70,26 @@ def test_has_org_unit_write_permission(self):
                 )
             )
         self.profile1.user_roles.clear()
+
+    def test_with_editable_org_unit_types(self):
+        org_unit_type_country = m.OrgUnitType.objects.create(name="Country")
+        org_unit_type_region = m.OrgUnitType.objects.create(name="Region")
+        org_unit_type_district = m.OrgUnitType.objects.create(name="District")
+
+        group_1 = Group.objects.create(name="Group 1")
+        user_role_1 = m.UserRole.objects.create(group=group_1, account=self.account)
+        user_role_1.editable_org_unit_types.set([org_unit_type_country])
+
+        group_2 = Group.objects.create(name="Group 2")
+        user_role_2 = m.UserRole.objects.create(group=group_2, account=self.account)
+        user_role_2.editable_org_unit_types.set([org_unit_type_region])
+
+        self.profile1.user_roles.set([user_role_1, user_role_2])
+        self.profile1.editable_org_unit_types.set([org_unit_type_district])
+
+        profile = m.Profile.objects.filter(id=self.profile1.pk).with_editable_org_unit_types().first()
+
+        self.assertEqual(profile.annotated_editable_org_unit_types_ids, [org_unit_type_district.pk])
+        self.assertCountEqual(
+            profile.annotated_user_roles_editable_org_unit_type_ids, [org_unit_type_country.pk, org_unit_type_region.pk]
+        )