feat(hierarchies): Recursive serializers (#468)

* Recursive serializers

* fix root hierarchies children with no parents

Author: samonaisi

apps/accounts/models.py

@@ -6,7 +6,6 @@
 
 from django.contrib.auth.models import AbstractUser, Group, Permission
 from django.contrib.contenttypes.models import ContentType
-from django.contrib.postgres.aggregates import ArrayAgg
 from django.contrib.postgres.fields import ArrayField
 from django.core.validators import MaxValueValidator
 from django.db import models, transaction
@@ -177,43 +176,6 @@ def update_or_create_root(cls, organization: "Organization"):
         )
         return root_group
 
-    @classmethod
-    def _get_hierarchy(cls, groups: dict[int, dict], group_id: int):
-        from apps.files.serializers import ImageSerializer
-
-        return {
-            "id": groups[group_id].id,
-            "slug": groups[group_id].slug,
-            "name": groups[group_id].name,
-            "publication_status": groups[group_id].publication_status,
-            "children": [
-                cls._get_hierarchy(groups, child)
-                for child in groups[group_id].children_ids
-                if child is not None and child in groups
-            ],
-            "roles": [group.name for group in groups[group_id].groups.all()],
-            "header_image": (
-                ImageSerializer(groups[group_id].header_image).data
-                if groups[group_id].header_image
-                else None
-            ),
-        }
-
-    def get_hierarchy(self, user: Optional["ProjectUser"] = None) -> dict:
-        # This would be better with a recursive serializer, but it doubles the query time
-        if user:
-            groups = (
-                user.get_people_group_queryset()
-                | PeopleGroup.objects.filter(is_root=True).distinct()
-            )
-        else:
-            groups = PeopleGroup.objects.all()
-        groups = groups.filter(organization=self.organization.pk).annotate(
-            children_ids=ArrayAgg("children")
-        )
-        groups = {group.id: group for group in groups}
-        return self._get_hierarchy(groups, self.id)
-
     def get_default_managers_permissions(self) -> QuerySet[Permission]:
         return Permission.objects.filter(content_type=self.content_type)
 

apps/accounts/serializers.py

@@ -82,9 +82,11 @@ class Meta:
 
     def get_people_groups(self, user: ProjectUser) -> list:
         organization = self.context.get("organization")
-        queryset = PeopleGroup.objects.filter(
-            groups__users=user, is_root=False
-        ).distinct()
+        queryset = (
+            PeopleGroup.objects.filter(groups__users=user, is_root=False)
+            .select_related("organization")
+            .distinct()
+        )
         if organization:
             queryset = queryset.filter(organization=organization).distinct()
         return PeopleGroupSuperLightSerializer(
@@ -199,6 +201,7 @@ def get_people_groups(self, user: ProjectUser) -> list:
         queryset = (
             request_user.get_people_group_queryset()
             .filter(groups__users=user, is_root=False)
+            .select_related("organization")
             .distinct()
         )
         if organization:
@@ -226,9 +229,11 @@ def get_can_mentor_on(self, user: ProjectUser) -> List[Dict]:
 class PeopleGroupSuperLightSerializer(
     AutoTranslatedModelSerializer, serializers.ModelSerializer
 ):
+    organization = serializers.SlugRelatedField(read_only=True, slug_field="code")
+
     class Meta:
         model = PeopleGroup
-        read_only_fields = ["id", "slug", "name"]
+        read_only_fields = ["id", "slug", "name", "organization"]
         fields = read_only_fields
 
 
@@ -264,6 +269,57 @@ class Meta:
         ]
 
 
