7699 use mixin for model

netbox-community · Oct 24, 2024 · cfdab0e · cfdab0e
1 parent c73902c
commit cfdab0e
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 61 deletions.
diff --git a/netbox/netbox/models/features.py b/netbox/netbox/models/features.py
@@ -2,6 +2,7 @@
 from collections import defaultdict
 from functools import cached_property
 
+from django.apps import apps
 from django.contrib.contenttypes.fields import GenericRelation
 from django.core.validators import ValidationError
 from django.db import models
@@ -23,6 +24,7 @@
 
 __all__ = (
     'BookmarksMixin',
+    'CachedLocationScopeMixin',
     'ChangeLoggingMixin',
     'CloningMixin',
     'ContactsMixin',
@@ -580,6 +582,68 @@ def sync_data(self):
         ))
 
 
+class CachedLocationScopeMixin(models.Model):
+    """
+    Cached associations for scope to enable efficient filtering - must define scope and scope_type on model
+    """
+    _location = models.ForeignKey(
+        to='dcim.Location',
+        on_delete=models.CASCADE,
+        related_name='_%(class)ss',
+        blank=True,
+        null=True
+    )
+    _site = models.ForeignKey(
+        to='dcim.Site',
+        on_delete=models.CASCADE,
+        related_name='_%(class)ss',
+        blank=True,
+        null=True
+    )
+    _region = models.ForeignKey(
+        to='dcim.Region',
+        on_delete=models.CASCADE,
+        related_name='_%(class)ss',
+        blank=True,
+        null=True
+    )
+    _sitegroup = models.ForeignKey(
+        to='dcim.SiteGroup',
+        on_delete=models.CASCADE,
+        related_name='_%(class)ss',
+        blank=True,
+        null=True
+    )
+
+    class Meta:
+        abstract = True
+
+    def save(self, *args, **kwargs):
+        # Cache objects associated with the terminating object (for filtering)
+        self.cache_related_objects()
+
+        super().save(*args, **kwargs)
+
+    def cache_related_objects(self):
+        self._region = self._sitegroup = self._site = self._location = None
+        if self.scope_type:
+            scope_type = self.scope_type.model_class()
+            if scope_type == apps.get_model('dcim', 'region'):
+                self._region = self.scope
+            elif scope_type == apps.get_model('dcim', 'sitegroup'):
+                self._sitegroup = self.scope
+            elif scope_type == apps.get_model('dcim', 'site'):
+                self._region = self.scope.region
+                self._sitegroup = self.scope.group
+                self._site = self.scope
+            elif scope_type == apps.get_model('dcim', 'location'):
+                self._region = self.scope.site.region
+                self._sitegroup = self.scope.site.group
+                self._site = self.scope.site
+                self._location = self.scope
+    cache_related_objects.alters_data = True
+
+
 #
 # Feature registration
 #

diff --git a/netbox/virtualization/migrations/0043_clusters_cached_relations.py b/netbox/virtualization/migrations/0043_clusters_cached_relations.py
@@ -32,7 +32,7 @@ class Migration(migrations.Migration):
                 blank=True,
                 null=True,
                 on_delete=django.db.models.deletion.CASCADE,
-                related_name='_clusters',
+                related_name='_%(class)ss',
                 to='dcim.location',
             ),
         ),
@@ -43,7 +43,7 @@ class Migration(migrations.Migration):
                 blank=True,
                 null=True,
                 on_delete=django.db.models.deletion.CASCADE,
-                related_name='_clusters',
+                related_name='_%(class)ss',
                 to='dcim.region',
             ),
         ),
@@ -54,7 +54,7 @@ class Migration(migrations.Migration):
                 blank=True,
                 null=True,
                 on_delete=django.db.models.deletion.CASCADE,
-                related_name='_clusters',
+                related_name='_%(class)ss',
                 to='dcim.site',
             ),
         ),
@@ -65,7 +65,7 @@ class Migration(migrations.Migration):
                 blank=True,
                 null=True,
                 on_delete=django.db.models.deletion.CASCADE,
-                related_name='_clusters',
+                related_name='_%(class)ss',
                 to='dcim.sitegroup',
             ),
         ),

diff --git a/netbox/virtualization/models/clusters.py b/netbox/virtualization/models/clusters.py
@@ -6,7 +6,7 @@
 
 from dcim.models import Device
 from netbox.models import OrganizationalModel, PrimaryModel
-from netbox.models.features import ContactsMixin
+from netbox.models.features import CachedLocationScopeMixin, ContactsMixin
 from virtualization.choices import *
 from virtualization.constants import CLUSTER_SCOPE_TYPES
 
@@ -44,7 +44,7 @@ class Meta:
         verbose_name_plural = _('cluster groups')
 
 
-class Cluster(ContactsMixin, PrimaryModel):
+class Cluster(ContactsMixin, CachedLocationScopeMixin, PrimaryModel):
     """
     A cluster of VirtualMachines. Each Cluster may optionally be associated with one or more Devices.
     """
@@ -103,36 +103,6 @@ class Cluster(ContactsMixin, PrimaryModel):
         related_query_name='cluster'
     )
 
-    # Cached associations to enable efficient filtering
-    _location = models.ForeignKey(
-        to='dcim.Location',
-        on_delete=models.CASCADE,
-        related_name='_clusters',
-        blank=True,
-        null=True
-    )
-    _site = models.ForeignKey(
-        to='dcim.Site',
-        on_delete=models.CASCADE,
-        related_name='_clusters',
-        blank=True,
-        null=True
-    )
-    _region = models.ForeignKey(
-        to='dcim.Region',
-        on_delete=models.CASCADE,
-        related_name='_clusters',
-        blank=True,
-        null=True
-    )
-    _sitegroup = models.ForeignKey(
-        to='dcim.SiteGroup',
-        on_delete=models.CASCADE,
-        related_name='_clusters',
-        blank=True,
-        null=True
-    )
-
     clone_fields = (
         'scope_type', 'scope_id', 'type', 'group', 'status', 'tenant',
     )
@@ -180,28 +150,3 @@ def clean(self):
                         "{count} devices are assigned as hosts for this cluster but are not in site {site}"
                     ).format(count=nonsite_devices, site=site)
                 })
-
-    def save(self, *args, **kwargs):
-        # Cache objects associated with the terminating object (for filtering)
-        self.cache_related_objects()
-
-        super().save(*args, **kwargs)
-
-    def cache_related_objects(self):
-        self._region = self._sitegroup = self._site = self._location = None
-        if self.scope_type:
-            scope_type = self.scope_type.model_class()
-            if scope_type == apps.get_model('dcim', 'region'):
-                self._region = self.scope
-            elif scope_type == apps.get_model('dcim', 'sitegroup'):
-                self._sitegroup = self.scope
-            elif scope_type == apps.get_model('dcim', 'site'):
-                self._region = self.scope.region
-                self._sitegroup = self.scope.group
-                self._site = self.scope
-            elif scope_type == apps.get_model('dcim', 'location'):
-                self._region = self.scope.site.region
-                self._sitegroup = self.scope.site.group
-                self._site = self.scope.site
-                self._location = self.scope
-    cache_related_objects.alters_data = True