Handle all associations

startup_scripts/40_regions.py

@@ -5,13 +5,21 @@ with open('/opt/netbox/initializers/regions.yml', 'r') as stream:
   yaml=YAML(typ='safe')
   regions = yaml.load(stream)
 
-  if regions is not None:
+  optional_assocs = {
-    for region_params in regions:
+    'parent': (Region, 'name')
-      if "parent" in region_params:
+  }
-        parent = Region.objects.get(name=region_params.pop('parent'))
-        region_params['parent'] = parent
 
-      region, created = Region.objects.get_or_create(**region_params)
+  if regions is not None:
+    for params in regions:
+
+      for assoc, details in optional_assocs.items():
+        if assoc in params:
+          model, field = details
+          query = dict(field=params.pop(assoc))
+
+          params[assoc] = model.objects.get(**query)
+
+      region, created = Region.objects.get_or_create(**params)
 
       if created:
         print("Created region", region.name)

startup_scripts/41_sites.py

@@ -1,19 +1,39 @@
-from dcim.models import Site
+from dcim.models import Region, Site
-from dcim.models import Region
+from extras.models import CustomField, CustomFieldValue
+from tenancy.models import Tenant
 from ruamel.yaml import YAML
 
 with open('/opt/netbox/initializers/sites.yml', 'r') as stream:
   yaml = YAML(typ='safe')
   sites = yaml.load(stream)
 
+  optional_assocs = {
+    'region': (Region, 'name'),
+    'tenant': (Tenant, 'name')
+  }
+
   if sites is not None:
-    for site_params in sites:
+    for params in sites:
-      if "region" in site_params:
-        region = Region.objects.get(name=site_params.pop('region'))
-        site_params['region'] = region
 
+      for assoc, details in optional_assocs.items():
+        if assoc in params:
+          model, field = details
+          query = { field: params.pop(assoc) }
 
-      site, created = Site.objects.get_or_create(**site_params)
+          params[assoc] = model.objects.get(**query)
+
+      site, created = Site.objects.get_or_create(**params)
 
       if created:
+        if custom_fields is not None:
+          for cf_name, cf_value in custom_fields.items():
+            custom_field = CustomField.objects.get(name=cf_name)
+            custom_field_value = CustomFieldValue.objects.create(
+              field=custom_field,
+              obj=device_type,
+              value=cf_value
+            )
+
+            site.custom_field_values.add(custom_field_value)
+
         print("Created site", site.name)

startup_scripts/42_manufacturers.py

@@ -6,8 +6,8 @@ with open('/opt/netbox/initializers/manufacturers.yml', 'r') as stream:
   manufacturers = yaml.load(stream)
 
   if manufacturers is not None:
-    for manufacturer_params in manufacturers:
+    for params in manufacturers:
-      manufacturer, created = Manufacturer.objects.get_or_create(**manufacturer_params)
+      manufacturer, created = Manufacturer.objects.get_or_create(**params)
 
       if created:
         print("Created Manufacturer", manufacturer.name)

startup_scripts/43_device_types.py

@@ -6,20 +6,43 @@ with open('/opt/netbox/initializers/device_types.yml', 'r') as stream:
   yaml = YAML(typ='safe')
   device_types = yaml.load(stream)
 
+  required_assocs = {
+    'manufacturer': (Manufacturer, 'name')
+  }
+
+  optional_assocs = {
+    'region': (Region, 'name'),
+    'tenant': (Tenant, 'name')
+  }
+
   if device_types is not None:
-    for device_type_params in device_types:
+    for params in device_types:
-      custom_fields = device_type_params.pop('custom_fields', None)
+      custom_fields = params.pop('custom_fields', None)
 
-      manufacturer = Manufacturer.objects.get(name=device_type_params.pop('manufacturer'))
+      for assoc, details in required.items():
-      device_type_params['manufacturer'] = manufacturer
+        model, field = details
+        query = dict(field=params.pop(assoc))
 
