12336 add locks to all views using mptt models

netbox/dcim/api/views.py

@@ -1,6 +1,5 @@
 from django.http import Http404, HttpResponse
 from django.shortcuts import get_object_or_404
-from django_pglocks import advisory_lock
 from drf_spectacular.types import OpenApiTypes
 from drf_spectacular.utils import extend_schema, OpenApiParameter
 from rest_framework.decorators import action
@@ -21,9 +20,9 @@ from netbox.api.authentication import IsAuthenticatedOrLoginNotRequired
 from netbox.api.metadata import ContentTypeMetadata
 from netbox.api.pagination import StripCountAnnotationsPaginator
 from netbox.api.renderers import TextRenderer
-from netbox.api.viewsets import NetBoxModelViewSet
+from netbox.api.viewsets import NetBoxModelViewSet, MPTTLockedMixin
 from netbox.api.viewsets.mixins import SequentialBulkCreatesMixin
-from netbox.constants import NESTED_SERIALIZER_PREFIX, ADVISORY_LOCK_KEYS
+from netbox.constants import NESTED_SERIALIZER_PREFIX
 from utilities.api import get_serializer_for_model
 from utilities.utils import count_related
 from virtualization.models import VirtualMachine
@@ -99,7 +98,7 @@ class PassThroughPortMixin(object):
 # Regions
 #
 
