diff --git a/netbox/ipam/graphql/gfk_mixins.py b/netbox/ipam/graphql/gfk_mixins.py index 055be911e..c9ca617d2 100644 --- a/netbox/ipam/graphql/gfk_mixins.py +++ b/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.models import VLAN +from ipam.graphql.types import FHRPGroupType, VLANType +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 = ( diff --git a/netbox/ipam/graphql/types.py b/netbox/ipam/graphql/types.py index 649e80c03..4514c3acb 100644 --- a/netbox/ipam/graphql/types.py +++ b/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):