Add vlan_groups to Region, SiteGroup, Site, Location, Rack, ClusterGroup, Cluster

netbox/dcim/graphql/types.py

@@ -1,6 +1,6 @@
 from dcim import filtersets, models
 from extras.graphql.mixins import ImageAttachmentsMixin
-from ipam.graphql.mixins import IPAddressesMixin
+from ipam.graphql.mixins import IPAddressesMixin, VLANGroupsMixin
 from netbox.graphql.types import BaseObjectType, ObjectType, TaggedObjectType
 
 __all__ = (
@@ -187,7 +187,7 @@ class InventoryItemType(TaggedObjectType):
         filterset_class = filtersets.InventoryItemFilterSet
 
 
-class LocationType(ImageAttachmentsMixin, ObjectType):
+class LocationType(VLANGroupsMixin, ImageAttachmentsMixin, ObjectType):
 
     class Meta:
         model = models.Location
@@ -277,7 +277,7 @@ class PowerPortTemplateType(BaseObjectType):
         return self.type or None
 
 
-class RackType(ImageAttachmentsMixin, TaggedObjectType):
+class RackType(VLANGroupsMixin, ImageAttachmentsMixin, TaggedObjectType):
 
     class Meta:
         model = models.Rack
@@ -323,7 +323,7 @@ class RearPortTemplateType(BaseObjectType):
         filterset_class = filtersets.RearPortTemplateFilterSet
 
 
-class RegionType(ObjectType):
+class RegionType(VLANGroupsMixin, ObjectType):
 
     class Meta:
         model = models.Region
@@ -331,7 +331,7 @@ class RegionType(ObjectType):
         filterset_class = filtersets.RegionFilterSet
 
 
-class SiteType(ImageAttachmentsMixin, TaggedObjectType):
+class SiteType(VLANGroupsMixin, ImageAttachmentsMixin, TaggedObjectType):
 
     class Meta:
         model = models.Site
@@ -339,7 +339,7 @@ class SiteType(ImageAttachmentsMixin, TaggedObjectType):
         filterset_class = filtersets.SiteFilterSet
 
 
-class SiteGroupType(ObjectType):
+class SiteGroupType(VLANGroupsMixin, ObjectType):
 
     class Meta:
         model = models.SiteGroup

netbox/dcim/models/racks.py

@@ -175,6 +175,12 @@ class Rack(PrimaryModel):
     comments = models.TextField(
         blank=True
     )
+    vlan_groups = GenericRelation(
+        to='ipam.VLANGroup',
+        content_type_field='scope_type',
+        object_id_field='scope_id',
+        related_query_name='rack'
+    )
     images = GenericRelation(
         to='extras.ImageAttachment'
     )

netbox/dcim/models/sites.py

@@ -53,6 +53,12 @@ class Region(NestedGroupModel):
         max_length=200,
         blank=True
     )
+    vlan_groups = GenericRelation(
+        to='ipam.VLANGroup',
+        content_type_field='scope_type',
+        object_id_field='scope_id',
+        related_query_name='region'
+    )
 
     def get_absolute_url(self):
         return reverse('dcim:region', args=[self.pk])
@@ -95,6 +101,12 @@ class SiteGroup(NestedGroupModel):
         max_length=200,
         blank=True
     )
+    vlan_groups = GenericRelation(
+        to='ipam.VLANGroup',
+        content_type_field='scope_type',
+        object_id_field='scope_id',
+        related_query_name='site_group'
+    )
 
     def get_absolute_url(self):
         return reverse('dcim:sitegroup', args=[self.pk])
@@ -210,6 +222,12 @@ class Site(PrimaryModel):
     comments = models.TextField(
         blank=True
     )
+    vlan_groups = GenericRelation(
+        to='ipam.VLANGroup',
+        content_type_field='scope_type',
+        object_id_field='scope_id',
+        related_query_name='site'
+    )
     images = GenericRelation(
         to='extras.ImageAttachment'
     )
@@ -267,6 +285,12 @@ class Location(NestedGroupModel):
         max_length=200,
         blank=True
     )
+    vlan_groups = GenericRelation(
+        to='ipam.VLANGroup',
+        content_type_field='scope_type',
+        object_id_field='scope_id',
+        related_query_name='location'
+    )
     images = GenericRelation(
         to='extras.ImageAttachment'
     )

netbox/ipam/graphql/mixins.py

@@ -2,6 +2,7 @@ import graphene
 
 __all__ = (
     'IPAddressesMixin',
+    'VLANGroupsMixin',
 )
 
 
@@ -10,3 +11,10 @@ class IPAddressesMixin:
 
     def resolve_ip_addresses(self, info):
         return self.ip_addresses.restrict(info.context.user, 'view')
+
+
+class VLANGroupsMixin:
+    vlan_groups = graphene.List('ipam.graphql.types.VLANGroupType')
+
+    def resolve_vlan_groups(self, info):
+        return self.vlan_groups.restrict(info.context.user, 'view')

netbox/virtualization/graphql/types.py

@@ -1,4 +1,4 @@
-from ipam.graphql.mixins import IPAddressesMixin
+from ipam.graphql.mixins import IPAddressesMixin, VLANGroupsMixin
 from virtualization import filtersets, models
 from netbox.graphql.types import ObjectType, TaggedObjectType
 
@@ -11,7 +11,7 @@ __all__ = (
 )
 
 
-class ClusterType(TaggedObjectType):
+class ClusterType(VLANGroupsMixin, TaggedObjectType):
 
     class Meta:
         model = models.Cluster
@@ -19,7 +19,7 @@ class ClusterType(TaggedObjectType):
         filterset_class = filtersets.ClusterFilterSet
 
 
-class ClusterGroupType(ObjectType):
+class ClusterGroupType(VLANGroupsMixin, ObjectType):
 
     class Meta:
         model = models.ClusterGroup

netbox/virtualization/models.py

@@ -81,6 +81,12 @@ class ClusterGroup(OrganizationalModel):
         max_length=200,
         blank=True
     )
+    vlan_groups = GenericRelation(
+        to='ipam.VLANGroup',
+        content_type_field='scope_type',
+        object_id_field='scope_id',
+        related_query_name='cluster_group'
+    )
 
     objects = RestrictedQuerySet.as_manager()
 
@@ -136,6 +142,12 @@ class Cluster(PrimaryModel):
     comments = models.TextField(
         blank=True
     )
+    vlan_groups = GenericRelation(
+        to='ipam.VLANGroup',
+        content_type_field='scope_type',
+        object_id_field='scope_id',
+        related_query_name='cluster'
+    )
 
     objects = RestrictedQuerySet.as_manager()