-class RegionViewSet(NetBoxModelViewSet):
+class RegionViewSet(MPTTLockedMixin, NetBoxModelViewSet):
     queryset = Region.objects.add_related_count(
         Region.objects.all(),
         Site,
@@ -110,24 +109,12 @@ class RegionViewSet(NetBoxModelViewSet):
     serializer_class = serializers.RegionSerializer
     filterset_class = filtersets.RegionFilterSet
 
-    @advisory_lock(ADVISORY_LOCK_KEYS['regions'])
-    def create(self, request, *args, **kwargs):
-        return super().create(request, *args, **kwargs)
-
-    @advisory_lock(ADVISORY_LOCK_KEYS['regions'])
-    def update(self, request, *args, **kwargs):
-        return super().update(request, *args, **kwargs)
-
-    @advisory_lock(ADVISORY_LOCK_KEYS['regions'])
-    def destroy(self, request, *args, **kwargs):
-        return super().destroy(request, *args, **kwargs)
-
 
 #
 # Site groups
 #
 
-class SiteGroupViewSet(NetBoxModelViewSet):
+class SiteGroupViewSet(MPTTLockedMixin, NetBoxModelViewSet):
     queryset = SiteGroup.objects.add_related_count(
         SiteGroup.objects.all(),
         Site,
@@ -162,7 +149,7 @@ class SiteViewSet(NetBoxModelViewSet):
 # Locations
 #
 
-class LocationViewSet(NetBoxModelViewSet):
+class LocationViewSet(MPTTLockedMixin, NetBoxModelViewSet):
     queryset = Location.objects.add_related_count(
         Location.objects.add_related_count(
             Location.objects.all(),
@@ -363,7 +350,7 @@ class DeviceBayTemplateViewSet(NetBoxModelViewSet):
     filterset_class = filtersets.DeviceBayTemplateFilterSet
 
 
-class InventoryItemTemplateViewSet(NetBoxModelViewSet):
+class InventoryItemTemplateViewSet(MPTTLockedMixin, NetBoxModelViewSet):
     queryset = InventoryItemTemplate.objects.prefetch_related('device_type__manufacturer', 'role')
     serializer_class = serializers.InventoryItemTemplateSerializer
     filterset_class = filtersets.InventoryItemTemplateFilterSet
@@ -551,7 +538,7 @@ class DeviceBayViewSet(NetBoxModelViewSet):
     brief_prefetch_fields = ['device']
 
 
-class InventoryItemViewSet(NetBoxModelViewSet):
+class InventoryItemViewSet(MPTTLockedMixin, NetBoxModelViewSet):
     queryset = InventoryItem.objects.prefetch_related('device', 'manufacturer', 'tags')
     serializer_class = serializers.InventoryItemSerializer
     filterset_class = filtersets.InventoryItemFilterSet

netbox/netbox/api/viewsets/__init__.py

@@ -3,6 +3,8 @@ import logging
 from django.core.exceptions import ObjectDoesNotExist, PermissionDenied
 from django.db import transaction
 from django.db.models import ProtectedError
+from django_pglocks import advisory_lock
+from netbox.constants import ADVISORY_LOCK_KEYS
 from rest_framework import mixins as drf_mixins
 from rest_framework.response import Response
 from rest_framework.viewsets import GenericViewSet
@@ -157,3 +159,22 @@ class NetBoxModelViewSet(
         logger.info(f"Deleting {model._meta.verbose_name} {instance} (PK: {instance.pk})")
 
         return super().perform_destroy(instance)
+
+
+class MPTTLockedMixin(GenericViewSet):
+    """
+    Puts pglock on objects that derive from MPTTModel for parallel API calling.
+    Note: If adding this to a view, must add the model name to ADVISORY_LOCK_KEYS
+    """
+
+    def create(self, request, *args, **kwargs):
+        with advisory_lock(ADVISORY_LOCK_KEYS[self.queryset.model._meta.verbose_name.lower()]):
+            return super().create(request, *args, **kwargs)
+
+    def update(self, request, *args, **kwargs):
+        with advisory_lock(ADVISORY_LOCK_KEYS[self.queryset.model._meta.verbose_name.lower()]):
+            return super().update(request, *args, **kwargs)
+
+    def destroy(self, request, *args, **kwargs):
+        with advisory_lock(ADVISORY_LOCK_KEYS[self.queryset.model._meta.verbose_name.lower()]):
+            return super().destroy(request, *args, **kwargs)

netbox/netbox/constants.py

@@ -15,5 +15,14 @@ ADVISORY_LOCK_KEYS = {
     'available-ips': 100200,
     'available-vlans': 100300,
     'available-asns': 100400,
-    'regions': 100500,
+
+    # MPTT locks
+    'region': 100500,
+    'sitegroup': 100501,
+    'location': 100502,
+    'tenantgroup': 100503,
+    'contactgroup': 100504,
+    'wirelesslangroup': 100505,
+    'inventoryitem': 100506,
+    'inventoryitemtemplate': 100507,
 }

netbox/tenancy/api/views.py

@@ -3,7 +3,7 @@ from rest_framework.routers import APIRootView
 from circuits.models import Circuit
 from dcim.models import Device, Rack, Site
 from ipam.models import IPAddress, Prefix, VLAN, VRF
-from netbox.api.viewsets import NetBoxModelViewSet
+from netbox.api.viewsets import NetBoxModelViewSet, MPTTLockedMixin
 from tenancy import filtersets
 from tenancy.models import *
 from utilities.utils import count_related
@@ -23,7 +23,7 @@ class TenancyRootView(APIRootView):
 # Tenants
 #
 
-class TenantGroupViewSet(NetBoxModelViewSet):
+class TenantGroupViewSet(MPTTLockedMixin, NetBoxModelViewSet):
     queryset = TenantGroup.objects.add_related_count(
         TenantGroup.objects.all(),
         Tenant,
@@ -58,7 +58,7 @@ class TenantViewSet(NetBoxModelViewSet):
 # Contacts
 #
 
-class ContactGroupViewSet(NetBoxModelViewSet):
+class ContactGroupViewSet(MPTTLockedMixin, NetBoxModelViewSet):
     queryset = ContactGroup.objects.add_related_count(
         ContactGroup.objects.all(),
         Contact,

netbox/wireless/api/views.py

@@ -1,6 +1,6 @@
 from rest_framework.routers import APIRootView
 
-from netbox.api.viewsets import NetBoxModelViewSet
+from netbox.api.viewsets import NetBoxModelViewSet, MPTTLockedMixin
 from wireless import filtersets
 from wireless.models import *
 from . import serializers
@@ -14,7 +14,7 @@ class WirelessRootView(APIRootView):
         return 'Wireless'
 
 
-class WirelessLANGroupViewSet(NetBoxModelViewSet):
+class WirelessLANGroupViewSet(MPTTLockedMixin, NetBoxModelViewSet):
     queryset = WirelessLANGroup.objects.add_related_count(
         WirelessLANGroup.objects.all(),
         WirelessLAN,

netbox/wireless/models.py

@@ -2,7 +2,6 @@ from django.core.exceptions import ValidationError
 from django.db import models
 from django.urls import reverse
 from django.utils.translation import gettext_lazy as _
-from mptt.models import MPTTModel
 
 from dcim.choices import LinkStatusChoices
 from dcim.constants import WIRELESS_IFACE_TYPES