Closes #3731: Change Graph.type to a ContentType foreign key field

netbox/extras/api/serializers.py

@@ -28,7 +28,9 @@ from .nested_serializers import *
 #
 
 class GraphSerializer(ValidatedModelSerializer):
-    type = ChoiceField(choices=GRAPH_TYPE_CHOICES)
+    type = ContentTypeField(
+        queryset=ContentType.objects.all()
+    )
 
     class Meta:
         model = Graph
@@ -38,7 +40,9 @@ class GraphSerializer(ValidatedModelSerializer):
 class RenderedGraphSerializer(serializers.ModelSerializer):
     embed_url = serializers.SerializerMethodField()
     embed_link = serializers.SerializerMethodField()
-    type = ChoiceField(choices=GRAPH_TYPE_CHOICES)
+    type = ContentTypeField(
+        queryset=ContentType.objects.all()
+    )
 
     class Meta:
         model = Graph

netbox/extras/constants.py

@@ -42,18 +42,6 @@ CUSTOMLINK_MODELS = [
     'virtualization.virtualmachine',
 ]
 
-# Graph types
-GRAPH_TYPE_INTERFACE = 100
-GRAPH_TYPE_DEVICE = 150
-GRAPH_TYPE_PROVIDER = 200
-GRAPH_TYPE_SITE = 300
-GRAPH_TYPE_CHOICES = (
-    (GRAPH_TYPE_INTERFACE, 'Interface'),
-    (GRAPH_TYPE_DEVICE, 'Device'),
-    (GRAPH_TYPE_PROVIDER, 'Provider'),
-    (GRAPH_TYPE_SITE, 'Site'),
-)
-
 # Models which support export templates
 EXPORTTEMPLATE_MODELS = [
     'circuits.circuit',

netbox/extras/migrations/0033_graph_type_to_fk.py

@@ -0,0 +1,46 @@
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+GRAPH_TYPE_CHOICES = (
+    (100, 'dcim', 'interface'),
+    (150, 'dcim', 'device'),
+    (200, 'circuits', 'provider'),
+    (300, 'dcim', 'site'),
+)
+
+
+def graph_type_to_fk(apps, schema_editor):
+    Graph = apps.get_model('extras', 'Graph')
+    ContentType = apps.get_model('contenttypes', 'ContentType')
+
+    # On a new installation (and during tests) content types might not yet exist. So, we only perform the bulk
+    # updates if a Graph has been created, which implies that we're working with a populated database.
+    if Graph.objects.exists():
+        for id, app_label, model in GRAPH_TYPE_CHOICES:
+            content_type = ContentType.objects.get(app_label=app_label, model=model)
+            Graph.objects.filter(type=id).update(type=content_type.pk)
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('extras', '0032_3569_webhook_fields'),
+    ]
+
+    operations = [
+        # We have to swap the legacy IDs to ContentType PKs *before* we alter the field, to avoid triggering an
+        # IntegrityError on the ForeignKey.
+        migrations.RunPython(
+            code=graph_type_to_fk
+        ),
+        migrations.AlterField(
+            model_name='graph',
+            name='type',
+            field=models.ForeignKey(
+                limit_choices_to={'model__in': ['device', 'interface', 'provider', 'site']},
+                on_delete=django.db.models.deletion.CASCADE,
+                to='contenttypes.ContentType'
+            ),
+        ),
+    ]

netbox/extras/models.py

@@ -408,8 +408,12 @@ class CustomLink(models.Model):
 #
 
 class Graph(models.Model):
