9817 add graphql ipaddressassignment assigned_object

netbox/ipam/graphql/gfk_mixins.py

@@ -1,12 +1,30 @@
 import graphene
 from dcim.graphql.types import InterfaceType
 from dcim.models import Interface
-from ipam.graphql.types import VLANType
+from ipam.graphql.types import FHRPGroupType, VLANType
-from ipam.models import VLAN
+from ipam.models import FHRPGroup, VLAN
 from virtualization.graphql.types import VMInterfaceType
 from virtualization.models import VMInterface
 
 
+class IPAddressAssignmentType(graphene.Union):
+    class Meta:
+        types = (
+            InterfaceType,
+            FHRPGroupType,
+            VMInterfaceType,
+        )
+
+    @classmethod
+    def resolve_type(cls, instance, info):
+        if type(instance) == Interface:
+            return InterfaceType
+        if type(instance) == FHRPGroup:
+            return FHRPGroupType
+        if type(instance) == VMInterface:
+            return VMInterfaceType
+
+
 class L2VPNAssignmentType(graphene.Union):
     class Meta:
         types = (

netbox/ipam/graphql/types.py

@@ -63,10 +63,11 @@ class FHRPGroupAssignmentType(BaseObjectType):
 
 
 class IPAddressType(NetBoxObjectType):
+    assigned_object = graphene.Field('ipam.graphql.gfk_mixins.IPAddressAssignmentType')
 
     class Meta:
         model = models.IPAddress
-        fields = '__all__'
+        exclude = ('assigned_object_type', 'assigned_object_id')
         filterset_class = filtersets.IPAddressFilterSet
 
     def resolve_role(self, info):