Fix setting CF data if CF object is missing

kr3ator · kr3ator · commit 81d9e4f5602d · 2022-04-06T14:09:07.000+02:00
diff --git a/startup_scripts/080_tenants.py b/startup_scripts/080_tenants.py
@@ -23,6 +23,6 @@
     tenant, created = Tenant.objects.get_or_create(**params)
 
     if created:
-        set_custom_fields_values(tenant, custom_field_data)
-
         print("👩‍💻 Created Tenant", tenant.name)
+
+    set_custom_fields_values(tenant, custom_field_data)
diff --git a/startup_scripts/110_sites.py b/startup_scripts/110_sites.py
@@ -24,6 +24,6 @@
     site, created = Site.objects.get_or_create(**params)
 
     if created:
-        set_custom_fields_values(site, custom_field_data)
-
         print("📍 Created site", site.name)
+
+    set_custom_fields_values(site, custom_field_data)
diff --git a/startup_scripts/140_racks.py b/startup_scripts/140_racks.py
@@ -36,6 +36,6 @@
     rack, created = Rack.objects.get_or_create(**params)
 
     if created:
-        set_custom_fields_values(rack, custom_field_data)
-
         print("🔳 Created rack", rack.site, rack.name)
+
+    set_custom_fields_values(rack, custom_field_data)
diff --git a/startup_scripts/150_power_panels.py b/startup_scripts/150_power_panels.py
@@ -31,6 +31,6 @@
     power_panel, created = PowerPanel.objects.get_or_create(**params)
 
     if created:
-        set_custom_fields_values(power_panel, custom_field_data)
-
         print("⚡ Created Power Panel", power_panel.site, power_panel.name)
+
+    set_custom_fields_values(power_panel, custom_field_data)
diff --git a/startup_scripts/160_power_feeds.py b/startup_scripts/160_power_feeds.py
@@ -31,6 +31,6 @@
     power_feed, created = PowerFeed.objects.get_or_create(**params)
 
     if created:
-        set_custom_fields_values(power_feed, custom_field_data)
-
         print("⚡ Created Power Feed", power_feed.name)
+
+    set_custom_fields_values(power_feed, custom_field_data)
diff --git a/startup_scripts/190_device_types.py b/startup_scripts/190_device_types.py
@@ -32,6 +32,6 @@
     device_type, created = DeviceType.objects.get_or_create(**params)
 
     if created:
-        set_custom_fields_values(device_type, custom_field_data)
-
         print("🔡 Created device type", device_type.manufacturer, device_type.model)
+
+    set_custom_fields_values(device_type, custom_field_data)
diff --git a/startup_scripts/200_devices.py b/startup_scripts/200_devices.py
@@ -47,6 +47,6 @@
     device, created = Device.objects.get_or_create(**params)
 
     if created:
-        set_custom_fields_values(device, custom_field_data)
-
         print("🖥️  Created device", device.name)
+
+    set_custom_fields_values(device, custom_field_data)
diff --git a/startup_scripts/210_dcim_interfaces.py b/startup_scripts/210_dcim_interfaces.py
@@ -22,6 +22,6 @@
     interface, created = Interface.objects.get_or_create(**params)
 
     if created:
-        set_custom_fields_values(interface, custom_field_data)
-
         print("🧷 Created interface", interface.name, interface.device.name)
+
+    set_custom_fields_values(interface, custom_field_data)
diff --git a/startup_scripts/230_route_targets.py b/startup_scripts/230_route_targets.py
@@ -24,6 +24,6 @@
     route_target, created = RouteTarget.objects.get_or_create(**params)
 
     if created:
-        set_custom_fields_values(route_target, custom_field_data)
-
         print("🎯 Created Route Target", route_target.name)
+
+    set_custom_fields_values(route_target, custom_field_data)
diff --git a/startup_scripts/240_vrfs.py b/startup_scripts/240_vrfs.py
@@ -24,6 +24,6 @@
     vrf, created = VRF.objects.get_or_create(**params)
 
     if created:
-        set_custom_fields_values(vrf, custom_field_data)
-
         print("📦 Created VRF", vrf.name)
+
+    set_custom_fields_values(vrf, custom_field_data)
diff --git a/startup_scripts/270_aggregates.py b/startup_scripts/270_aggregates.py
@@ -37,6 +37,6 @@
     aggregate, created = Aggregate.objects.get_or_create(**params)
 
     if created:
-        set_custom_fields_values(aggregate, custom_field_data)
-
         print("🗞️ Created Aggregate", aggregate.prefix)
+
+    set_custom_fields_values(aggregate, custom_field_data)
diff --git a/startup_scripts/310_clusters.py b/startup_scripts/310_clusters.py
@@ -37,6 +37,6 @@
     cluster, created = Cluster.objects.get_or_create(**params)
 
     if created:
-        set_custom_fields_values(cluster, custom_field_data)
-
         print("🗄️ Created cluster", cluster.name)
+
+    set_custom_fields_values(cluster, custom_field_data)
diff --git a/startup_scripts/320_vlan_groups.py b/startup_scripts/320_vlan_groups.py
@@ -35,6 +35,6 @@
     vlan_group, created = VLANGroup.objects.get_or_create(**params)
 
     if created:
