move get_queryset() to common mixin

netbox/dcim/api/views.py

@@ -24,7 +24,7 @@ from dcim.models import (
     VirtualChassis,
 )
 from extras.api.serializers import RenderedGraphSerializer
-from extras.api.views import CustomFieldModelViewSet
+from extras.api.views import ConfigContextQuerySetMixin, CustomFieldModelViewSet
 from extras.models import Graph
 from ipam.models import Prefix, VLAN
 from utilities.api import (
@@ -336,28 +336,13 @@ class PlatformViewSet(ModelViewSet):
 # Devices
 #
 
-class DeviceViewSet(CustomFieldModelViewSet):
+class DeviceViewSet(CustomFieldModelViewSet, ConfigContextQuerySetMixin):
     queryset = Device.objects.prefetch_related(
         'device_type__manufacturer', 'device_role', 'tenant', 'platform', 'site', 'rack', 'parent_bay',
         'virtual_chassis__master', 'primary_ip4__nat_outside', 'primary_ip6__nat_outside', 'tags',
     )
     filterset_class = filters.DeviceFilterSet
 
-    def get_queryset(self):
-        """
-        Build the proper queryset based on the request context
-
-        If the `brief` query param equates to True or the `exclude` query param
-        includes `config_context` as a value, return the base queryset.
-
-        Else, return the queryset annotated with config context data
-        """
-
-        request = self.get_serializer_context()['request']
-        if request.query_params.get('brief') or 'config_context' in request.query_params.get('exclude', []):
-            return self.queryset
-        return self.queryset.annotate_config_context_data()
-
     def get_serializer_class(self):
         """
         Select the specific serializer based on the request context.

netbox/extras/api/views.py

@@ -26,6 +26,29 @@ from utilities.utils import copy_safe_request
 from . import serializers
 
 
+class ConfigContextQuerySetMixin:
+    """
+    Used by views that work with config context models (device and virtual machine).
+    Provides a get_queryset() method which deals with adding the config context
+    data annotation or not.
+    """
+
+    def get_queryset(self):
+        """
+        Build the proper queryset based on the request context
+
+        If the `brief` query param equates to True or the `exclude` query param
+        includes `config_context` as a value, return the base queryset.
+
+        Else, return the queryset annotated with config context data
+        """
+
+        request = self.get_serializer_context()['request']
+        if request.query_params.get('brief') or 'config_context' in request.query_params.get('exclude', []):
+            return self.queryset
+        return self.queryset.annotate_config_context_data()
+
+
 class ExtrasRootView(APIRootView):
     """
     Extras API root view

netbox/virtualization/api/views.py

@@ -6,7 +6,7 @@ from rest_framework.routers import APIRootView
 
 from dcim.models import Device
 from extras.api.serializers import RenderedGraphSerializer
-from extras.api.views import CustomFieldModelViewSet
+from extras.api.views import ConfigContextQuerySetMixin, CustomFieldModelViewSet
 from extras.models import Graph
 from utilities.api import ModelViewSet
 from utilities.utils import get_subquery
@@ -58,27 +58,12 @@ class ClusterViewSet(CustomFieldModelViewSet):
 # Virtual machines
 #
 
-class VirtualMachineViewSet(CustomFieldModelViewSet):
+class VirtualMachineViewSet(CustomFieldModelViewSet, ConfigContextQuerySetMixin):
     queryset = VirtualMachine.objects.prefetch_related(
         'cluster__site', 'role', 'tenant', 'platform', 'primary_ip4', 'primary_ip6', 'tags'
     )
     filterset_class = filters.VirtualMachineFilterSet
 
-    def get_queryset(self):
-        """
-        Build the proper queryset based on the request context
-
-        If the `brief` query param equates to True or the `exclude` query param
-        includes `config_context` as a value, return the base queryset.
-
-        Else, return the queryset annotated with config context data
-        """
-
-        request = self.get_serializer_context()['request']
-        if request.query_params.get('brief') or 'config_context' in request.query_params.get('exclude', []):
-            return self.queryset
-        return self.queryset.annotate_config_context_data()
-
     def get_serializer_class(self):
         """
         Select the specific serializer based on the request context.