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.

Author: kkthxbye-code

netbox/dcim/api/views.py

@@ -19,6 +19,7 @@
 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):
         """

netbox/netbox/api/pagination.py

@@ -16,7 +16,7 @@ def __init__(self):
     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 @@ def get_limit(self, request):
 
         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 @@ def get_previous_link(self):
             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()