implemented registry for extras model functionality

netbox/dcim/models/__init__.py

@@ -21,6 +21,7 @@ from dcim.constants import *
 from dcim.fields import ASNField
 from dcim.elevations import RackElevationSVG
 from extras.models import ConfigContextModel, CustomFieldModel, ObjectChange, TaggedItem
+from extras.utils import extras_functionality
 from utilities.fields import ColorField, NaturalOrderingField
 from utilities.models import ChangeLoggedModel
 from utilities.utils import serialize_object, to_meters
@@ -1220,6 +1221,7 @@ class Platform(ChangeLoggedModel):
         )
 
 
+@extras_functionality(['webhooks', 'custom_fields', 'export_templates', 'custom_links', 'graphs'])
 class Device(ChangeLoggedModel, ConfigContextModel, CustomFieldModel):
     """
     A Device represents a piece of physical hardware mounted within a Rack. Each Device is assigned a DeviceType,

netbox/extras/api/serializers.py

@@ -1,5 +1,6 @@
 from django.contrib.contenttypes.models import ContentType
 from django.core.exceptions import ObjectDoesNotExist
+from django.db.models import Q
 from drf_yasg.utils import swagger_serializer_method
 from rest_framework import serializers
 
@@ -13,6 +14,7 @@ from extras.constants import *
 from extras.models import (
     ConfigContext, ExportTemplate, Graph, ImageAttachment, ObjectChange, ReportResult, Tag,
 )
+from extras.utils import FunctionalityQueryset
 from tenancy.api.nested_serializers import NestedTenantSerializer, NestedTenantGroupSerializer
 from tenancy.models import Tenant, TenantGroup
 from users.api.nested_serializers import NestedUserSerializer
@@ -31,7 +33,7 @@ from .nested_serializers import *
 
 class GraphSerializer(ValidatedModelSerializer):
     type = ContentTypeField(
-        queryset=ContentType.objects.filter(GRAPH_MODELS),
+        queryset=ContentType.objects.filter(FunctionalityQueryset('graphs').get_queryset()),
     )
 
     class Meta:
@@ -67,7 +69,7 @@ class RenderedGraphSerializer(serializers.ModelSerializer):
 
 class ExportTemplateSerializer(ValidatedModelSerializer):
     content_type = ContentTypeField(
-        queryset=ContentType.objects.filter(EXPORTTEMPLATE_MODELS),
+        queryset=ContentType.objects.filter(Q(FunctionalityQueryset('export_templates').get_queryset())),
     )
     template_language = ChoiceField(
         choices=TemplateLanguageChoices,

netbox/extras/constants.py

@@ -2,127 +2,127 @@ from django.db.models import Q
 
 
 # Models which support custom fields
-CUSTOMFIELD_MODELS = Q(
-    Q(app_label='circuits', model__in=[
-        'circuit',
-        'provider',
-    ]) |
-    Q(app_label='dcim', model__in=[
-        'device',
-        'devicetype',
-        'powerfeed',
-        'rack',
-        'site',
-    ]) |
-    Q(app_label='ipam', model__in=[
-        'aggregate',
-        'ipaddress',
-        'prefix',
-        'service',
-        'vlan',
-        'vrf',
-    ]) |
-    Q(app_label='secrets', model__in=[
-        'secret',
-    ]) |
-    Q(app_label='tenancy', model__in=[
-        'tenant',
-    ]) |
-    Q(app_label='virtualization', model__in=[
-        'cluster',
-        'virtualmachine',
-    ])
-)
-
-# Custom links
-CUSTOMLINK_MODELS = Q(
-    Q(app_label='circuits', model__in=[
-        'circuit',
-        'provider',
-    ]) |
-    Q(app_label='dcim', model__in=[
-        'cable',
-        'device',
-        'devicetype',
-        'powerpanel',
-        'powerfeed',
-        'rack',
-        'site',
-    ]) |
-    Q(app_label='ipam', model__in=[
-        'aggregate',
-        'ipaddress',
-        'prefix',
-        'service',
-        'vlan',
-        'vrf',
-    ]) |
-    Q(app_label='secrets', model__in=[
-        'secret',
-    ]) |
-    Q(app_label='tenancy', model__in=[
-        'tenant',
-    ]) |
-    Q(app_label='virtualization', model__in=[
-        'cluster',
-        'virtualmachine',
-    ])
-)
-
-# Models which can have Graphs associated with them
-GRAPH_MODELS = Q(
-    Q(app_label='circuits', model__in=[
-        'provider',
-    ]) |
-    Q(app_label='dcim', model__in=[
-        'device',
-        'interface',
-        'site',
-    ])
-)
-
-# Models which support export templates
-EXPORTTEMPLATE_MODELS = Q(
-    Q(app_label='circuits', model__in=[
-        'circuit',
-        'provider',
-    ]) |
-    Q(app_label='dcim', model__in=[
-        'cable',
-        'consoleport',
-        'device',
-        'devicetype',
-        'interface',
-        'inventoryitem',
-        'manufacturer',
-        'powerpanel',
-        'powerport',
-        'powerfeed',
-        'rack',
-        'rackgroup',
-        'region',
-        'site',
-        'virtualchassis',
-    ]) |
-    Q(app_label='ipam', model__in=[
-        'aggregate',
-        'ipaddress',
-        'prefix',
-        'service',
-        'vlan',
-        'vrf',
-    ]) |
-    Q(app_label='secrets', model__in=[
-        'secret',
-    ]) |
-    Q(app_label='tenancy', model__in=[
-        'tenant',
-    ]) |
-    Q(app_label='virtualization', model__in=[
-        'cluster',
-        'virtualmachine',
-    ])
-)
+#CUSTOMFIELD_MODELS = Q(
+#    Q(app_label='circuits', model__in=[
+#        'circuit',
+#        'provider',
+#    ]) |
+#    Q(app_label='dcim', model__in=[
+#        'device',
+#        'devicetype',
+#        'powerfeed',
+#        'rack',
+#        'site',
+#    ]) |
+#    Q(app_label='ipam', model__in=[
+#        'aggregate',
+#        'ipaddress',
+#        'prefix',
+#        'service',
+#        'vlan',
+#        'vrf',
+#    ]) |
+#    Q(app_label='secrets', model__in=[
+#        'secret',
+#    ]) |
+#    Q(app_label='tenancy', model__in=[
+#        'tenant',
+#    ]) |
+#    Q(app_label='virtualization', model__in=[
+#        'cluster',
+#        'virtualmachine',
+#    ])
+#)
+#
+## Custom links
+#CUSTOMLINK_MODELS = Q(
+#    Q(app_label='circuits', model__in=[
+#        'circuit',
+#        'provider',
+#    ]) |
+#    Q(app_label='dcim', model__in=[
+#        'cable',
+#        'device',
+#        'devicetype',
+#        'powerpanel',
+#        'powerfeed',
+#        'rack',
+#        'site',
+#    ]) |
+#    Q(app_label='ipam', model__in=[
+#        'aggregate',
+#        'ipaddress',
+#        'prefix',
+#        'service',
+#        'vlan',
+#        'vrf',
+#    ]) |
+#    Q(app_label='secrets', model__in=[
+#        'secret',
+#    ]) |
+#    Q(app_label='tenancy', model__in=[
+#        'tenant',
+#    ]) |
+#    Q(app_label='virtualization', model__in=[
+#        'cluster',
+#        'virtualmachine',
+#    ])
+#)
+#
+## Models which can have Graphs associated with them
+#GRAPH_MODELS = Q(
+#    Q(app_label='circuits', model__in=[
+#        'provider',
+#    ]) |
+#    Q(app_label='dcim', model__in=[
+#        'device',
+#        'interface',
+#        'site',
+#    ])
+#)
+#
+## Models which support export templates
+#EXPORTTEMPLATE_MODELS = Q(
+#    Q(app_label='circuits', model__in=[
+#        'circuit',
+#        'provider',
+#    ]) |
+#    Q(app_label='dcim', model__in=[
+#        'cable',
+#        'consoleport',
+#        'device',
+#        'devicetype',
+#        'interface',
+#        'inventoryitem',
+#        'manufacturer',
+#        'powerpanel',
+#        'powerport',
+#        'powerfeed',
+#        'rack',
+#        'rackgroup',
+#        'region',
+#        'site',
+#        'virtualchassis',
+#    ]) |
+#    Q(app_label='ipam', model__in=[
+#        'aggregate',
+#        'ipaddress',
+#        'prefix',
+#        'service',
+#        'vlan',
+#        'vrf',
+#    ]) |
+#    Q(app_label='secrets', model__in=[
+#        'secret',
+#    ]) |
+#    Q(app_label='tenancy', model__in=[
+#        'tenant',
+#    ]) |
+#    Q(app_label='virtualization', model__in=[
+#        'cluster',
+#        'virtualmachine',
+#    ])
+#)
 
 # Report logging levels
 LOG_DEFAULT = 0
@@ -141,48 +141,58 @@ LOG_LEVEL_CODES = {
 HTTP_CONTENT_TYPE_JSON = 'application/json'
 
 # Models which support registered webhooks
-WEBHOOK_MODELS = Q(
-    Q(app_label='circuits', model__in=[
-        'circuit',
-        'provider',
-    ]) |
-    Q(app_label='dcim', model__in=[
-        'cable',
-        'consoleport',
-        'consoleserverport',
-        'device',
-        'devicebay',
-        'devicetype',
-        'frontport',
-        'interface',
-        'inventoryitem',
-        'manufacturer',
-        'poweroutlet',
-        'powerpanel',
-        'powerport',
-        'powerfeed',
-        'rack',
-        'rearport',
-        'region',
-        'site',
-        'virtualchassis',
-    ]) |
-    Q(app_label='ipam', model__in=[
-        'aggregate',
-        'ipaddress',
-        'prefix',
-        'service',
-        'vlan',
-        'vrf',
-    ]) |
-    Q(app_label='secrets', model__in=[
-        'secret',
-    ]) |
-    Q(app_label='tenancy', model__in=[
-        'tenant',
-    ]) |
-    Q(app_label='virtualization', model__in=[
-        'cluster',
-        'virtualmachine',
-    ])
-)
+#WEBHOOK_MODELS = Q(
+#    Q(app_label='circuits', model__in=[
+#        'circuit',
+#        'provider',
+#    ]) |
+#    Q(app_label='dcim', model__in=[
+#        'cable',
+#        'consoleport',
+#        'consoleserverport',
+#        'device',
+#        'devicebay',
+#        'devicetype',
+#        'frontport',
+#        'interface',
+#        'inventoryitem',
+#        'manufacturer',
+#        'poweroutlet',
+#        'powerpanel',
+#        'powerport',
+#        'powerfeed',
+#        'rack',
+#        'rearport',
+#        'region',
+#        'site',
+#        'virtualchassis',
+#    ]) |
+#    Q(app_label='ipam', model__in=[
+#        'aggregate',
+#        'ipaddress',
+#        'prefix',
+#        'service',
+#        'vlan',
+#        'vrf',
+#    ]) |
+#    Q(app_label='secrets', model__in=[
+#        'secret',
+#    ]) |
+#    Q(app_label='tenancy', model__in=[
+#        'tenant',
+#    ]) |
+#    Q(app_label='virtualization', model__in=[
+#        'cluster',
+#        'virtualmachine',
+#    ])
+#)
+
+
+# Registerable extras functionalities
+EXTRAS_FUNCTIONALITIES = [
+    'custom_fields',
+    'custom_links',
+    'graphs',
+    'export_templates',
+    'webhooks'
+]

netbox/extras/models.py

@@ -22,6 +22,7 @@ from utilities.utils import deepmerge, render_jinja2
 from .choices import *
 from .constants import *
 from .querysets import ConfigContextQuerySet
+from .utils import FunctionalityQueryset
 
 
 __all__ = (
@@ -58,7 +59,7 @@ class Webhook(models.Model):
         to=ContentType,
         related_name='webhooks',
         verbose_name='Object types',
-        limit_choices_to=WEBHOOK_MODELS,
+        limit_choices_to=FunctionalityQueryset('webhooks'),
         help_text="The object(s) to which this Webhook applies."
     )
     name = models.CharField(
@@ -223,7 +224,7 @@ class CustomField(models.Model):
         to=ContentType,
         related_name='custom_fields',
         verbose_name='Object(s)',
-        limit_choices_to=CUSTOMFIELD_MODELS,
+        limit_choices_to=FunctionalityQueryset('custom_fields'),
         help_text='The object(s) to which this field applies.'
     )
     type = models.CharField(
@@ -470,7 +471,7 @@ class CustomLink(models.Model):
     content_type = models.ForeignKey(
         to=ContentType,
         on_delete=models.CASCADE,
-        limit_choices_to=CUSTOMLINK_MODELS
+        limit_choices_to=FunctionalityQueryset('custom_links')
     )
     name = models.CharField(
         max_length=100,
@@ -518,7 +519,7 @@ class Graph(models.Model):
     type = models.ForeignKey(
         to=ContentType,
         on_delete=models.CASCADE,
-        limit_choices_to=GRAPH_MODELS
+        limit_choices_to=FunctionalityQueryset('graphs')
     )
     weight = models.PositiveSmallIntegerField(
         default=1000
@@ -581,7 +582,7 @@ class ExportTemplate(models.Model):
     content_type = models.ForeignKey(
         to=ContentType,
         on_delete=models.CASCADE,
-        limit_choices_to=EXPORTTEMPLATE_MODELS
+        limit_choices_to=FunctionalityQueryset('export_templates')
     )
     name = models.CharField(
         max_length=100

netbox/extras/utils.py

@@ -1,6 +1,11 @@
+import collections
+
+from django.db.models import Q
 from taggit.managers import _TaggableManager
 from utilities.querysets import DummyQuerySet
 
+from extras.constants import EXTRAS_FUNCTIONALITIES
+
 
 def is_taggable(obj):
     """
