diff --git a/iaso/api/profiles/profiles.py b/iaso/api/profiles/profiles.py index a940f2d1ef..df7e49b72d 100644 --- a/iaso/api/profiles/profiles.py +++ b/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", diff --git a/iaso/models/base.py b/iaso/models/base.py index cc7c6c506a..e4a6d7b6d5 100644 --- a/iaso/models/base.py +++ b/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): diff --git a/iaso/tests/models/test_profile.py b/iaso/tests/models/test_profile.py index 207a8012c7..c0aee1dff7 100644 --- a/iaso/tests/models/test_profile.py +++ b/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] + )