+class PeopleGroupHierarchySerializer(
+    AutoTranslatedModelSerializer,
+    serializers.ModelSerializer,
+):
+    children = serializers.SerializerMethodField()
+    header_image = ImageSerializer(read_only=True)
+    roles = serializers.SlugRelatedField(
+        many=True,
+        slug_field="name",
+        read_only=True,
+        source="groups",
+    )
+
+    class Meta:
+        model = PeopleGroup
+        read_only_fields = [
+            "id",
+            "slug",
+            "name",
+            "publication_status",
+            "header_image",
+            "children",
+            "roles",
+        ]
+        fields = read_only_fields
+
+    def get_children(
+        self, people_group: PeopleGroup
+    ) -> List[Dict[str, Union[str, int]]]:
+        context = self.context
+        request = context.get("request")
+        mapping = context.get("mapping")
+        if not mapping:
+            base_queryset = request.user.get_people_group_queryset().filter(
+                organization=people_group.organization
+            )
+            mapping = {group.id: group for group in base_queryset}
+            context["mapping"] = mapping
+        children_ids = list(people_group.children.all().values_list("id", flat=True))
+        if people_group.is_root:
+            children_ids += list(
+                PeopleGroup.objects.filter(
+                    organization=people_group.organization,
+                    parent__isnull=True,
+                    is_root=False,
+                ).values_list("id", flat=True)
+            )
+        children = [mapping.get(child) for child in children_ids if child in mapping]
+        return PeopleGroupHierarchySerializer(children, many=True, context=context).data
+
+
 class PeopleGroupAddTeamMembersSerializer(serializers.Serializer):
     people_group = HiddenPrimaryKeyRelatedField(
         required=False, write_only=True, queryset=PeopleGroup.objects.all()
@@ -386,23 +442,23 @@ def get_hierarchy(self, obj: PeopleGroup) -> List[Dict[str, Union[str, int]]]:
         while obj.parent and not obj.parent.is_root:
             obj = obj.parent
             if obj in queryset:
-                hierarchy.append({"id": obj.id, "slug": obj.slug, "name": obj.name})
+                hierarchy.append(
+                    PeopleGroupSuperLightSerializer(obj, context=self.context).data
+                )
         return [{"order": i, **h} for i, h in enumerate(hierarchy[::-1])]
 
     def get_children(self, obj: PeopleGroup) -> List[Dict[str, Union[str, int]]]:
         request = self.context.get("request")
         queryset = (
-            request.user.get_people_group_queryset() & obj.children.all().distinct()
+            request.user.get_people_group_queryset()
+            .select_related("organization")
+            .filter(parent=obj)
+            .order_by("name")
+            .distinct()
         )
-        return [
-            {
-                "id": child.id,
-                "slug": child.slug,
-                "name": child.name,
-                "organization": child.organization.code,
-            }
-            for child in queryset.order_by("name")
-        ]
+        return PeopleGroupSuperLightSerializer(
+            queryset, many=True, context=self.context
+        ).data
 
     def validate_featured_projects(self, projects: List[Project]) -> List[Project]:
         request = self.context.get("request")

apps/accounts/views.py

@@ -65,6 +65,7 @@
     EmptyPayloadResponseSerializer,
     PeopleGroupAddFeaturedProjectsSerializer,
     PeopleGroupAddTeamMembersSerializer,
+    PeopleGroupHierarchySerializer,
     PeopleGroupLightSerializer,
     PeopleGroupRemoveFeaturedProjectsSerializer,
     PeopleGroupRemoveTeamMembersSerializer,
@@ -830,7 +831,10 @@ def project(self, request, *args, **kwargs):
     def hierarchy(self, request, *args, **kwargs):
         people_group = self.get_object()
         return Response(
-            people_group.get_hierarchy(self.request.user), status=status.HTTP_200_OK
+            PeopleGroupHierarchySerializer(
+                people_group, context={"request": request}
+            ).data,
+            status=status.HTTP_200_OK,
         )
 
 

apps/organizations/models.py

@@ -2,7 +2,6 @@
 
 from django.contrib.auth.models import Group, Permission
 from django.contrib.contenttypes.models import ContentType
-from django.contrib.postgres.aggregates import ArrayAgg
 from django.contrib.postgres.fields import ArrayField
 from django.db import models
 from django.db.models import Q, QuerySet, UniqueConstraint
@@ -557,36 +556,6 @@ def update_or_create_root(cls, organization: "Organization"):
         )
         return root_group
 