-        set_custom_fields_values(vlan_group, custom_field_data)
-
         print("🏘️ Created VLAN Group", vlan_group.name)
+
+    set_custom_fields_values(vlan_group, custom_field_data)
diff --git a/startup_scripts/330_vlans.py b/startup_scripts/330_vlans.py
@@ -31,6 +31,6 @@
     vlan, created = VLAN.objects.get_or_create(**params)
 
     if created:
-        set_custom_fields_values(vlan, custom_field_data)
-
         print("🏠 Created VLAN", vlan.name)
+
+    set_custom_fields_values(vlan, custom_field_data)
diff --git a/startup_scripts/340_virtual_machines.py b/startup_scripts/340_virtual_machines.py
@@ -41,6 +41,6 @@
     virtual_machine, created = VirtualMachine.objects.get_or_create(**params)
 
     if created:
-        set_custom_fields_values(virtual_machine, custom_field_data)
-
         print("🖥️ Created virtual machine", virtual_machine.name)
+
+    set_custom_fields_values(virtual_machine, custom_field_data)
diff --git a/startup_scripts/350_virtualization_interfaces.py b/startup_scripts/350_virtualization_interfaces.py
@@ -22,6 +22,6 @@
     interface, created = VMInterface.objects.get_or_create(**params)
 
     if created:
-        set_custom_fields_values(interface, custom_field_data)
-
         print("🧷 Created interface", interface.name, interface.virtual_machine.name)
+
+    set_custom_fields_values(interface, custom_field_data)
diff --git a/startup_scripts/360_prefixes.py b/startup_scripts/360_prefixes.py
@@ -34,6 +34,6 @@
     prefix, created = Prefix.objects.get_or_create(**params)
 
     if created:
-        set_custom_fields_values(prefix, custom_field_data)
-
         print("📌 Created Prefix", prefix.prefix)
+
+    set_custom_fields_values(prefix, custom_field_data)
diff --git a/startup_scripts/370_ip_addresses.py b/startup_scripts/370_ip_addresses.py
@@ -58,6 +58,6 @@
     ip_address, created = IPAddress.objects.get_or_create(**params)
 
     if created:
-        set_custom_fields_values(ip_address, custom_field_data)
-
         print("🧬 Created IP Address", ip_address.address)
+
+    set_custom_fields_values(ip_address, custom_field_data)
diff --git a/startup_scripts/420_providers.py b/startup_scripts/420_providers.py
@@ -14,6 +14,6 @@
     provider, created = Provider.objects.get_or_create(**params)
 
     if created:
-        set_custom_fields_values(provider, custom_field_data)
-
         print("📡 Created provider", provider.name)
+
+    set_custom_fields_values(provider, custom_field_data)
diff --git a/startup_scripts/440_circuit_types.py b/startup_scripts/440_circuit_types.py
@@ -14,6 +14,6 @@
     circuit_type, created = CircuitType.objects.get_or_create(**params)
 
     if created:
-        set_custom_fields_values(circuit_type, custom_field_data)
-
         print("⚡ Created Circuit Type", circuit_type.name)
+
+    set_custom_fields_values(circuit_type, custom_field_data)
diff --git a/startup_scripts/450_circuits.py b/startup_scripts/450_circuits.py
@@ -32,6 +32,6 @@
     circuit, created = Circuit.objects.get_or_create(**params)
 
     if created:
-        set_custom_fields_values(circuit, custom_field_data)
-
         print("⚡ Created Circuit", circuit.cid)
+
+    set_custom_fields_values(circuit, custom_field_data)
diff --git a/startup_scripts/startup_script_utils/custom_fields.py b/startup_scripts/startup_script_utils/custom_fields.py
@@ -1,9 +1,38 @@
+from django.contrib.contenttypes.models import ContentType
+from django.core.exceptions import ObjectDoesNotExist
+from extras.models import CustomField
+
+
 def set_custom_fields_values(entity, custom_field_data):
     if not custom_field_data:
         return
 
-    entity.custom_field_data = custom_field_data
-    return entity.save()
+    missing_cfs = []
+    save = False
+    for key, value in custom_field_data.items():
+        try:
+            cf = CustomField.objects.get(name=key)
+        except ObjectDoesNotExist:
+            missing_cfs.append(key)
+        else:
+            ct = ContentType.objects.get_for_model(entity)
+            if ct not in cf.content_types.all():
+                print(
+                    f"⚠️ Custom field {key} is not enabled for {entity}'s model!"
+                    "Please check the 'on_objects' for that custom field in custom_fields.yml"
+                )
+            elif key not in entity.custom_field_data:
+                entity.custom_field_data[key] = value
+                save = True
+
+    if missing_cfs:
+        raise Exception(
+            f"⚠️ Custom field(s) '{missing_cfs}' requested for {entity} but not found in Netbox!"
+            "Please chceck the custom_fields.yml"
+        )
+
+    if save:
+        entity.save()
 
 
 def pop_custom_fields(params):
