Designate feature mixin classes & employ class_prepared signal to register features

netbox/extras/utils.py

@@ -57,21 +57,24 @@ class FeatureQuery:
         return query
 
 
+def register_features(model, features):
+    if 'model_features' not in registry:
+        registry['model_features'] = {
+            f: collections.defaultdict(list) for f in EXTRAS_FEATURES
+        }
+    for feature in features:
+        if feature not in EXTRAS_FEATURES:
+            raise ValueError(f"{feature} is not a valid extras feature!")
+        app_label, model_name = model._meta.label_lower.split('.')
+        registry['model_features'][feature][app_label].append(model_name)
+
+
 def extras_features(*features):
     """
     Decorator used to register extras provided features to a model
     """
     def wrapper(model_class):
         # Initialize the model_features store if not already defined
-        if 'model_features' not in registry:
-            registry['model_features'] = {
-                f: collections.defaultdict(list) for f in EXTRAS_FEATURES
-            }
-        for feature in features:
-            if feature in EXTRAS_FEATURES:
-                app_label, model_name = model_class._meta.label_lower.split('.')
-                registry['model_features'][feature][app_label].append(model_name)
-            else:
-                raise ValueError('{} is not a valid extras feature!'.format(feature))
+        register_features(model_class, features)
         return model_class
     return wrapper

netbox/netbox/models/__init__.py

@@ -0,0 +1,129 @@
+from django.contrib.contenttypes.fields import GenericRelation
+from django.core.validators import ValidationError
+from django.db import models
+from mptt.models import MPTTModel, TreeForeignKey
+
+from utilities.mptt import TreeManager
+from utilities.querysets import RestrictedQuerySet
+from netbox.models.features import *
+
+__all__ = (
+    'BigIDModel',
+    'ChangeLoggedModel',
+    'NestedGroupModel',
+    'OrganizationalModel',
+    'PrimaryModel',
+)
+
+
+#
+# Base model classes
+#
+
+class BigIDModel(models.Model):
+    """
+    Abstract base model for all data objects. Ensures the use of a 64-bit PK.
+    """
+    id = models.BigAutoField(
+        primary_key=True
+    )
+
+    class Meta:
+        abstract = True
+
+
+class ChangeLoggedModel(ChangeLoggingMixin, CustomValidationMixin, BigIDModel):
+    """
+    Base model for all objects which support change logging.
+    """
+    objects = RestrictedQuerySet.as_manager()
+
+    class Meta:
+        abstract = True
+
+
+class PrimaryModel(ChangeLoggingMixin, CustomFieldsMixin, CustomValidationMixin, TagsMixin, BigIDModel):
+    """
+    Primary models represent real objects within the infrastructure being modeled.
+    """
+    journal_entries = GenericRelation(
+        to='extras.JournalEntry',
+        object_id_field='assigned_object_id',
+        content_type_field='assigned_object_type'
+    )
+
+    objects = RestrictedQuerySet.as_manager()
+
+    class Meta:
+        abstract = True
+
+
+class NestedGroupModel(ChangeLoggingMixin, CustomFieldsMixin, CustomValidationMixin, TagsMixin, BigIDModel, MPTTModel):
+    """
+    Base model for objects which are used to form a hierarchy (regions, locations, etc.). These models nest
+    recursively using MPTT. Within each parent, each child instance must have a unique name.
+    """
+    parent = TreeForeignKey(
+        to='self',
+        on_delete=models.CASCADE,
+        related_name='children',
+        blank=True,
+        null=True,
+        db_index=True
+    )
+    name = models.CharField(
+        max_length=100
+    )
+    description = models.CharField(
+        max_length=200,
+        blank=True
+    )
+
+    objects = TreeManager()
+
+    class Meta:
+        abstract = True
+
+    class MPTTMeta:
+        order_insertion_by = ('name',)
+
+    def __str__(self):
+        return self.name
+
+    def clean(self):
+        super().clean()
+
+        # An MPTT model cannot be its own parent
+        if self.pk and self.parent_id == self.pk:
+            raise ValidationError({
+                "parent": "Cannot assign self as parent."
+            })
+
+
+class OrganizationalModel(ChangeLoggingMixin, CustomFieldsMixin, CustomValidationMixin, TagsMixin, BigIDModel):
+    """
+    Organizational models are those which are used solely to categorize and qualify other objects, and do not convey
+    any real information about the infrastructure being modeled (for example, functional device roles). Organizational
+    models provide the following standard attributes:
+    - Unique name
+    - Unique slug (automatically derived from name)
+    - Optional description
+    """
+    name = models.CharField(
+        max_length=100,
+        unique=True
+    )
+    slug = models.SlugField(
+        max_length=100,
+        unique=True
+    )
+    description = models.CharField(
+        max_length=200,
+        blank=True
+    )
+
+    objects = RestrictedQuerySet.as_manager()
+
+    class Meta:
+        abstract = True
+        ordering = ('name',)

netbox/netbox/models.py → netbox/netbox/models/features.py

@@ -1,34 +1,37 @@
 import logging
 
-from django.contrib.contenttypes.fields import GenericRelation
+from django.db.models.signals import class_prepared
+from django.dispatch import receiver
+
 from django.core.serializers.json import DjangoJSONEncoder
 from django.core.validators import ValidationError
 from django.db import models
-from mptt.models import MPTTModel, TreeForeignKey
 from taggit.managers import TaggableManager
 
 from extras.choices import ObjectChangeActionChoices
+from extras.utils import register_features
 from netbox.signals import post_clean
