Merge pull request #1273 from BLSQ/IA-2907-api-tree-picker

IA-2907 New API endpoint for Org Unit Tree

iaso/admin.py

@@ -1,38 +1,17 @@
 from typing import Any
 from typing import Protocol
 
+from django import forms as django_forms
 from django.contrib.admin import widgets
 from django.contrib.gis import admin, forms
 from django.contrib.gis.db import models as geomodels
+from django.contrib.postgres.fields import ArrayField
 from django.db import models
 from django.utils.html import format_html_join, format_html
 from django.utils.safestring import mark_safe
 from django_json_widget.widgets import JSONEditorWidget
 
 from iaso.models.json_config import Config  # type: ignore
-
-
-class IasoJSONEditorWidget(JSONEditorWidget):
-    class Media:
-        css = {"all": ("css/admin-json-widget.css",)}
-
-    def __init__(self, attrs=None, mode="code", options=None, width=None, height=None):
-        if height == None:
-            height = "400px"
-
-        default_options = {
-            "modes": ["text", "code"],
-            "mode": mode,
-            "search": True,
-        }
-        if options:
-            default_options.update(options)
-
-        super(IasoJSONEditorWidget, self).__init__(
-            attrs=attrs, mode=mode, options=default_options, width=width, height=height
-        )
-
-
 from .models import (
     Account,
     AccountFeatureFlag,
@@ -87,11 +66,58 @@ def __init__(self, attrs=None, mode="code", options=None, width=None, height=Non
     Payment,
     PaymentLot,
 )
-from .models.microplanning import Team, Planning, Assignment
 from .models.data_store import JsonDataStore
+from .models.microplanning import Team, Planning, Assignment
 from .utils.gis import convert_2d_point_to_3d
 
 
+class IasoJSONEditorWidget(JSONEditorWidget):
+    class Media:
+        css = {"all": ("css/admin-json-widget.css",)}
+
+    def __init__(self, attrs=None, mode="code", options=None, width=None, height=None):
+        if height == None:
+            height = "400px"
+
+        default_options = {
+            "modes": ["text", "code"],
+            "mode": mode,
+            "search": True,
+        }
+        if options:
+            default_options.update(options)
+
+        super(IasoJSONEditorWidget, self).__init__(
+            attrs=attrs, mode=mode, options=default_options, width=width, height=height
+        )
+
+
+class ArrayFieldMultipleChoiceField(django_forms.MultipleChoiceField):
+    """
+    Display a multi-select field for ArrayField:
+
+    formfield_overrides = {
+        ArrayField: {
+            "form_class": ArrayFieldMultipleChoiceField,
+        }
+    }
+
+    formfield_overrides = {
+        ArrayField: {
+            "form_class": ArrayFieldMultipleChoiceField,
+            "widget": forms.CheckboxSelectMultiple,
+        }
+    }
+    """
+
+    def __init__(self, *args, **kwargs):
+        kwargs.pop("max_length", None)
+        base_field = kwargs.pop("base_field", None)
+        kwargs["choices"] = base_field.choices
+        kwargs["choices"].pop(0)
+        super().__init__(*args, **kwargs)
+
+
 class AdminAttributes(Protocol):
     """Workaround to avoid mypy errors, see https://github.com/python/mypy/issues/2087#issuecomment-462726600"""
 
@@ -786,10 +812,19 @@ class PaymentLotAdmin(admin.ModelAdmin):
     formfield_overrides = {models.JSONField: {"widget": IasoJSONEditorWidget}}
 
 
+@admin.register(DataSource)
+class DataSourceAdmin(admin.ModelAdmin):
+    formfield_overrides = {
+        ArrayField: {
+            "form_class": ArrayFieldMultipleChoiceField,
+            "widget": forms.CheckboxSelectMultiple,
+        }
+    }
+
+
 admin.site.register(Account)
 admin.site.register(AccountFeatureFlag)
 admin.site.register(Device)
-admin.site.register(DataSource)
 admin.site.register(DeviceOwnership)
 admin.site.register(MatchingAlgorithm)
 admin.site.register(ExternalCredentials)

iaso/api/data_sources.py

@@ -31,6 +31,7 @@ class Meta:
             "projects",
             "default_version",
             "credentials",
+            "tree_config_status_fields",
         ]
 
     url = serializers.SerializerMethodField()
@@ -243,7 +244,12 @@ def get_queryset(self):
         linked_to = self.kwargs.get("linkedTo", None)
         profile = self.request.user.iaso_profile
         order = self.request.GET.get("order", "name").split(",")
-        sources = DataSource.objects.filter(projects__account=profile.account).distinct()
+        sources = (
+            DataSource.objects.select_related("default_version", "credentials")
+            .prefetch_related("projects", "versions")
+            .filter(projects__account=profile.account)
+            .distinct()
+        )
         if linked_to:
             org_unit = OrgUnit.objects.get(pk=linked_to)
             useful_sources = org_unit.source_set.values_list("algorithm_run__version_2__data_source_id", flat=True)

iaso/api/org_unit_tree/filters.py

