diff --git a/netbox/dcim/api/views.py b/netbox/dcim/api/views.py index 956b0cbff..427aecd5f 100644 --- a/netbox/dcim/api/views.py +++ b/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. diff --git a/netbox/extras/api/views.py b/netbox/extras/api/views.py index a63dbe44d..74c9ea889 100644 --- a/netbox/extras/api/views.py +++ b/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 diff --git a/netbox/virtualization/api/views.py b/netbox/virtualization/api/views.py index d19f0f9fa..55393b110 100644 --- a/netbox/virtualization/api/views.py +++ b/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.