Просмотр исходного кода

Clean up limit_to for ForeignKeys referencing ContentType

Jeremy Stretch 6 лет назад
Родитель
Сommit
cd3924520d

+ 46 - 16
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

+ 12 - 2
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'),
+        ),
     ]

+ 12 - 7
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

+ 11 - 0
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 ['<app>.<model>', '<app>.<model>', ...] 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