Catch post_migrate signal to update ObjectTypes

This commit is contained in:
Jeremy Stretch 2025-07-22 10:59:48 -04:00
parent b64304d3dd
commit 749eed8f02
2 changed files with 36 additions and 17 deletions

View File

@ -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):

View File

@ -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
#