From cd3924520db38a2f550f4b7fbe81e3e2518761e1 Mon Sep 17 00:00:00 2001 From: Jeremy Stretch Date: Wed, 17 Apr 2019 13:36:05 -0400 Subject: [PATCH] Clean up limit_to for ForeignKeys referencing ContentType --- netbox/extras/constants.py | 62 ++++++++++++++----- netbox/extras/migrations/0022_custom_links.py | 14 ++++- netbox/extras/models.py | 19 +++--- netbox/utilities/utils.py | 11 ++++ 4 files changed, 81 insertions(+), 25 deletions(-) diff --git a/netbox/extras/constants.py b/netbox/extras/constants.py index 4daf32839..cf1822352 100644 --- a/netbox/extras/constants.py +++ b/netbox/extras/constants.py @@ -1,13 +1,24 @@ # Models which support custom fields -CUSTOMFIELD_MODELS = ( - 'provider', 'circuit', # Circuits - 'site', 'rack', 'devicetype', 'device', # DCIM - 'aggregate', 'prefix', 'ipaddress', 'vlan', 'vrf', 'service', # IPAM - 'secret', # Secrets - 'tenant', # Tenancy - 'cluster', 'virtualmachine', # Virtualization -) +CUSTOMFIELD_MODELS = [ + 'circuits.circuit', + 'circuits.provider', + 'dcim.device', + 'dcim.devicetype', + 'dcim.powerfeed', + 'dcim.rack', + 'dcim.site', + 'ipam.aggregate', + 'ipam.ipaddress', + 'ipam.prefix', + 'ipam.service', + 'ipam.vlan', + 'ipam.vrf', + 'secrets.secret', + 'tenancy.tenant', + 'virtualization.cluster', + 'virtualization.virtualmachine', +] # Custom field types CF_TYPE_TEXT = 100 @@ -36,7 +47,7 @@ CF_FILTER_CHOICES = ( ) # Custom links -CUSTOM_LINK_MODELS = [ +CUSTOMLINK_MODELS = [ 'circuits.circuit', 'circuits.provider', 'dcim.cable', @@ -87,13 +98,32 @@ GRAPH_TYPE_CHOICES = ( # Models which support export templates EXPORTTEMPLATE_MODELS = [ - 'provider', 'circuit', # Circuits - 'site', 'region', 'rack', 'rackgroup', 'manufacturer', 'devicetype', 'device', # DCIM - 'consoleport', 'powerport', 'interface', 'cable', 'virtualchassis', # DCIM - 'aggregate', 'prefix', 'ipaddress', 'vlan', 'vrf', 'service', # IPAM - 'secret', # Secrets - 'tenant', # Tenancy - 'cluster', 'virtualmachine', # Virtualization + 'circuits.circuit', + 'circuits.provider', + 'dcim.cable', + 'dcim.consoleport', + 'dcim.device', + 'dcim.devicetype', + 'dcim.interface', + 'dcim.manufacturer', + 'dcim.powerpanel', + 'dcim.powerport', + 'dcim.powerfeed', + 'dcim.rack', + 'dcim.rackgroup', + 'dcim.region', + 'dcim.site', + 'dcim.virtualchassis', + 'ipam.aggregate', + 'ipam.ipaddress', + 'ipam.prefix', + 'ipam.service', + 'ipam.vlan', + 'ipam.vrf', + 'secrets.secret', + 'tenancy.tenant', + 'virtualization.cluster', + 'virtualization.virtualmachine', ] # ExportTemplate language choices diff --git a/netbox/extras/migrations/0022_custom_links.py b/netbox/extras/migrations/0022_custom_links.py index 7762606de..45f43059d 100644 --- a/netbox/extras/migrations/0022_custom_links.py +++ b/netbox/extras/migrations/0022_custom_links.py @@ -1,5 +1,3 @@ -# Generated by Django 2.2 on 2019-04-15 19:28 - from django.db import migrations, models import django.db.models.deletion import extras.models @@ -30,4 +28,16 @@ class Migration(migrations.Migration): 'ordering': ['group_name', 'weight', 'name'], }, ), + + # Update limit_choices_to for CustomFields and ExportTemplates + migrations.AlterField( + model_name='customfield', + name='obj_type', + field=models.ManyToManyField(limit_choices_to=extras.models.get_custom_field_models, related_name='custom_fields', to='contenttypes.ContentType'), + ), + migrations.AlterField( + model_name='exporttemplate', + name='content_type', + field=models.ForeignKey(limit_choices_to=extras.models.get_export_template_models, on_delete=django.db.models.deletion.CASCADE, to='contenttypes.ContentType'), + ), ] diff --git a/netbox/extras/models.py b/netbox/extras/models.py index 0f151fb39..49a4f7022 100644 --- a/netbox/extras/models.py +++ b/netbox/extras/models.py @@ -17,7 +17,7 @@ from taggit.models import TagBase, GenericTaggedItemBase from dcim.constants import CONNECTION_STATUS_CONNECTED from utilities.fields import ColorField -from utilities.utils import deepmerge, foreground_color +from utilities.utils import deepmerge, foreground_color, model_names_to_filter_dict from .constants import * from .querysets import ConfigContextQuerySet @@ -134,12 +134,16 @@ class CustomFieldModel(models.Model): return OrderedDict([(field, None) for field in fields]) +def get_custom_field_models(): + return model_names_to_filter_dict(CUSTOMFIELD_MODELS) + + class CustomField(models.Model): obj_type = models.ManyToManyField( to=ContentType, related_name='custom_fields', verbose_name='Object(s)', - limit_choices_to={'model__in': CUSTOMFIELD_MODELS}, + limit_choices_to=get_custom_field_models, help_text='The object(s) to which this field applies.' ) type = models.PositiveSmallIntegerField( @@ -305,10 +309,7 @@ class CustomFieldChoice(models.Model): # def get_custom_link_models(): - # TODO: This should match on the app_label as well as the model name to avoid potential duplicate names - return { - 'model__in': [model.split('.')[1] for model in CUSTOM_LINK_MODELS], - } + return model_names_to_filter_dict(CUSTOMLINK_MODELS) class CustomLink(models.Model): @@ -404,11 +405,15 @@ class Graph(models.Model): # Export templates # +def get_export_template_models(): + return model_names_to_filter_dict(EXPORTTEMPLATE_MODELS) + + class ExportTemplate(models.Model): content_type = models.ForeignKey( to=ContentType, on_delete=models.CASCADE, - limit_choices_to={'model__in': EXPORTTEMPLATE_MODELS} + limit_choices_to=get_export_template_models ) name = models.CharField( max_length=100 diff --git a/netbox/utilities/utils.py b/netbox/utilities/utils.py index c323bf473..ff8ec6297 100644 --- a/netbox/utilities/utils.py +++ b/netbox/utilities/utils.py @@ -60,6 +60,17 @@ def dynamic_import(name): return mod +def model_names_to_filter_dict(names): + """ + Accept a list of content types in the format ['.', '.', ...] and return a dictionary + suitable for QuerySet filtering. + """ + # TODO: This should match on the app_label as well as the model name to avoid potential duplicate names + return { + 'model__in': [model.split('.')[1] for model in names], + } + + def serialize_object(obj, extra=None): """ Return a generic JSON representation of an object using Django's built-in serializer. (This is used for things like