-    @classmethod
-    def _get_hierarchy(cls, categories: dict[int, dict], category_id: int):
-        from apps.files.serializers import ImageSerializer
-
-        return {
-            "id": categories[category_id].id,
-            "slug": categories[category_id].slug,
-            "name": categories[category_id].name,
-            "background_color": categories[category_id].background_color,
-            "foreground_color": categories[category_id].foreground_color,
-            "background_image": (
-                ImageSerializer(categories[category_id].background_image).data
-                if categories[category_id].background_image
-                else None
-            ),
-            "children": [
-                cls._get_hierarchy(categories, child)
-                for child in categories[category_id].children_ids
-                if child is not None
-            ],
-        }
-
-    def get_hierarchy(self):
-        # This would be better with a recursive serializer, but it doubles the query time
-        categories = ProjectCategory.objects.filter(
-            organization=self.organization.pk
-        ).annotate(children_ids=ArrayAgg("children"))
-        categories = {category.id: category for category in categories}
-        return self._get_hierarchy(categories, self.id)
-
 
 class CategoryFollow(HasOwner, OrganizationRelated, models.Model):
     """Represent a user following a project category.

apps/organizations/serializers.py

@@ -427,6 +427,20 @@ def get_related_organizations(self) -> List[Organization]:
         return [self.instance.organization] if self.instance else []
 
 
+class ProjectCategorySuperLightSerializer(
+    AutoTranslatedModelSerializer,
+    serializers.ModelSerializer,
+):
+
+    class Meta:
+        model = ProjectCategory
+        fields = [
+            "id",
+            "slug",
+            "name",
+        ]
+
+
 class ProjectCategoryLightSerializer(
     AutoTranslatedModelSerializer,
     OrganizationRelatedSerializer,
@@ -551,6 +565,53 @@ def get_string_images_kwargs(
         }
 
 
+class ProjectCategoryHierarchySerializer(
+    AutoTranslatedModelSerializer,
+    OrganizationRelatedSerializer,
+    serializers.ModelSerializer,
+):
+    children = serializers.SerializerMethodField()
+    background_image = ImageSerializer(read_only=True)
+
+    class Meta:
+        model = ProjectCategory
+        read_only_fields = [
+            "id",
+            "slug",
+            "name",
+            "background_color",
+            "foreground_color",
+            "background_image",
+            "children",
+        ]
+        fields = read_only_fields
+
+    def get_children(
+        self, category: ProjectCategory
+    ) -> List[Dict[str, Union[str, int]]]:
+        context = self.context
+        mapping = context.get("mapping")
+        if not mapping:
+            queryset = ProjectCategory.objects.filter(
+                organization=category.organization
+            )
+            mapping = {cat.id: cat for cat in queryset}
+            context["mapping"] = mapping
+        children_ids = list(category.children.all().values_list("id", flat=True))
+        if category.is_root:
+            children_ids += list(
+                ProjectCategory.objects.filter(
+                    organization=category.organization,
+                    parent__isnull=True,
+                    is_root=False,
+                ).values_list("id", flat=True)
+            )
+        children = [mapping.get(child) for child in children_ids if child in mapping]
+        return ProjectCategoryHierarchySerializer(
+            children, many=True, context=context
+        ).data
+
+
 class ProjectCategorySerializer(
     StringsImagesSerializer,
     AutoTranslatedModelSerializer,
@@ -614,18 +675,16 @@ def get_hierarchy(self, obj: ProjectCategory) -> List[Dict[str, Union[str, int]]
         hierarchy = []
         while obj.parent and not obj.parent.is_root:
             obj = obj.parent
-            hierarchy.append({"id": obj.id, "slug": obj.slug, "name": obj.name})
+            hierarchy.append(
+                ProjectCategorySuperLightSerializer(obj, context=self.context).data
+            )
         return [{"order": i, **h} for i, h in enumerate(hierarchy[::-1])]
 
     def get_children(self, obj: ProjectCategory) -> List[Dict[str, Union[str, int]]]:
-        return [
-            {
-                "id": child.id,
-                "slug": child.slug,
-                "name": child.name,
-            }
-            for child in obj.children.all().order_by("name")
-        ]
+        queryset = obj.children.all().order_by("name")
+        return ProjectCategorySuperLightSerializer(
+            queryset, many=True, context=self.context
+        ).data
 
     def get_projects_count(self, obj: ProjectCategory) -> int:
         return obj.projects.count()