Catch post_migrate signal to update ObjectTypes

netbox/core/models/contenttypes.py

@@ -1,7 +1,6 @@
 from django.contrib.contenttypes.models import ContentType, ContentTypeManager
 from django.contrib.postgres.fields import ArrayField
 from django.db import models
-from django.db.models import Q
 from django.utils.translation import gettext as _
 
 from netbox.registry import registry
@@ -16,31 +15,23 @@ class ObjectTypeManager(ContentTypeManager):
 
     def public(self):
         """
-        Filter the base queryset to return only ContentTypes corresponding to "public" models; those which are listed
-        in registry['models'] and intended for reference by other objects.
+        Filter the base queryset to return only ObjectTypes corresponding to "public" models; those which are intended
+        for reference by other objects.
         """
-        q = Q()
-        for app_label, model_list in registry['models'].items():
-            q |= Q(app_label=app_label, model__in=model_list)
-        return self.get_queryset().filter(q)
+        return self.get_queryset().filter(public=True)
 
     def with_feature(self, feature):
         """
         Return the ContentTypes only for models which are registered as supporting the specified feature. For example,
-        we can find all ContentTypes for models which support webhooks with
+        we can find all ContentTypes for models which support event rules with:
 
-            ContentType.objects.with_feature('event_rules')
+            ObjectType.objects.with_feature('event_rules')
         """
         if feature not in registry['model_features']:
             raise KeyError(
                 f"{feature} is not a registered model feature! Valid features are: {registry['model_features'].keys()}"
             )
-
-        q = Q()
-        for app_label, model_list in registry['model_features'][feature].items():
-            q |= Q(app_label=app_label, model__in=model_list)
-
-        return self.get_queryset().filter(q)
+        return self.get_queryset().filter(features__contains=[feature])
 
 
 class ObjectType(ContentType):

netbox/core/signals.py

@@ -3,18 +3,19 @@ import logging
 from django.contrib.contenttypes.models import ContentType
 from django.core.exceptions import ValidationError
 from django.db.models.fields.reverse_related import ManyToManyRel, ManyToOneRel
-from django.db.models.signals import m2m_changed, post_save, pre_delete
+from django.db.models.signals import m2m_changed, post_migrate, post_save, pre_delete
 from django.dispatch import receiver, Signal
 from django.utils.translation import gettext_lazy as _
 from django_prometheus.models import model_deletes, model_inserts, model_updates
 
 from core.choices import JobStatusChoices, ObjectChangeActionChoices
 from core.events import *
+from core.models import ObjectType
 from extras.events import enqueue_event
 from extras.utils import run_validators
 from netbox.config import get_config
 from netbox.context import current_request, events_queue
-from netbox.models.features import ChangeLoggingMixin
+from netbox.models.features import ChangeLoggingMixin, FEATURES_MAP
 from utilities.exceptions import AbortRequest
 from .models import ConfigRevision, DataSource, ObjectChange
 
@@ -38,6 +39,33 @@ post_sync = Signal()
 clear_events = Signal()
 
 
+#
+# Model registration
+#
+
+@receiver(post_migrate)
+def update_object_types(sender, **kwargs):
+    models = sender.get_models()
+
+    for model in models:
+        app_label, model_name = model._meta.label_lower.split('.')
+
+        # Determine whether model is public
+        is_public = not getattr(model, '_netbox_private', False)
+
+        # Determine NetBox features supported by the model
+        features = [
+            feature for feature, cls in FEATURES_MAP.items() if issubclass(model, cls)
+        ]
+
+        # TODO: Update ObjectTypes in bulk
+        # Update the ObjectType for the model
+        ObjectType.objects.filter(app_label=app_label, model=model_name).update(
+            public=is_public,
+            features=features,
+        )
+
+
 #
 # Change logging & event handling
 #