-from utilities.mptt import TreeManager
-from utilities.querysets import RestrictedQuerySet
 from utilities.utils import serialize_object
 
 __all__ = (
-    'BigIDModel',
-    'ChangeLoggedModel',
-    'NestedGroupModel',
-    'OrganizationalModel',
-    'PrimaryModel',
+    'ChangeLoggingMixin',
+    'CustomFieldsMixin',
+    'CustomLinksMixin',
+    'CustomValidationMixin',
+    'ExportTemplatesMixin',
+    'JobResultsMixin',
+    'TagsMixin',
+    'WebhooksMixin',
 )
 
 
 #
-# Mixins
+# Feature mixins
 #
 
 class ChangeLoggingMixin(models.Model):
     """
-    Provides change logging support.
+    Provides change logging support for a model. Adds the `created` and `last_updated` fields.
     """
     created = models.DateField(
         auto_now_add=True,
@@ -74,7 +77,7 @@ class ChangeLoggingMixin(models.Model):
 
 class CustomFieldsMixin(models.Model):
     """
-    Provides support for custom fields.
+    Enables support for custom fields.
     """
     custom_field_data = models.JSONField(
         encoder=DjangoJSONEncoder,
@@ -128,6 +131,14 @@ class CustomFieldsMixin(models.Model):
                 raise ValidationError(f"Missing required custom field '{cf.name}'.")
 
 
+class CustomLinksMixin(models.Model):
+    """
+    Enables support for custom links.
+    """
+    class Meta:
+        abstract = True
+
+
 class CustomValidationMixin(models.Model):
     """
     Enables user-configured validation rules for built-in models by extending the clean() method.
@@ -142,125 +153,55 @@ class CustomValidationMixin(models.Model):
         post_clean.send(sender=self.__class__, instance=self)
 
 
-class TagsMixin(models.Model):
-    """
-    Enable the assignment of Tags.
-    """
-    tags = TaggableManager(
-        through='extras.TaggedItem'
-    )
-
-    class Meta:
-        abstract = True
-
-
-#
-# Base model classes
-
-class BigIDModel(models.Model):
+class ExportTemplatesMixin(models.Model):
     """
-    Abstract base model for all data objects. Ensures the use of a 64-bit PK.
+    Enables support for export templates.
     """
-    id = models.BigAutoField(
-        primary_key=True
-    )
-
     class Meta:
         abstract = True
 
 
-class ChangeLoggedModel(ChangeLoggingMixin, CustomValidationMixin, BigIDModel):
+class JobResultsMixin(models.Model):
     """
-    Base model for all objects which support change logging.
+    Enable the assignment of JobResults to a model.
     """
-    objects = RestrictedQuerySet.as_manager()
-
     class Meta:
         abstract = True
 
 
-class PrimaryModel(ChangeLoggingMixin, CustomFieldsMixin, CustomValidationMixin, TagsMixin, BigIDModel):
+class TagsMixin(models.Model):
     """
-    Primary models represent real objects within the infrastructure being modeled.
+    Enable the assignment of Tags to a model.
     """
-    journal_entries = GenericRelation(
-        to='extras.JournalEntry',
-        object_id_field='assigned_object_id',
-        content_type_field='assigned_object_type'
+    tags = TaggableManager(
+        through='extras.TaggedItem'
     )
 
-    objects = RestrictedQuerySet.as_manager()
-
     class Meta:
         abstract = True
 
 
-class NestedGroupModel(ChangeLoggingMixin, CustomFieldsMixin, CustomValidationMixin, TagsMixin, BigIDModel, MPTTModel):
+class WebhooksMixin(models.Model):
     """
-    Base model for objects which are used to form a hierarchy (regions, locations, etc.). These models nest
-    recursively using MPTT. Within each parent, each child instance must have a unique name.
+    Enables support for webhooks.
     """
-    parent = TreeForeignKey(
-        to='self',
-        on_delete=models.CASCADE,
-        related_name='children',
-        blank=True,
-        null=True,
-        db_index=True
-    )
-    name = models.CharField(
-        max_length=100
-    )
-    description = models.CharField(
-        max_length=200,
-        blank=True
-    )
-
-    objects = TreeManager()
-
     class Meta:
         abstract = True
 
-    class MPTTMeta:
-        order_insertion_by = ('name',)
 
-    def __str__(self):
-        return self.name
-
-    def clean(self):
-        super().clean()
-
-        # An MPTT model cannot be its own parent
-        if self.pk and self.parent_id == self.pk:
-            raise ValidationError({
-                "parent": "Cannot assign self as parent."
-            })
-
-
-class OrganizationalModel(ChangeLoggingMixin, CustomFieldsMixin, CustomValidationMixin, TagsMixin, BigIDModel):
-    """
-    Organizational models are those which are used solely to categorize and qualify other objects, and do not convey
-    any real information about the infrastructure being modeled (for example, functional device roles). Organizational
-    models provide the following standard attributes:
-    - Unique name
-    - Unique slug (automatically derived from name)
-    - Optional description
-    """
-    name = models.CharField(
-        max_length=100,
-        unique=True
-    )
-    slug = models.SlugField(
-        max_length=100,
-        unique=True
-    )
-    description = models.CharField(
-        max_length=200,
-        blank=True
-    )
+FEATURES_MAP = (
+    ('custom_fields', CustomFieldsMixin),
+    ('custom_links', CustomLinksMixin),
+    ('export_templates', ExportTemplatesMixin),
+    ('job_results', JobResultsMixin),
+    ('tags', TagsMixin),
+    ('webhooks', WebhooksMixin),
+)
 
-    objects = RestrictedQuerySet.as_manager()
 
-    class Meta:
-        abstract = True
-        ordering = ('name',)
+@receiver(class_prepared)
+def _register_features(sender, **kwargs):
+    features = {
+        feature for feature, cls in FEATURES_MAP if issubclass(sender, cls)
+    }
+    register_features(sender, features)