diff --git a/netbox/ipam/graphql/gfk_mixins.py b/netbox/ipam/graphql/gfk_mixins.py new file mode 100644 index 000000000..055be911e --- /dev/null +++ b/netbox/ipam/graphql/gfk_mixins.py @@ -0,0 +1,25 @@ +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 virtualization.graphql.types import VMInterfaceType +from virtualization.models import VMInterface + + +class L2VPNAssignmentType(graphene.Union): + class Meta: + types = ( + InterfaceType, + VLANType, + VMInterfaceType, + ) + + @classmethod + def resolve_type(cls, instance, info): + if type(instance) == Interface: + return InterfaceType + if type(instance) == VLAN: + return VLANType + if type(instance) == VMInterface: + return VMInterfaceType diff --git a/netbox/ipam/graphql/types.py b/netbox/ipam/graphql/types.py index 5af2ca72a..649e80c03 100644 --- a/netbox/ipam/graphql/types.py +++ b/netbox/ipam/graphql/types.py @@ -1,5 +1,6 @@ import graphene +from graphene_django import DjangoObjectType from ipam import filtersets, models from netbox.graphql.scalars import BigInt from netbox.graphql.types import BaseObjectType, OrganizationalObjectType, NetBoxObjectType @@ -163,7 +164,9 @@ class L2VPNType(NetBoxObjectType): class L2VPNTerminationType(NetBoxObjectType): + assigned_object = graphene.Field('ipam.graphql.gfk_mixins.L2VPNAssignmentType') + class Meta: model = models.L2VPNTermination - fields = '__all__' + exclude = ('assigned_object_type', 'assigned_object_id') filtersets_class = filtersets.L2VPNTerminationFilterSet