From 8ef74192ec78596788f425ac76bb162510569ea3 Mon Sep 17 00:00:00 2001
From: kkthxbye-code <festll234@gmail.com>
Date: Mon, 13 Jun 2022 20:45:08 +0200
Subject: [PATCH] Implement a custom paginator for DeviceViewSet Running count
 on the annotated query for loading config_context is slow. The custom
 paginator removes the annotation before getting the count.

---
 netbox/dcim/api/views.py        |  2 ++
 netbox/netbox/api/pagination.py | 18 +++++++++++++++++-
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/netbox/dcim/api/views.py b/netbox/dcim/api/views.py
index e99ef333a..c4c25f654 100644
--- a/netbox/dcim/api/views.py
+++ b/netbox/dcim/api/views.py
@@ -19,6 +19,7 @@ from ipam.models import Prefix, VLAN
 from netbox.api.authentication import IsAuthenticatedOrLoginNotRequired
 from netbox.api.exceptions import ServiceUnavailable
 from netbox.api.metadata import ContentTypeMetadata
+from netbox.api.pagination import StripCountAnnotationsPaginator
 from netbox.api.viewsets import NetBoxModelViewSet
 from netbox.config import get_config
 from utilities.api import get_serializer_for_model
@@ -392,6 +393,7 @@ class DeviceViewSet(ConfigContextQuerySetMixin, NetBoxModelViewSet):
         'virtual_chassis__master', 'primary_ip4__nat_outside', 'primary_ip6__nat_outside', 'tags',
     )
     filterset_class = filtersets.DeviceFilterSet
+    pagination_class = StripCountAnnotationsPaginator
 
     def get_serializer_class(self):
         """
diff --git a/netbox/netbox/api/pagination.py b/netbox/netbox/api/pagination.py
index d89e32124..5ecade264 100644
--- a/netbox/netbox/api/pagination.py
+++ b/netbox/netbox/api/pagination.py
@@ -16,7 +16,7 @@ class OptionalLimitOffsetPagination(LimitOffsetPagination):
     def paginate_queryset(self, queryset, request, view=None):
 
         if isinstance(queryset, QuerySet):
-            self.count = queryset.count()
+            self.count = self.get_queryset_count(queryset)
         else:
             # We're dealing with an iterable, not a QuerySet
             self.count = len(queryset)
@@ -52,6 +52,9 @@ class OptionalLimitOffsetPagination(LimitOffsetPagination):
 
         return self.default_limit
 
+    def get_queryset_count(self, queryset):
+        return queryset.count()
+
     def get_next_link(self):
 
         # Pagination has been disabled
@@ -67,3 +70,16 @@ class OptionalLimitOffsetPagination(LimitOffsetPagination):
             return None
 
         return super().get_previous_link()
+
+
+class StripCountAnnotationsPaginator(OptionalLimitOffsetPagination):
+    """
+    Strips the annotations on the queryset before getting the count
+    to optimize pagination of complex queries.
+    """
+    def get_queryset_count(self, queryset):
+        # Clone the queryset to avoid messing up the actual query
+        cloned_queryset = queryset.all()
+        cloned_queryset.query.annotations.clear()
+
+        return cloned_queryset.count()