Closes #11611: Refactor API viewset classes and introduce NetBoxReadOnlyModelViewSet

jeremystretch · jeremystretch · commit 7accdd52d836 · 2023-02-19T20:09:51.000-05:00
diff --git a/docs/release-notes/version-3.5.md b/docs/release-notes/version-3.5.md
@@ -13,3 +13,4 @@
 
 * [#10604](https://github.com/netbox-community/netbox/issues/10604) - Remove unused `extra_tabs` block from `object.html` generic template
 * [#10923](https://github.com/netbox-community/netbox/issues/10923) - Remove unused `NetBoxModelCSVForm` class (replaced by `NetBoxModelImportForm`)
+* [#11611](https://github.com/netbox-community/netbox/issues/11611) - Refactor API viewset classes and introduce NetBoxReadOnlyModelViewSet
diff --git a/netbox/netbox/api/viewsets/__init__.py b/netbox/netbox/api/viewsets/__init__.py
@@ -1,21 +1,17 @@
 import logging
 
-from django.contrib.contenttypes.models import ContentType
 from django.core.exceptions import ObjectDoesNotExist, PermissionDenied
 from django.db import transaction
 from django.db.models import ProtectedError
-from django.http import Http404
+from rest_framework import mixins as drf_mixins
 from rest_framework.response import Response
-from rest_framework.viewsets import ModelViewSet
+from rest_framework.viewsets import GenericViewSet
 
-from extras.models import ExportTemplate
-from netbox.api.exceptions import SerializerNotFound
-from netbox.constants import NESTED_SERIALIZER_PREFIX
-from utilities.api import get_serializer_for_model
 from utilities.exceptions import AbortRequest
-from .mixins import *
+from . import mixins
 
 __all__ = (
+    'NetBoxReadOnlyModelViewSet',
     'NetBoxModelViewSet',
 )
 
@@ -30,13 +26,47 @@
 }
 
 
-class NetBoxModelViewSet(BulkUpdateModelMixin, BulkDestroyModelMixin, ObjectValidationMixin, ModelViewSet):
+class BaseViewSet(GenericViewSet):
     """
-    Extend DRF's ModelViewSet to support bulk update and delete functions.
+    Base class for all API ViewSets. This is responsible for the enforcement of object-based permissions.
     """
-    brief = False
-    brief_prefetch_fields = []
+    def initial(self, request, *args, **kwargs):
+        super().initial(request, *args, **kwargs)
 
+        # Restrict the view's QuerySet to allow only the permitted objects
+        if request.user.is_authenticated:
+            if action := HTTP_ACTIONS[request.method]:
+                self.queryset = self.queryset.restrict(request.user, action)
+
+
+class NetBoxReadOnlyModelViewSet(
+    mixins.BriefModeMixin,
+    mixins.CustomFieldsMixin,
+    mixins.ExportTemplatesMixin,
+    drf_mixins.RetrieveModelMixin,
+    drf_mixins.ListModelMixin,
+    BaseViewSet
+):
+    pass
+
+
+class NetBoxModelViewSet(
+    mixins.BulkUpdateModelMixin,
+    mixins.BulkDestroyModelMixin,
+    mixins.ObjectValidationMixin,
+    mixins.BriefModeMixin,
+    mixins.CustomFieldsMixin,
+    mixins.ExportTemplatesMixin,
+    drf_mixins.CreateModelMixin,
+    drf_mixins.RetrieveModelMixin,
+    drf_mixins.UpdateModelMixin,
+    drf_mixins.DestroyModelMixin,
+    drf_mixins.ListModelMixin,
+    BaseViewSet
+):
+    """
+    Extend DRF's ModelViewSet to support bulk update and delete functions.
+    """
     def get_object_with_snapshot(self):
         """
         Save a pre-change snapshot of the object immediately after retrieving it. This snapshot will be used to
@@ -48,71 +78,14 @@ def get_object_with_snapshot(self):
         return obj
 
     def get_serializer(self, *args, **kwargs):
-
         # If a list of objects has been provided, initialize the serializer with many=True
         if isinstance(kwargs.get('data', {}), list):
             kwargs['many'] = True
 
         return super().get_serializer(*args, **kwargs)
 
-    def get_serializer_class(self):
-        logger = logging.getLogger('netbox.api.views.ModelViewSet')
-
-        # If using 'brief' mode, find and return the nested serializer for this model, if one exists
-        if self.brief:
-            logger.debug("Request is for 'brief' format; initializing nested serializer")
-            try:
-                serializer = get_serializer_for_model(self.queryset.model, prefix=NESTED_SERIALIZER_PREFIX)
-                logger.debug(f"Using serializer {serializer}")
-                return serializer
-            except SerializerNotFound:
-                logger.debug(f"Nested serializer for {self.queryset.model} not found!")
-
-        # Fall back to the hard-coded serializer class
-        logger.debug(f"Using serializer {self.serializer_class}")
-        return self.serializer_class
-
-    def get_serializer_context(self):
-        """
-        For models which support custom fields, populate the `custom_fields` context.
-        """
-        context = super().get_serializer_context()
-
-        if hasattr(self.queryset.model, 'custom_fields'):
-            content_type = ContentType.objects.get_for_model(self.queryset.model)
-            context.update({
-                'custom_fields': content_type.custom_fields.all(),
-            })
-
-        return context
-
-    def get_queryset(self):
-        # If using brief mode, clear all prefetches from the queryset and append only brief_prefetch_fields (if any)
-        if self.brief:
-            return super().get_queryset().prefetch_related(None).prefetch_related(*self.brief_prefetch_fields)
-
-        return super().get_queryset()
-
-    def initialize_request(self, request, *args, **kwargs):
-        # Check if brief=True has been passed
-        if request.method == 'GET' and request.GET.get('brief'):
-            self.brief = True
-
-        return super().initialize_request(request, *args, **kwargs)
-
-    def initial(self, request, *args, **kwargs):
-        super().initial(request, *args, **kwargs)
-
-        if not request.user.is_authenticated:
-            return
-
-        # Restrict the view's QuerySet to allow only the permitted objects
-        action = HTTP_ACTIONS[request.method]
-        if action:
-            self.queryset = self.queryset.restrict(request.user, action)
-
     def dispatch(self, request, *args, **kwargs):
-        logger = logging.getLogger('netbox.api.views.ModelViewSet')
+        logger = logging.getLogger(f'netbox.api.views.{self.__class__.__name__}')
 
         try:
             return super().dispatch(request, *args, **kwargs)
@@ -136,21 +109,11 @@ def dispatch(self, request, *args, **kwargs):
                 **kwargs
             )
 
-    def list(self, request, *args, **kwargs):
-        # Overrides ListModelMixin to allow processing ExportTemplates.
-        if 'export' in request.GET:
-            content_type = ContentType.objects.get_for_model(self.get_serializer_class().Meta.model)
-            et = ExportTemplate.objects.filter(content_types=content_type, name=request.GET['export']).first()
-            if et is None:
-                raise Http404
-            queryset = self.filter_queryset(self.get_queryset())
-            return et.render_to_response(queryset)
-
-        return super().list(request, *args, **kwargs)
+    # Creates
 
     def perform_create(self, serializer):
         model = self.queryset.model
-        logger = logging.getLogger('netbox.api.views.ModelViewSet')
+        logger = logging.getLogger(f'netbox.api.views.{self.__class__.__name__}')
         logger.info(f"Creating new {model._meta.verbose_name}")
 
         # Enforce object-level permissions on save()
@@ -161,14 +124,16 @@ def perform_create(self, serializer):
         except ObjectDoesNotExist:
             raise PermissionDenied()
 
+    # Updates
+
     def update(self, request, *args, **kwargs):
         # Hotwire get_object() to ensure we save a pre-change snapshot
         self.get_object = self.get_object_with_snapshot
         return super().update(request, *args, **kwargs)
 
     def perform_update(self, serializer):
         model = self.queryset.model
-        logger = logging.getLogger('netbox.api.views.ModelViewSet')
+        logger = logging.getLogger(f'netbox.api.views.{self.__class__.__name__}')
         logger.info(f"Updating {model._meta.verbose_name} {serializer.instance} (PK: {serializer.instance.pk})")
 
         # Enforce object-level permissions on save()
@@ -179,14 +144,16 @@ def perform_update(self, serializer):
         except ObjectDoesNotExist:
             raise PermissionDenied()
 
+    # Deletes
+
     def destroy(self, request, *args, **kwargs):
         # Hotwire get_object() to ensure we save a pre-change snapshot
         self.get_object = self.get_object_with_snapshot
         return super().destroy(request, *args, **kwargs)
 
     def perform_destroy(self, instance):
         model = self.queryset.model
-        logger = logging.getLogger('netbox.api.views.ModelViewSet')
+        logger = logging.getLogger(f'netbox.api.views.{self.__class__.__name__}')
         logger.info(f"Deleting {model._meta.verbose_name} {instance} (PK: {instance.pk})")
 
         return super().perform_destroy(instance)
diff --git a/netbox/netbox/api/viewsets/mixins.py b/netbox/netbox/api/viewsets/mixins.py
@@ -1,17 +1,99 @@
+import logging
+
+from django.contrib.contenttypes.models import ContentType
 from django.core.exceptions import ObjectDoesNotExist
 from django.db import transaction
+from django.http import Http404
 from rest_framework import status
 from rest_framework.response import Response
 
+from extras.models import ExportTemplate
+from netbox.api.exceptions import SerializerNotFound
 from netbox.api.serializers import BulkOperationSerializer
+from netbox.constants import NESTED_SERIALIZER_PREFIX
+from utilities.api import get_serializer_for_model
 
 __all__ = (
+    'BriefModeMixin',
     'BulkUpdateModelMixin',
+    'CustomFieldsMixin',
+    'ExportTemplatesMixin',
     'BulkDestroyModelMixin',
     'ObjectValidationMixin',
 )
 
 
+class BriefModeMixin:
+    """
+    Enables brief mode support, so that the client can invoke a model's nested serializer by passing e.g.
+        GET /api/dcim/sites/?brief=True
+    """
+    brief = False
+    brief_prefetch_fields = []
+
+    def initialize_request(self, request, *args, **kwargs):
+        # Annotate whether brief mode is active
+        self.brief = request.method == 'GET' and request.GET.get('brief')
+
+        return super().initialize_request(request, *args, **kwargs)
+
+    def get_serializer_class(self):
+        logger = logging.getLogger(f'netbox.api.views.{self.__class__.__name__}')
+
+        # If using 'brief' mode, find and return the nested serializer for this model, if one exists
+        if self.brief:
+            logger.debug("Request is for 'brief' format; initializing nested serializer")
+            try:
+                return get_serializer_for_model(self.queryset.model, prefix=NESTED_SERIALIZER_PREFIX)
+            except SerializerNotFound:
+                logger.debug(
+                    f"Nested serializer for {self.queryset.model} not found! Using serializer {self.serializer_class}"
+                )
+
+        return self.serializer_class
+
+    def get_queryset(self):
+        qs = super().get_queryset()
+
+        # If using brief mode, clear all prefetches from the queryset and append only brief_prefetch_fields (if any)
+        if self.brief:
+            return qs.prefetch_related(None).prefetch_related(*self.brief_prefetch_fields)
+
+        return qs
+
+
+class CustomFieldsMixin:
+    """
+    For models which support custom fields, populate the `custom_fields` context.
+    """
+    def get_serializer_context(self):
+        context = super().get_serializer_context()
+
+        if hasattr(self.queryset.model, 'custom_fields'):
+            content_type = ContentType.objects.get_for_model(self.queryset.model)
+            context.update({
+                'custom_fields': content_type.custom_fields.all(),
+            })
+
+        return context
+
+
+class ExportTemplatesMixin:
+    """
+    Enable ExportTemplate support for list views.
+    """
+    def list(self, request, *args, **kwargs):
+        if 'export' in request.GET:
+            content_type = ContentType.objects.get_for_model(self.get_serializer_class().Meta.model)
+            et = ExportTemplate.objects.filter(content_types=content_type, name=request.GET['export']).first()
+            if et is None:
+                raise Http404
+            queryset = self.filter_queryset(self.get_queryset())
+            return et.render_to_response(queryset)
+
+        return super().list(request, *args, **kwargs)
+
+
 class BulkUpdateModelMixin:
     """
     Support bulk modification of objects using the list endpoint for a model. Accepts a PATCH action with a list of one

Original file line number	Diff line number	Diff line change
`@@ -13,3 +13,4 @@`
`13`	`13`
`14`	`14`	* [#10604](https://github.com/netbox-community/netbox/issues/10604) - Remove unused `extra_tabs` block from `object.html` generic template
`15`	`15`	* [#10923](https://github.com/netbox-community/netbox/issues/10923) - Remove unused `NetBoxModelCSVForm` class (replaced by `NetBoxModelImportForm`)
	`16`	`+* [#11611](https://github.com/netbox-community/netbox/issues/11611) - Refactor API viewset classes and introduce NetBoxReadOnlyModelViewSet`