-    type = models.PositiveSmallIntegerField(
-        choices=GRAPH_TYPE_CHOICES
+    type = models.ForeignKey(
+        to=ContentType,
+        on_delete=models.CASCADE,
+        limit_choices_to={
+            'model__in': ['device', 'interface', 'provider', 'site']
+        }
     )
     weight = models.PositiveSmallIntegerField(
         default=1000

netbox/extras/tests/test_api.py

@@ -4,7 +4,6 @@ from rest_framework import status
 
 from dcim.models import Device, DeviceRole, DeviceType, Manufacturer, Platform, Region, Site
 from extras.api.views import ScriptViewSet
-from extras.constants import GRAPH_TYPE_SITE
 from extras.models import ConfigContext, Graph, ExportTemplate, Tag
 from extras.scripts import BooleanVar, IntegerVar, Script, StringVar
 from tenancy.models import Tenant, TenantGroup
@@ -17,14 +16,21 @@ class GraphTest(APITestCase):
 
         super().setUp()
 
+        site_ct = ContentType.objects.get_for_model(Site)
         self.graph1 = Graph.objects.create(
-            type=GRAPH_TYPE_SITE, name='Test Graph 1', source='http://example.com/graphs.py?site={{ obj.name }}&foo=1'
+            type=site_ct,
+            name='Test Graph 1',
+            source='http://example.com/graphs.py?site={{ obj.name }}&foo=1'
         )
         self.graph2 = Graph.objects.create(
-            type=GRAPH_TYPE_SITE, name='Test Graph 2', source='http://example.com/graphs.py?site={{ obj.name }}&foo=2'
+            type=site_ct,
+            name='Test Graph 2',
+            source='http://example.com/graphs.py?site={{ obj.name }}&foo=2'
         )
         self.graph3 = Graph.objects.create(
-            type=GRAPH_TYPE_SITE, name='Test Graph 3', source='http://example.com/graphs.py?site={{ obj.name }}&foo=3'
+            type=site_ct,
+            name='Test Graph 3',
+            source='http://example.com/graphs.py?site={{ obj.name }}&foo=3'
         )
 
     def test_get_graph(self):
@@ -44,7 +50,7 @@ class GraphTest(APITestCase):
     def test_create_graph(self):
 
         data = {
-            'type': GRAPH_TYPE_SITE,
+            'type': 'dcim.site',
             'name': 'Test Graph 4',
             'source': 'http://example.com/graphs.py?site={{ obj.name }}&foo=4',
         }
@@ -55,7 +61,7 @@ class GraphTest(APITestCase):
         self.assertHttpStatus(response, status.HTTP_201_CREATED)
         self.assertEqual(Graph.objects.count(), 4)
         graph4 = Graph.objects.get(pk=response.data['id'])
-        self.assertEqual(graph4.type, data['type'])
+        self.assertEqual(graph4.type, ContentType.objects.get_for_model(Site))
         self.assertEqual(graph4.name, data['name'])
         self.assertEqual(graph4.source, data['source'])
 
@@ -63,17 +69,17 @@ class GraphTest(APITestCase):
 
         data = [
             {
-                'type': GRAPH_TYPE_SITE,
+                'type': 'dcim.site',
                 'name': 'Test Graph 4',
                 'source': 'http://example.com/graphs.py?site={{ obj.name }}&foo=4',
             },
             {
-                'type': GRAPH_TYPE_SITE,
+                'type': 'dcim.site',
                 'name': 'Test Graph 5',
                 'source': 'http://example.com/graphs.py?site={{ obj.name }}&foo=5',
             },
             {
-                'type': GRAPH_TYPE_SITE,
+                'type': 'dcim.site',
                 'name': 'Test Graph 6',
                 'source': 'http://example.com/graphs.py?site={{ obj.name }}&foo=6',
             },
@@ -91,7 +97,7 @@ class GraphTest(APITestCase):
     def test_update_graph(self):
 
         data = {
-            'type': GRAPH_TYPE_SITE,
+            'type': 'dcim.site',
             'name': 'Test Graph X',
             'source': 'http://example.com/graphs.py?site={{ obj.name }}&foo=99',
         }
@@ -102,7 +108,7 @@ class GraphTest(APITestCase):
         self.assertHttpStatus(response, status.HTTP_200_OK)
         self.assertEqual(Graph.objects.count(), 3)
         graph1 = Graph.objects.get(pk=response.data['id'])
-        self.assertEqual(graph1.type, data['type'])
+        self.assertEqual(graph1.type, ContentType.objects.get_for_model(Site))
         self.assertEqual(graph1.name, data['name'])
         self.assertEqual(graph1.source, data['source'])