-      device_type, created = DeviceType.objects.get_or_create(**device_type_params)
+        params[assoc] = model.objects.get(**query)
+
+      for assoc, details in optional_assocs.items():
+        if assoc in params:
+          model, field = details
+          query = dict(field=params.pop(assoc))
+
+          params[assoc] = model.objects.get(**query)
+
+      device_type, created = DeviceType.objects.get_or_create(**params)
 
       if created:
         if custom_fields is not None:
           for cf_name, cf_value in custom_fields.items():
             custom_field = CustomField.objects.get(name=cf_name)
-            custom_field_value = CustomFieldValue.objects.create(field=custom_field, obj=device_type, value=cf_value)
+            custom_field_value = CustomFieldValue.objects.create(
+              field=custom_field,
+              obj=device_type,
+              value=cf_value
+            )
 
             device_type.custom_field_values.add(custom_field_value)
 

startup_scripts/44_rack_roles.py

@@ -7,14 +7,15 @@ with open('/opt/netbox/initializers/rack_roles.yml', 'r') as stream:
   rack_roles = yaml.load(stream)
 
   if rack_roles is not None:
-    for rack_role_params in rack_roles:
+    for params in rack_roles:
-      color = rack_role_params.pop('color', None)
+      if 'color' in params:
+        color = params.pop('color')
 
         for color_tpl in COLOR_CHOICES:
           if color in color_tpl:
-          rack_role_params['color'] = color_tpl[0]
+          params['color'] = color_tpl[0]
 
-      rack_role, created = RackRole.objects.get_or_create(**rack_role_params)
+      rack_role, created = RackRole.objects.get_or_create(**params)
 
       if created:
         print("Created rack role", rack_role.name)

startup_scripts/45_racks.py

@@ -8,33 +8,52 @@ with open('/opt/netbox/initializers/racks.yml', 'r') as stream:
   yaml = YAML(typ='safe')
   racks = yaml.load(stream)
 
-  optional_assocs = dict(role=RackRole, tenant=Tenant, group=RackGroup)
+  required_assocs = {
+    'site': (Site, 'name')
+  }
+
+  optional_assocs = {
+    'role': (RackRole, 'name'),
+    'tenant': (Tenant, 'name'),
+    'group': (RackGroup, 'name')
+  }
 
   if racks is not None:
-    for rack_params in racks:
+    for params in racks:
-      custom_fields = rack_params.pop('custom_fields', None)
+      custom_fields = params.pop('custom_fields', None)
 
-      rack_params['site'] = Site.objects.get(name=rack_params.pop('site'))
+      for assoc, details in required.items():
+        model, field = details
+        query = dict(field=params.pop(assoc))
 
-      for param_name, model in optional_assoc.items():
+        params[assoc] = model.objects.get(**query)
-        if param_name in rack_params:
+
-          rack_params[param_name] = model.objects.get(name=rack_params.pop(param_name))
+      for assoc, details in optional_assocs.items():
+        if assoc in params:
+          model, field = details
+          query = dict(field=params.pop(assoc))
+
+          params[assoc] = model.objects.get(**query)
 
       for rack_type in RACK_TYPE_CHOICES:
-        if rack_params['type'] in rack_type:
+        if params['type'] in rack_type:
-          rack_params['type'] = rack_type[0]
+          params['type'] = rack_type[0]
 
       for rack_width in RACK_WIDTH_CHOICES:
-        if rack_params['width'] in rack_width:
+        if params['width'] in rack_width:
-          rack_params['width'] = rack_width[0]
+          params['width'] = rack_width[0]
 
-      rack, created = Rack.objects.get_or_create(**rack_params)
+      rack, created = Rack.objects.get_or_create(**params)
 
       if created:
         if custom_fields is not None:
           for cf_name, cf_value in custom_fields.items():
             custom_field = CustomField.objects.get(name=cf_name)
-            custom_field_value = CustomFieldValue.objects.create(field=custom_field, obj=rack, value=cf_value)
+            custom_field_value = CustomFieldValue.objects.create(
+              field=custom_field,
+              obj=device_type,
+              value=cf_value
+            )
 
             rack.custom_field_values.add(custom_field_value)
 

