Automatically create ObjectTypes

netbox/core/migrations/0017_concrete_objecttype.py

@@ -18,7 +18,8 @@ def populate_object_types(apps, schema_editor):
             apps.get_model(ct.app_label, ct.model)
         except LookupError:
             continue
-        ObjectType(pk=ct.pk).save_base(raw=True)
+        # TODO assign public/features
+        ObjectType(pk=ct.pk, features=[]).save_base(raw=True)
 
 
 class Migration(migrations.Migration):
@@ -58,8 +59,7 @@ class Migration(migrations.Migration):
                     'features',
                     django.contrib.postgres.fields.ArrayField(
                         base_field=models.CharField(max_length=50),
-                        blank=True,
-                        null=True,
+                        default=list,
                         size=None
                     )
                 ),

netbox/core/models/contenttypes.py

@@ -1,5 +1,6 @@
 from django.contrib.contenttypes.models import ContentType, ContentTypeManager
 from django.contrib.postgres.fields import ArrayField
+from django.core.exceptions import ObjectDoesNotExist
 from django.db import models
 from django.utils.translation import gettext as _
 
@@ -13,8 +14,57 @@ __all__ = (
 )
 
 
+class ObjectTypeQuerySet(models.QuerySet):
+
+    def create(self, **kwargs):
+        # If attempting to create a new ObjectType for a given app_label & model, replace those kwargs
+        # with a reference to the ContentType (if one exists).
+        if (app_label := kwargs.get('app_label')) and (model := kwargs.get('model')):
+            try:
+                kwargs['contenttype_ptr'] = ContentType.objects.get(app_label=app_label, model=model)
+                kwargs.pop('app_label')
+                kwargs.pop('model')
+            except ObjectDoesNotExist:
+                pass
+        return super().create(**kwargs)
+
+
 class ObjectTypeManager(ContentTypeManager):
 
+    def get_queryset(self):
+        return ObjectTypeQuerySet(self.model, using=self._db)
+
+    def create(self, **kwargs):
+        return self.get_queryset().create(**kwargs)
+
+    def get_for_model(self, model, for_concrete_model=True):
+        """
+        Return the ContentType object for a given model, creating the
+        ContentType if necessary. Lookups are cached so that subsequent lookups
+        for the same model don't hit the database.
+        """
+        opts = self._get_opts(model, for_concrete_model)
+        try:
+            return self._get_from_cache(opts)
+        except KeyError:
+            pass
+
+        # The ContentType entry was not found in the cache, therefore we
+        # proceed to load or create it.
+        try:
+            # Start with get() and not get_or_create() in order to use
+            # the db_for_read (see #20401).
+            ct = self.get(app_label=opts.app_label, model=opts.model_name)
+        except self.model.DoesNotExist:
+            # Not found in the database; we proceed to create it. This time
+            # use get_or_create to take care of any race conditions.
+            ct, __ = self.get_or_create(
+                app_label=opts.app_label,
+                model=opts.model_name,
+            )
+        self._add_to_cache(self.db, ct)
+        return ct
+
     def public(self):
         """
         Filter the base queryset to return only ObjectTypes corresponding to "public" models; those which are intended
@@ -53,8 +103,7 @@ class ObjectType(ContentType):
     )
     features = ArrayField(
         base_field=models.CharField(max_length=50),
-        blank=True,
-        null=True,
+        default=list,
     )
 
     objects = ObjectTypeManager()

netbox/core/signals.py

@@ -1,7 +1,8 @@
 import logging
 
 from django.contrib.contenttypes.models import ContentType
-from django.core.exceptions import ValidationError
+from django.core.exceptions import ObjectDoesNotExist, ValidationError
+from django.db import ProgrammingError
 from django.db.models.fields.reverse_related import ManyToManyRel, ManyToOneRel
 from django.db.models.signals import m2m_changed, post_migrate, post_save, pre_delete
 from django.dispatch import receiver, Signal
@@ -39,8 +40,18 @@ post_sync = Signal()
 clear_events = Signal()
 
 
+def model_is_public(model):
+    return not getattr(model, '_netbox_private', False)
+
+
+def get_model_features(model):
+    return [
+        feature for feature, cls in FEATURES_MAP.items() if issubclass(model, cls)
+    ]
+
+
 #
-# Model registration
+# Object types
 #
 
 @receiver(post_migrate)
@@ -51,19 +62,41 @@ def update_object_types(sender, **kwargs):
         app_label, model_name = model._meta.label_lower.split('.')
 
         # Determine whether model is public
-        is_public = not getattr(model, '_netbox_private', False)
+        is_public = model_is_public(model)
 
         # Determine NetBox features supported by the model
-        features = [
-            feature for feature, cls in FEATURES_MAP.items() if issubclass(model, cls)
-        ]
+        features = get_model_features(model)
 
-        # 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,
-        )
+        # Create/update the ObjectType for the model
+        try:
+            ot = ObjectType.objects.get_by_natural_key(app_label=app_label, model=model_name)
+            ot.public = is_public
+            ot.features = features
+        except ObjectDoesNotExist:
+            ct = ContentType.objects.get_for_model(model)
+            ot = ObjectType(
+                contenttype_ptr=ct,
+                app_label=app_label,
+                model=model_name,
+                public=is_public,
+                features=features,
+            )
+        ot.save()
+
+
+@receiver(post_save, sender=ContentType)
+def create_object_type(sender, instance, created, **kwargs):
+    if created:
+        model = instance.model_class()
+        try:
+            ObjectType.objects.create(
+                contenttype_ptr=instance,
+                public=model_is_public(model),
+                features=get_model_features(model),
+            )
+        except ProgrammingError:
+            # Will fail during migrations if ObjectType hasn't been created yet
+            pass
 
 
 #