fix: prefetch right geo custom zones - permissions added and fix applied to detection_object views

Author: stephenprn

core/permissions/geo_custom_zone.py

@@ -0,0 +1,57 @@
+from typing import Optional
+from core.models.geo_custom_zone import GeoCustomZone, GeoCustomZoneStatus
+from core.models.geo_custom_zone_category import GeoCustomZoneCategory
+from core.models.user import User, UserRole
+from core.permissions.base import BasePermission
+from core.repository.geo_custom_zone import GeoCustomZoneRepository
+
+from django.db.models import QuerySet
+from django.db.models import Prefetch
+
+
+class GeoCustomZonePermission(
+    BasePermission[GeoCustomZone],
+):
+    def __init__(
+        self, user: User, initial_queryset: Optional[QuerySet[GeoCustomZone]] = None
+    ):
+        self.repository = GeoCustomZoneRepository(initial_queryset=initial_queryset)
+        self.user = user
+
+    def _get_prefetch(self, lookup_root: str = ""):
+        if self.user.user_role == UserRole.SUPER_ADMIN:
+            geo_custom_zones_prefetch = Prefetch(
+                f"{lookup_root}geo_custom_zones",
+                queryset=GeoCustomZone.objects.filter(
+                    geo_custom_zone_status=GeoCustomZoneStatus.ACTIVE,
+                ),
+            )
+            geo_custom_zones_category_prefetch = Prefetch(
+                f"{lookup_root}geo_custom_zones__geo_custom_zone_category",
+                queryset=GeoCustomZoneCategory.objects.filter(
+                    geo_custom_zones__geo_custom_zone_status=GeoCustomZoneStatus.ACTIVE,
+                ),
+            )
+        else:
+            geo_custom_zones_prefetch = Prefetch(
+                f"{lookup_root}geo_custom_zones",
+                queryset=GeoCustomZone.objects.filter(
+                    geo_custom_zone_status=GeoCustomZoneStatus.ACTIVE,
+                    user_groups_custom_geo_zones__user_user_groups__user=self.user.id,
+                ),
+            )
+            geo_custom_zones_category_prefetch = Prefetch(
+                f"{lookup_root}geo_custom_zones__geo_custom_zone_category",
+                queryset=GeoCustomZoneCategory.objects.filter(
+                    geo_custom_zones__geo_custom_zone_status=GeoCustomZoneStatus.ACTIVE,
+                    geo_custom_zones__user_groups_custom_geo_zones__user_user_groups__user=self.user.id,
+                ),
+            )
+
+        return geo_custom_zones_prefetch, geo_custom_zones_category_prefetch
+
+    def get_detection_object_prefetch(self):
+        return self._get_prefetch()
+
+    def get_detection_prefetch(self):
+        return self._get_prefetch("detection_object__")

core/repository/geo_custom_zone.py

@@ -0,0 +1,49 @@
+from typing import List, Optional
+from core.models.geo_custom_zone import GeoCustomZone
+from core.repository.base import (
+    BaseRepository,
+    DateRepoFilter,
+    TimestampedBaseRepositoryMixin,
+    UuidBaseRepositoryMixin,
+)
+from django.db.models import QuerySet
+
+
+class GeoCustomZoneRepository(
+    BaseRepository[GeoCustomZone],
+    TimestampedBaseRepositoryMixin[GeoCustomZone],
+    UuidBaseRepositoryMixin[GeoCustomZone],
+):
+    def __init__(self, initial_queryset: Optional[QuerySet[GeoCustomZone]] = None):
+        self.model = GeoCustomZone
+        self.initial_queryset = (
+            initial_queryset if initial_queryset is not None else self.model.objects
+        )
+
+    def filter_(
+        self,
+        queryset: QuerySet[GeoCustomZone],
+        filter_created_at: Optional[DateRepoFilter] = None,
+        filter_updated_at: Optional[DateRepoFilter] = None,
+        filter_uuid_in: Optional[List[str]] = None,
+        filter_uuid_notin: Optional[List[str]] = None,
+        order_bys: Optional[List[str]] = None,
+        *args,
+        **kwargs,
+    ) -> QuerySet[GeoCustomZone]:
+        # mixin filters
+
+        queryset = self._filter_timestamped(
+            queryset=queryset,
+            filter_created_at=filter_created_at,
+            filter_updated_at=filter_updated_at,
+        )
+        queryset = self._filter_uuid(
+            queryset=queryset,
+            filter_uuid_in=filter_uuid_in,
+            filter_uuid_notin=filter_uuid_notin,
+        )
+
+        queryset = self.order_by(queryset=queryset, order_bys=order_bys)
+
+        return queryset