@@ -0,0 +1,35 @@
+import django_filters
+
+from django import forms
+from django.db.models.query import QuerySet
+from django.utils.translation import gettext_lazy as _
+
+from rest_framework.exceptions import ValidationError
+
+from iaso.models import OrgUnit, DataSource
+
+
+class OrgUnitTreeFilter(django_filters.rest_framework.FilterSet):
+    ignore_empty_names = django_filters.BooleanFilter(method="filter_empty_names", label=_("Ignore empty names"))
+    parent_id = django_filters.NumberFilter(field_name="parent_id", label=_("Parent ID"))
+    data_source_id = django_filters.NumberFilter(method="filter_data_source_id", label=_("Data source ID"))
+    validation_status = django_filters.MultipleChoiceFilter(
+        choices=OrgUnit.VALIDATION_STATUS_CHOICES, label=_("Validation status"), widget=forms.CheckboxSelectMultiple
+    )
+    search = django_filters.CharFilter(field_name="name", lookup_expr="icontains")
+
+    class Meta:
+        model = OrgUnit
+        fields = ["version"]
+
+    def filter_empty_names(self, queryset: QuerySet, _, use_empty_names: bool) -> QuerySet:
+        return queryset.exclude(name="") if use_empty_names else queryset
+
+    def filter_data_source_id(self, queryset: QuerySet, _, data_source_id: int) -> QuerySet:
+        try:
+            source = DataSource.objects.get(id=data_source_id)
+        except OrgUnit.DoesNotExist:
+            raise ValidationError({"data_source_id": [f"DataSource with id {data_source_id} does not exist."]})
+        if source.default_version:
+            return queryset.filter(version=source.default_version)
+        return queryset.filter(version__data_source_id=data_source_id)

iaso/api/org_unit_tree/pagination.py

@@ -0,0 +1,5 @@
+from iaso.api.common import Paginator
+
+
+class OrgUnitTreePagination(Paginator):
+    page_size = 10

iaso/api/org_unit_tree/serializers.py

@@ -0,0 +1,37 @@
+from typing import Union
+
+from rest_framework import serializers
+
+from iaso.api.common import DynamicFieldsModelSerializer
+from iaso.models import OrgUnit
+
+
+class OrgUnitTreeSerializer(DynamicFieldsModelSerializer, serializers.ModelSerializer):
+    has_children = serializers.SerializerMethodField()
+    org_unit_type_short_name = serializers.SerializerMethodField()
+
+    @classmethod
+    def get_has_children(cls, org_unit: OrgUnit) -> bool:
+        return org_unit.children_count > 0
+
+    def get_org_unit_type_short_name(self, org_unit: OrgUnit) -> Union[str, None]:
+        return org_unit.org_unit_type.short_name if org_unit.org_unit_type else None
+
+    class Meta:
+        model = OrgUnit
+        fields = [
+            "id",
+            "name",
+            "validation_status",
+            "has_children",
+            "org_unit_type_id",
+            "org_unit_type_short_name",
+        ]
+        default_fields = [
+            "id",
+            "name",
+            "validation_status",
+            "has_children",
+            "org_unit_type_id",
+            "org_unit_type_short_name",
+        ]

iaso/api/org_unit_tree/views.py

