From e38ba299283fd602413fa6efab2a4e5f91f48fb2 Mon Sep 17 00:00:00 2001 From: Jeremy Stretch Date: Wed, 23 Jul 2025 13:00:25 -0400 Subject: [PATCH] Automatically create ObjectTypes --- .../migrations/0017_concrete_objecttype.py | 6 +- netbox/core/models/contenttypes.py | 53 ++++++++++++++++- netbox/core/signals.py | 57 +++++++++++++++---- 3 files changed, 99 insertions(+), 17 deletions(-) diff --git a/netbox/core/migrations/0017_concrete_objecttype.py b/netbox/core/migrations/0017_concrete_objecttype.py index 1c3fec669..42b624376 100644 --- a/netbox/core/migrations/0017_concrete_objecttype.py +++ b/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 ) ), diff --git a/netbox/core/models/contenttypes.py b/netbox/core/models/contenttypes.py index a49bb0882..9418172e1 100644 --- a/netbox/core/models/contenttypes.py +++ b/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() diff --git a/netbox/core/signals.py b/netbox/core/signals.py index c64d11bc9..50a360421 100644 --- a/netbox/core/signals.py +++ b/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 #