core/views/detection/detection_list.py

@@ -22,10 +22,8 @@
     DetectionPrescriptionStatus,
     DetectionValidationStatus,
 )
-from core.models.geo_custom_zone import GeoCustomZone, GeoCustomZoneStatus
-from core.models.geo_custom_zone_category import GeoCustomZoneCategory
 from core.models.tile_set import TileSet, TileSetStatus, TileSetType
-from core.models.user import UserRole
+from core.permissions.geo_custom_zone import GeoCustomZonePermission
 from core.permissions.tile_set import TileSetPermission
 from core.permissions.user import UserPermission
 from core.repository.base import NumberRepoFilter, RepoFilterLookup
@@ -279,27 +277,9 @@ def filter_queryset(self, queryset):
             "detection_object__commune__geometry",
         )
 
-        # TODO: move this in permission file and apply this to geo
-        if self.request.user.user_role == UserRole.SUPER_ADMIN:
-            geo_custom_zones_prefetch = Prefetch("detection_object__geo_custom_zones")
-            geo_custom_zones_category_prefetch = Prefetch(
-                "detection_object__geo_custom_zones__geo_custom_zone_category"
-            )
-        else:
-            geo_custom_zones_prefetch = Prefetch(
-                "detection_object__geo_custom_zones",
-                queryset=GeoCustomZone.objects.filter(
-                    geo_custom_zone_status=GeoCustomZoneStatus.ACTIVE,
-                    user_groups_custom_geo_zones__user_user_groups__user=self.request.user.id,
-                ),
-            )
-            geo_custom_zones_category_prefetch = Prefetch(
-                "detection_object__geo_custom_zones__geo_custom_zone_category",
-                queryset=GeoCustomZoneCategory.objects.filter(
-                    geo_custom_zones__geo_custom_zone_status=GeoCustomZoneStatus.ACTIVE,
-                    geo_custom_zones__user_groups_custom_geo_zones__user_user_groups__user=self.request.user.id,
-                ),
-            )
+        geo_custom_zones_prefetch, geo_custom_zones_category_prefetch = (
+            GeoCustomZonePermission(user=self.request.user).get_detection_prefetch()
+        )
 
         queryset = queryset.prefetch_related(
             "detection_object",

core/views/detection_object.py

@@ -6,6 +6,7 @@
 from core.contants.geo import SRID
 from core.models.detection_object import DetectionObject
 from core.models.tile_set import TileSetType
+from core.permissions.geo_custom_zone import GeoCustomZonePermission
 from core.permissions.tile_set import TileSetPermission
 from core.serializers.detection_object import (
     DetectionObjectDetailSerializer,
@@ -91,8 +92,15 @@ def get_queryset(self):
         queryset = queryset.defer("parcel__geometry")
 
         if self.action == "retrieve":
+            geo_custom_zones_prefetch, geo_custom_zones_category_prefetch = (
+                GeoCustomZonePermission(
+                    user=self.request.user
+                ).get_detection_object_prefetch()
+            )
+
             queryset = queryset.prefetch_related(
-                "geo_custom_zones", "geo_sub_custom_zones"
+                geo_custom_zones_prefetch,
+                geo_custom_zones_category_prefetch,
             )
             queryset = queryset.defer(
                 "geo_custom_zones__geometry", "geo_sub_custom_zones__geometry"