import graphene from graphene import relay from graphene_django import DjangoObjectType from graphene_django.filter import DjangoFilterConnectionField from rooms.models import Room, RoomEntry from user.schema import UserNode class RoomEntryNode(DjangoObjectType): pk = graphene.Int() author = UserNode() class Meta: model = RoomEntry filter_fields = ['slug'] interfaces = (relay.Node,) def resolve_pk(self, *args, **kwargs): return self.id class RoomNode(DjangoObjectType): pk = graphene.Int() entry_count = graphene.Int() class Meta: model = Room filter_fields = ['slug', 'appearance'] interfaces = (relay.Node,) def resolve_pk(self, *args, **kwargs): return self.id def resolve_entry_count(self, *args, **kwargs): return self.room_entries.count() class RoomsQuery(object): room = graphene.Field(RoomNode, slug=graphene.String(), id=graphene.Int(), appearance=graphene.String()) room_entry = relay.Node.Field(RoomEntryNode) all_rooms = DjangoFilterConnectionField(RoomNode) all_room_entries = DjangoFilterConnectionField(RoomEntryNode) def resolve_room(self, info, **kwargs): slug = kwargs.get('slug') room_id = kwargs.get('id') if room_id is not None: return Room.objects.get(pk=room_id) if slug is not None: return Room.objects.get(slug=slug) return None