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 UserType class RoomEntryType(DjangoObjectType): pk = graphene.Int() author = UserType() class Meta: model = RoomEntry filter_fields = ['slug'] interfaces = (relay.Node,) def resolve_pk(self, *args, **kwargs): return self.id class RoomType(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(RoomType, slug=graphene.String(), id=graphene.Int(), appearance=graphene.String()) room_entry = relay.Node.Field(RoomEntryType) all_rooms = DjangoFilterConnectionField(RoomType) all_room_entries = DjangoFilterConnectionField(RoomEntryType) 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