From cc3c3e7e9a965777e084463bdb62affc55c4a36c Mon Sep 17 00:00:00 2001
From: Dana Lambert <danalambert@catalyst.net.nz>
Date: Wed, 20 Oct 2021 12:36:48 +1300
Subject: [PATCH] Add soil variant filter

---
 backend/right_tree/api/filters.py | 15 +++++++++++++--
 backend/right_tree/api/views.py   |  5 +++--
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/backend/right_tree/api/filters.py b/backend/right_tree/api/filters.py
index 0c9e337..a6bada7 100644
--- a/backend/right_tree/api/filters.py
+++ b/backend/right_tree/api/filters.py
@@ -3,7 +3,7 @@ import json
 from django.contrib.gis.geos import Point
 from django.db.models import Q
 
-from .models import EcologicalRegion, EcologicalDistrictLayer, SoilOrder
+from .models import EcologicalRegion, EcologicalDistrictLayer, SoilOrder, SoilVariant
 
 
 def coordinate_filter(request, queryset):
@@ -23,6 +23,17 @@ def coordinate_filter(request, queryset):
         # Filter by ecological regions and soil orders
         return queryset.filter(
             Q(ecological_regions__in=filtered_regions) &
-            Q(soil_order__in=filtered_soil_orders))
+            Q(soil_order__in=filtered_soil_orders)).distinct()
+
+    return queryset
+
+
+def soil_variant_filter(request, queryset):
+    soil_variant = request.query_params.get('soilVariant')
+
+    if soil_variant in {"D", "W", "M"}:
+        soil_variant_ids = SoilVariant.objects.filter(Q(name__startswith=soil_variant) | Q(
+            name__startswith="M")).values_list('id', flat=True).distinct()
+        return queryset.filter(soil_variants__in=soil_variant_ids).distinct()
 
     return queryset
diff --git a/backend/right_tree/api/views.py b/backend/right_tree/api/views.py
index c1b9949..eb7f952 100644
--- a/backend/right_tree/api/views.py
+++ b/backend/right_tree/api/views.py
@@ -1,11 +1,11 @@
 from rest_framework import viewsets
 from right_tree.api.models import Plant
 from right_tree.api.serializers import PlantSerializer
-from .filters import coordinate_filter
+from .filters import coordinate_filter, soil_variant_filter
 
 class PlantViewSet(viewsets.ModelViewSet):
     """
-    API endpoint that allows hours to be grouped or edited.
+    Filtered viewset for plants.
     """
     queryset = Plant.objects.all()
     serializer_class = PlantSerializer
@@ -16,6 +16,7 @@ class PlantViewSet(viewsets.ModelViewSet):
         """
         queryset = Plant.objects.all()
         queryset = coordinate_filter(self.request, queryset)
+        queryset = soil_variant_filter(self.request, queryset)
 
         return queryset
    
\ No newline at end of file