move maintenance mode and agent count checks to annotate to load client tree faster

api/tacticalrmm/clients/migrations/0021_remove_client_agent_count_remove_site_agent_count.py

@@ -0,0 +1,21 @@
+# Generated by Django 4.0.3 on 2022-04-15 01:53
+
+from django.db import migrations
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('clients', '0020_auto_20211226_0547'),
+    ]
+
+    operations = [
+        migrations.RemoveField(
+            model_name='client',
+            name='agent_count',
+        ),
+        migrations.RemoveField(
+            model_name='site',
+            name='agent_count',
+        ),
+    ]

api/tacticalrmm/clients/models.py

@@ -20,7 +20,6 @@ class Client(BaseAuditModel):
     name = models.CharField(max_length=255, unique=True)
     block_policy_inheritance = models.BooleanField(default=False)
     failing_checks = models.JSONField(default=_default_failing_checks_data)
-    agent_count = models.PositiveIntegerField(default=0)
     workstation_policy = models.ForeignKey(
         "automation.Policy",
         related_name="workstation_clients",
@@ -82,14 +81,6 @@ class Client(BaseAuditModel):
     def __str__(self):
         return self.name
 
-    @property
-    def has_maintenanace_mode_agents(self) -> bool:
-        return (
-            Agent.objects.defer(*AGENT_DEFER)
-            .filter(site__client=self, maintenance_mode=True)
-            .exists()
-        )
-
     @property
     def live_agent_count(self) -> int:
         return Agent.objects.defer(*AGENT_DEFER).filter(site__client=self).count()
@@ -109,7 +100,6 @@ class Site(BaseAuditModel):
     name = models.CharField(max_length=255)
     block_policy_inheritance = models.BooleanField(default=False)
     failing_checks = models.JSONField(default=_default_failing_checks_data)
-    agent_count = models.PositiveIntegerField(default=0)
     workstation_policy = models.ForeignKey(
         "automation.Policy",
         related_name="workstation_sites",
@@ -166,10 +156,6 @@ class Site(BaseAuditModel):
     def __str__(self):
         return self.name
 
-    @property
-    def has_maintenanace_mode_agents(self) -> bool:
-        return self.agents.defer(*AGENT_DEFER).filter(maintenance_mode=True).exists()  # type: ignore
-
     @property
     def live_agent_count(self) -> int:
         return self.agents.defer(*AGENT_DEFER).count()  # type: ignore

api/tacticalrmm/clients/serializers.py

@@ -5,7 +5,9 @@ from rest_framework.serializers import (
     ValidationError,
 )
 
+from django.db.models import OuterRef, Exists, Count, Prefetch
 from .models import Client, ClientCustomField, Deployment, Site, SiteCustomField
+from agents.models import Agent
 
 
 class SiteCustomFieldSerializer(ModelSerializer):
@@ -30,7 +32,8 @@ class SiteCustomFieldSerializer(ModelSerializer):
 class SiteSerializer(ModelSerializer):
     client_name = ReadOnlyField(source="client.name")
     custom_fields = SiteCustomFieldSerializer(many=True, read_only=True)
-    maintenance_mode = ReadOnlyField(source="has_maintenanace_mode_agents")
+    maintenance_mode = ReadOnlyField()
+    agent_count = ReadOnlyField()
 
     class Meta:
         model = Site
@@ -92,13 +95,19 @@ class ClientCustomFieldSerializer(ModelSerializer):
 class ClientSerializer(ModelSerializer):
     sites = SerializerMethodField()
     custom_fields = ClientCustomFieldSerializer(many=True, read_only=True)
-    maintenance_mode = ReadOnlyField(source="has_maintenanace_mode_agents")
+    maintenance_mode = ReadOnlyField()
+    agent_count = ReadOnlyField()
 
     def get_sites(self, obj):
         return SiteSerializer(
             obj.sites.select_related("client")
             .filter_by_role(self.context["user"])
-            .prefetch_related("custom_fields__field"),
+            .annotate(
+                maintenance_mode=Exists(
+                    Agent.objects.filter(site=OuterRef("pk"), maintenance_mode=True)
+                )
+            )
+            .annotate(agent_count=Count("agents")),
             many=True,
         ).data
 

api/tacticalrmm/clients/views.py

@@ -2,11 +2,11 @@ import datetime as dt
 import re
 import uuid
 
-import pytz
 from agents.models import Agent
 from core.utils import get_core_settings
 from django.shortcuts import get_object_or_404
 from django.utils import timezone as djangotime
+from django.db.models import OuterRef, Exists, Count, Prefetch
 from rest_framework.exceptions import PermissionDenied
 from rest_framework.permissions import AllowAny, IsAuthenticated
 from rest_framework.response import Response
@@ -36,7 +36,20 @@ class GetAddClients(APIView):
                 "workstation_policy", "server_policy", "alert_template"
             )
             .filter_by_role(request.user)  # type: ignore
-            .prefetch_related("custom_fields__field", "sites")
+            .prefetch_related(
+                Prefetch(
+                    "custom_fields",
+                    queryset=ClientCustomField.objects.select_related("field"),
+                ),
+            )
+            .annotate(
+                maintenance_mode=Exists(
+                    Agent.objects.filter(
+                        site__client=OuterRef("pk"), maintenance_mode=True
+                    )
+                )
+            )
+            .annotate(agent_count=Count("sites__agents"))
         )
         return Response(
             ClientSerializer(clients, context={"user": request.user}, many=True).data

api/tacticalrmm/core/tasks.py

@@ -100,14 +100,12 @@ def cache_db_fields_task() -> None:
     for site in Site.objects.all():
         agents = site.agents.defer(*AGENT_DEFER)
         site.failing_checks = _get_failing_data(agents)
-        site.agent_count = agents.count()
-        site.save(update_fields=["failing_checks", "agent_count"])
+        site.save(update_fields=["failing_checks"])
 
     for client in Client.objects.all():
         agents = Agent.objects.defer(*AGENT_DEFER).filter(site__client=client)
         client.failing_checks = _get_failing_data(agents)
-        client.agent_count = agents.count()
-        client.save(update_fields=["failing_checks", "agent_count"])
+        client.save(update_fields=["failing_checks"])
 
     for agent in Agent.objects.defer(*AGENT_DEFER):
         if (