@@ -0,0 +1,93 @@
+import django_filters
+
+from django.db.models import Count, Q
+
+from rest_framework import filters, permissions, viewsets
+from rest_framework import serializers
+from rest_framework.decorators import action
+from rest_framework.exceptions import ValidationError
+
+from iaso.api.org_unit_tree.filters import OrgUnitTreeFilter
+from iaso.api.org_unit_tree.pagination import OrgUnitTreePagination
+from iaso.api.org_unit_tree.serializers import OrgUnitTreeSerializer
+from iaso.models import OrgUnit
+
+
+class OrgUnitTreeQuerystringSerializer(serializers.Serializer):
+    force_full_tree = serializers.BooleanField(required=False)
+    parent_id = serializers.IntegerField(required=False)
+    data_source_id = serializers.IntegerField(required=False)
+    validation_status = serializers.MultipleChoiceField(choices=OrgUnit.VALIDATION_STATUS_CHOICES)
+
+
+class OrgUnitTreeViewSet(viewsets.ModelViewSet):
+    """
+    Explore the OrgUnit tree level by level.
+    """
+
+    filter_backends = [filters.OrderingFilter, django_filters.rest_framework.DjangoFilterBackend]
+    filterset_class = OrgUnitTreeFilter
+    http_method_names = ["get", "options", "head", "trace"]
+    ordering_fields = ["id", "name"]
+    pagination_class = None  # Since results are displayed level by level, results are not paginated in the list view.
+    permission_classes = [permissions.AllowAny]
+    serializer_class = OrgUnitTreeSerializer
+
+    def get_queryset(self):
+        user = self.request.user
+
+        querystring = OrgUnitTreeQuerystringSerializer(data=self.request.query_params)
+        querystring.is_valid(raise_exception=True)
+        force_full_tree = querystring.validated_data.get("force_full_tree")
+        parent_id = querystring.validated_data.get("parent_id")
+        data_source_id = querystring.validated_data.get("data_source_id")
+        validation_status = querystring.validated_data.get("validation_status", set())
+
+        if user.is_anonymous:
+            if not data_source_id:
+                raise ValidationError({"data_source_id": ["A `data_source_id` must be provided for anonymous users."]})
+            qs = OrgUnit.objects.all()  # `qs` will be filtered by `data_source_id` in `OrgUnitTreeFilter`.
+        elif user.is_superuser or force_full_tree:
+            qs = OrgUnit.objects.filter(version=user.iaso_profile.account.default_version)
+        else:
+            qs = OrgUnit.objects.filter_for_user(user)
+
+        can_view_full_tree = any([user.is_anonymous, user.is_superuser, force_full_tree])
+        display_root_level = not parent_id
+
+        if display_root_level and self.action == "list":
+            if can_view_full_tree:
+                qs = qs.filter(parent__isnull=True)
+            elif user.is_authenticated:
+                if user.iaso_profile.org_units.exists():
+                    # Root level of the tree for this user (the user may be restricted to a subpart of the tree).
+                    qs = qs.filter(id__in=user.iaso_profile.org_units.all())
+
+        qs = qs.only("id", "name", "validation_status", "version", "org_unit_type", "parent")
+        qs = qs.order_by("name")
+        qs = qs.select_related("org_unit_type")
+
+        if validation_status == {OrgUnit.VALIDATION_VALID}:
+            exclude_filter = ~Q(orgunit__validation_status__in=[OrgUnit.VALIDATION_REJECTED, OrgUnit.VALIDATION_NEW])
+            children_count = Count("orgunit", filter=exclude_filter)
+        else:
+            children_count = Count("orgunit")
+        qs = qs.annotate(children_count=children_count)
+
+        return qs
+
+    @action(detail=False, methods=["get"])
+    def search(self, request):
+        """
+        Search the OrgUnit tree.
+
+        ```
+        /api/orgunits/tree/search/?search=congo&page=2&limit=10
+        ```
+        """
+        org_units = self.get_queryset()
+        paginator = OrgUnitTreePagination()
+        filtered_org_units = self.filterset_class(request.query_params, org_units).qs
+        paginated_org_units = paginator.paginate_queryset(filtered_org_units, request)
+        serializer = OrgUnitTreeSerializer(paginated_org_units, many=True, context={"request": request})
+        return paginator.get_paginated_response(serializer.data)

iaso/api/org_units.py

@@ -336,6 +336,9 @@ def list_to_gpkg(self, queryset, filename):
 
     @action(methods=["GET"], detail=False)
     def treesearch(self, request, **kwargs):
+        """
+        TODO: delete this route when it's been replaced by `OrgUnitTreeViewSet`.
+        """
         queryset = self.get_queryset().order_by("name")
         params = request.GET
         parent_id = params.get("parent_id")

iaso/api/profiles.py

@@ -305,6 +305,7 @@ def get_row(profile: Profile, **_) -> List[Any]:
 
         filename = "users"
         response: Union[HttpResponse, StreamingHttpResponse]
+        queryset = queryset.order_by("id")
 
         if file_format == FileFormatEnum.XLSX:
             filename = filename + ".xlsx"

iaso/migrations/0280_datasource_tree_config_status_fields.py

@@ -0,0 +1,24 @@
+# Generated by Django 4.2.13 on 2024-05-14 14:23
+
+import django.contrib.postgres.fields
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("iaso", "0279_merge_20240417_1319"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="datasource",
+            name="tree_config_status_fields",
+            field=django.contrib.postgres.fields.ArrayField(
+                base_field=models.CharField(blank=True, choices=[], max_length=30),
+                blank=True,
+                default=list,
+                help_text="List of statuses used for display configuration of the OrgUnit tree.",
+                size=None,
+            ),
+        ),
+    ]

iaso/models/data_source.py

@@ -1,3 +1,4 @@
+from django.contrib.postgres.fields import ArrayField
 from django.db import models
 
 
@@ -19,10 +20,30 @@ class DataSource(models.Model):
     created_at = models.DateTimeField(auto_now_add=True)
     updated_at = models.DateTimeField(auto_now=True)
     default_version = models.ForeignKey("SourceVersion", null=True, blank=True, on_delete=models.SET_NULL)
+    tree_config_status_fields = ArrayField(
+        models.CharField(max_length=30, blank=True, choices=[]),
+        default=list,
+        blank=True,
+        help_text="List of statuses used for display configuration of the OrgUnit tree.",
+    )
 
     def __str__(self):
         return "%s " % (self.name,)
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # Override `__init__` to avoid a circular import.
+        # This could be replaced by a callable in Django ≥ 5.
+        # https://docs.djangoproject.com/en/5.0/releases/5.0/#more-options-for-declaring-field-choices
+        from iaso.models import OrgUnit
+
+        self._meta.get_field("tree_config_status_fields").base_field.choices = OrgUnit.VALIDATION_STATUS_CHOICES
+
+    def clean(self, *args, **kwargs):
+        super().clean()
+        self.tree_config_status_fields = list(set(self.tree_config_status_fields))
+
     def as_dict(self):
         versions = SourceVersion.objects.filter(data_source_id=self.id)
         return {