From ca8824d1e3870735895defd008de39f8f4ee3b4a Mon Sep 17 00:00:00 2001
From: sadnub <josh@torchlake.com>
Date: Sun, 14 Aug 2022 21:49:38 -0400
Subject: [PATCH] fix clients not filtering by role in policy overview

---
 api/tacticalrmm/automation/serializers.py | 27 ++++++++++++++++++++---
 api/tacticalrmm/automation/views.py       | 17 +++++++++++---
 2 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/api/tacticalrmm/automation/serializers.py b/api/tacticalrmm/automation/serializers.py
index cfca78ee..9ea41fce 100644
--- a/api/tacticalrmm/automation/serializers.py
+++ b/api/tacticalrmm/automation/serializers.py
@@ -7,8 +7,11 @@ from rest_framework.serializers import (
 from agents.serializers import AgentHostnameSerializer
 from autotasks.models import TaskResult
 from checks.models import CheckResult
-from clients.models import Client
-from clients.serializers import ClientMinimumSerializer, SiteMinimumSerializer
+from clients.models import Client, Site
+from clients.serializers import (
+    ClientMinimumSerializer,
+    SiteMinimumSerializer,
+)
 from winupdate.serializers import WinUpdatePolicySerializer
 
 from .models import Policy
@@ -85,11 +88,29 @@ class PolicyRelatedSerializer(ModelSerializer):
         )
 
 
+class PolicyOverviewSiteSerializer(ModelSerializer):
+    workstation_policy = PolicySerializer(read_only=True)
+    server_policy = PolicySerializer(read_only=True)
+
+    class Meta:
+        model = Site
+        fields = ("pk", "name", "workstation_policy", "server_policy")
+
+
 class PolicyOverviewSerializer(ModelSerializer):
+    sites = SerializerMethodField()
+    workstation_policy = PolicySerializer(read_only=True)
+    server_policy = PolicySerializer(read_only=True)
+
+    def get_sites(self, obj):
+        return PolicyOverviewSiteSerializer(
+            obj.filtered_sites,
+            many=True,
+        ).data
+
     class Meta:
         model = Client
         fields = ("pk", "name", "sites", "workstation_policy", "server_policy")
-        depth = 2
 
 
 class PolicyCheckStatusSerializer(ModelSerializer):
diff --git a/api/tacticalrmm/automation/views.py b/api/tacticalrmm/automation/views.py
index bcd4d317..5b1d50df 100644
--- a/api/tacticalrmm/automation/views.py
+++ b/api/tacticalrmm/automation/views.py
@@ -7,10 +7,11 @@ from rest_framework.views import APIView
 from agents.models import Agent
 from autotasks.models import TaskResult
 from checks.models import CheckResult
-from clients.models import Client
+from clients.models import Client, Site
 from tacticalrmm.permissions import _has_perm_on_client, _has_perm_on_site
 from winupdate.models import WinUpdatePolicy
 from winupdate.serializers import WinUpdatePolicySerializer
+from django.db.models import Prefetch
 
 from .models import Policy
 from .permissions import AutomationPolicyPerms
@@ -108,8 +109,18 @@ class PolicyCheck(APIView):
 class OverviewPolicy(APIView):
     def get(self, request):
 
-        clients = Client.objects.filter_by_role(request.user).select_related(
-            "workstation_policy", "server_policy"
+        clients = (
+            Client.objects.filter_by_role(request.user)
+            .select_related("workstation_policy", "server_policy")
+            .prefetch_related(
+                Prefetch(
+                    "sites",
+                    queryset=Site.objects.select_related(
+                        "workstation_policy", "server_policy"
+                    ),
+                    to_attr="filtered_sites",
+                )
+            )
         )
         return Response(PolicyOverviewSerializer(clients, many=True).data)