9817 ipan graphql gfk

netbox/ipam/graphql/gfk_mixins.py

@@ -1,10 +1,17 @@
 import graphene
-from dcim.graphql.types import InterfaceType
-from dcim.models import Interface
+from dcim.graphql.types import (
+    InterfaceType,
+    LocationType,
+    RackType,
+    RegionType,
+    SiteGroupType,
+    SiteType,
+)
+from dcim.models import Interface, Location, Rack, Region, Site, SiteGroup
 from ipam.graphql.types import FHRPGroupType, VLANType
-from ipam.models import FHRPGroup, VLAN
-from virtualization.graphql.types import VMInterfaceType
-from virtualization.models import VMInterface
+from ipam.models import VLAN, FHRPGroup
+from virtualization.graphql.types import ClusterGroupType, ClusterType, VMInterfaceType
+from virtualization.models import Cluster, ClusterGroup, VMInterface
 
 
 class IPAddressAssignmentType(graphene.Union):
@@ -41,3 +48,48 @@ class L2VPNAssignmentType(graphene.Union):
             return VLANType
         if type(instance) == VMInterface:
             return VMInterfaceType
+
+
+class FHRPGroupInterfaceType(graphene.Union):
+    class Meta:
+        types = (
+            InterfaceType,
+            VMInterfaceType,
+        )
+
+    @classmethod
+    def resolve_type(cls, instance, info):
+        if type(instance) == Interface:
+            return InterfaceType
+        if type(instance) == VMInterface:
+            return VMInterfaceType
+
+
+class VLANGroupScopeType(graphene.Union):
+    class Meta:
+        types = (
+            ClusterType,
+            ClusterGroupType,
+            LocationType,
+            RackType,
+            RegionType,
+            SiteType,
+            SiteGroupType,
+        )
+
+    @classmethod
+    def resolve_type(cls, instance, info):
+        if type(instance) == Cluster:
+            return ClusterType
+        if type(instance) == ClusterGroup:
+            return ClusterGroupType
+        if type(instance) == Location:
+            return LocationType
+        if type(instance) == Rack:
+            return RackType
+        if type(instance) == Region:
+            return RegionType
+        if type(instance) == Site:
+            return SiteType
+        if type(instance) == SiteGroup:
+            return SiteGroupType

netbox/ipam/graphql/types.py

@@ -55,10 +55,11 @@ class FHRPGroupType(NetBoxObjectType):
 
 
 class FHRPGroupAssignmentType(BaseObjectType):
+    interface = graphene.Field('ipam.graphql.gfk_mixins.FHRPGroupInterfaceType')
 
     class Meta:
         model = models.FHRPGroupAssignment
-        fields = '__all__'
+        exclude = ('interface_type', 'interface_id')
         filterset_class = filtersets.FHRPGroupAssignmentFilterSet
 
 
@@ -142,10 +143,11 @@ class VLANType(NetBoxObjectType):
 
 
 class VLANGroupType(OrganizationalObjectType):
+    scope = graphene.Field('ipam.graphql.gfk_mixins.VLANGroupScopeType')
 
     class Meta:
         model = models.VLANGroup
-        fields = '__all__'
+        exclude = ('scope_type', 'scope_id')
         filterset_class = filtersets.VLANGroupFilterSet