@@ -13,3 +18,59 @@ def is_taggable(obj):
         if isinstance(obj.tags, DummyQuerySet):
             return True
     return False
+
+
+class Registry:
+    """
+    Singleton object used to store important data
+    """
+    instance = None
+
+    def __new__(cls):
+        if cls.instance is not None:
+            return cls.instance
+        else:
+            cls.instance = super().__new__(cls)
+            cls.model_functionality_store = {f: collections.defaultdict(list) for f in EXTRAS_FUNCTIONALITIES}
+            return cls.instance
+
+
+class FunctionalityQueryset:
+    """
+    Helper class that delays evaluation of the registry contents for the functionaility store
+    until it has been populated.
+    """
+
+    def __init__(self, functionality):
+        self.functionality = functionality
+
+    def __call__(self):
+        return self.get_queryset()
+
+    def get_queryset(self):
+        """
+        Given an extras functionality, return a Q object for content type lookup
+        """
+        query = Q()
+        registry = Registry()
+        for app_label, models in registry.model_functionality_store[self.functionality].items():
+            query |= Q(app_label=app_label, model__in=models)
+
+        return query
+
+
+def extras_functionality(functionalities):
+    """
+    Decorator used to register extras provided functionalities to a model
+    """
+    def wrapper(model_class):
+        if isinstance(functionalities, list) and functionalities:
+            registry = Registry()
+            model_class._extras_functionality = []
+            for functionality in functionalities:
+                if functionality in EXTRAS_FUNCTIONALITIES:
+                    model_class._extras_functionality.append(functionality)
+                    app_label, model_name = model_class._meta.label_lower.split('.')
+                    registry.model_functionality_store[functionality][app_label].append(model_name)
+        return model_class
+    return wrapper