startup_scripts/46_device_roles.py

@@ -7,14 +7,16 @@ with open('/opt/netbox/initializers/device_roles.yml', 'r') as stream:
   device_roles = yaml.load(stream)
 
   if device_roles is not None:
-    for device_role_params in device_roles:
+    for params in device_roles:
-      color = device_role_params.pop('color')
+
+      if 'color' in params:
+        color = params.pop('color')
 
         for color_tpl in COLOR_CHOICES:
           if color in color_tpl:
-          device_role_params['color'] = color_tpl[0]
+            params['color'] = color_tpl[0]
 
-      device_role, created = DeviceRole.objects.get_or_create(**device_role_params)
+      device_role, created = DeviceRole.objects.get_or_create(**params)
 
       if created:
         print("Created device role", device_role.name)

startup_scripts/47_devices.py

@@ -1,6 +1,7 @@
 from dcim.models import Site, Rack, DeviceRole, DeviceType, Device, Platform
 from dcim.constants import RACK_FACE_CHOICES
 from ipam.models import IPAddress
+from virtualization.models import Cluster
 from tenancy.models import Tenant
 from extras.models import CustomField, CustomFieldValue
 from ruamel.yaml import YAML
@@ -9,37 +10,54 @@ with open('/opt/netbox/initializers/devices.yml', 'r') as stream:
   yaml = YAML(typ='safe')
   devices = yaml.load(stream)
 
+  required_assocs = {
+    'device_role': (DeviceRole, 'name'),
+    'device_type': (DeviceType, 'model'),
+    'site': (Site, 'name')
+  }
+
   optional_assocs = {
-    'platform': Platform,
+    'tenant': (Tenant, 'name'),
-    'tenant': Tenant
+    'platform': (Platform, 'name'),
-    'rack': Rack
+    'rack': (Rack, 'name'),
-    'primary_ip4': IPAddress
+    'cluster': (Cluster, 'name'),
-    'primary_ip6': IPAddress
+    'primary_ip4': (IPAddress, 'address')
+    'primary_ip6': (IPAddress, 'address')
   }
 
   if devices is not None:
-    for device_params in devices:
+    for params in devices:
-      custom_fields = device_params.pop('custom_fields', None)
+      custom_fields = params.pop('custom_fields', None)
 
-      device_params['device_role'] = DeviceRole.objects.get(name=device_params.pop('device_role'))
+      for assoc, details in required.items():
-      device_params['device_type'] = DeviceType.objects.get(model=device_params.pop('device_type'))
+        model, field = details
-      device_params['site'] = Site.objects.get(name=device_params.pop('site'))
+        query = dict(field=params.pop(assoc))
 
-      for param_name, model in optional_assoc.items():
+        params[assoc] = model.objects.get(**query)
-        if param_name in device_params:
-          device_params[param_name] = model.objects.get(name=device_params.pop(param_name))
 
+      for assoc, details in optional_assocs.items():
+        if assoc in params:
+          model, field = details
+          query = dict(field=params.pop(assoc))
+
+          params[assoc] = model.objects.get(**query)
+
+      if 'face' in params:
         for rack_face in RACK_FACE_CHOICES:
-        if device_params['face'] in rack_face:
+          if params['face'] in rack_face:
-          device_params['face'] = rack_face[0]
+            params['face'] = rack_face[0]
 
-      device, created = Device.objects.get_or_create(**device_params)
+      device, created = Device.objects.get_or_create(**params)
 
       if created:
         if custom_fields is not None:
           for cf_name, cf_value in custom_fields.items():
             custom_field = CustomField.objects.get(name=cf_name)
-            custom_field_value = CustomFieldValue.objects.create(field=custom_field, obj=device, value=cf_value)
+            custom_field_value = CustomFieldValue.objects.create(
+              field=custom_field,
+              obj=device_type,
+              value=cf_value
+            )
 
             device.custom_field_values.add(custom_field_value)