schema.py

import graphene
from graphene import relay
from graphene_sqlalchemy import SQLAlchemyObjectType, SQLAlchemyConnectionField
from database.models import User as UserModel, Application as ApplicationModel, Role as RoleModel
from database.database import db_session, get_user_datastore
from flask_security.utils import hash_password, verify_and_update_password

user_datastore = get_user_datastore()


class User(SQLAlchemyObjectType):
    class Meta:
        model = UserModel
        exclude_fields = ("password")
        interfaces = (relay.Node, )


class UserMutation(graphene.Mutation):
    class Arguments:
        username = graphene.String(required=True)
        password = graphene.String(required=True)
        email = graphene.String(required=True)

    user = graphene.Field(User)

    def mutate(self, info, username, password, email):
        user = user_datastore.create_user(username=username,
                                          email=email,
                                          password=hash_password(password))
        user_datastore.commit()
        return UserMutation(user=user)


class EditUser(graphene.Mutation):
    class Arguments:
        username = graphene.String(required=True)
        password = graphene.String()
        email = graphene.String()

    ok = graphene.Boolean()

    def mutate(self, info, username, password=None, email=None):
        user = user_datastore.find_user(username=username)
        ok = False
        if user is not None:
            if email is not None:
                user.email = email
            if password is not None:
                user.password = hash_password(password)
            user_datastore.commit()
            ok = True
        return EditUser(ok=ok)


class Application(SQLAlchemyObjectType):
    class Meta:
        model = ApplicationModel
        interfaces = (relay.Node, )


class CreateApplication(graphene.Mutation):
    class Arguments:
        name = graphene.String(required=True)

    application = graphene.Field(Application)

    def mutate(self, info, name):
        application = ApplicationModel(name=name)
        user_datastore.db.session().add(application)
        user_datastore.commit()
        return CreateApplication(application=application)


class Role(SQLAlchemyObjectType):
    class Meta:
        model = RoleModel
        interfaces = (relay.Node, )


class CreateRole(graphene.Mutation):
    class Arguments:
        name = graphene.String(required=True)
        description = graphene.String(required=False)

    role = graphene.Field(Role)

    def mutate(self, info, name, description=None):
        role = user_datastore.create_role(name=name,
                                          description=description)
        user_datastore.commit()
        return CreateRole(role=role)


class AddRoleToUser(graphene.Mutation):
    class Arguments:
        username = graphene.String(required=True)
        role = graphene.String(required=True)

    user = graphene.Field(User)
    role = graphene.Field(Role)

    def mutate(self, info, username, role):
        user = user_datastore.db.session().query(UserModel).filter_by(username=username).first()
        role = user_datastore.find_role(role)
        user_datastore.add_role_to_user(user, role)
        user_datastore.commit()
        return AddRoleToUser(user=user,role=role)


class RemoveRoleFromUser(graphene.Mutation):
    class Arguments:
        username = graphene.String(required=True)
        role = graphene.String(required=True)

    ok = graphene.Boolean()

    def mutate(self, info, username, role):
        ok = False
        user = user_datastore.db.session().query(UserModel).filter_by(username=username).first()
        if user is not None:
            user_datastore.remove_role_from_user(user, role)
            user_datastore.commit()
            ok = True
        return RemoveRoleFromUser(ok=ok)


class AddApplicationToUser(graphene.Mutation):
    class Arguments:
        username = graphene.String(required=True)
        application = graphene.String(required=True)

    user = graphene.Field(User)
    application = graphene.Field(Application)

    def mutate(self, info, username, application):
        user = user_datastore.db.session().query(UserModel).filter_by(username=username).first()
        application = user_datastore.db.session().query(ApplicationModel).filter_by(name=application).first()
        user.applications.append(application)
        user_datastore.commit()
        return AddApplicationToUser(user=user,application=application)


class DeleteUser(graphene.Mutation):
    class Arguments:
        username = graphene.String(required=True)

    ok = graphene.Boolean()

    def mutate(self, info, username):
        ok = False
        user = user_datastore.db.session().query(UserModel).filter_by(username=username).first()
        if user:
            user_datastore.delete_user(user)
            user_datastore.commit()
            ok = True
        return DeleteUser(ok=ok)


class Query(graphene.ObjectType):
    node = relay.Node.Field()
    all_users = SQLAlchemyConnectionField(User)
    all_roles = SQLAlchemyConnectionField(Role)
    get_user = graphene.Field(User,username=graphene.String())
    verify_password = graphene.Field(graphene.Boolean, username=graphene.String(), password=graphene.String())

    def resolve_get_user(root, info, username):
        query = User.get_query(info)
        return query.filter(UserModel.username == username).first()

    def resolve_verify_password(root, info, username, password):
        query = User.get_query(info)
        user = query.filter(UserModel.username == username).first()
        return verify_and_update_password(password, user)


class Mutation(graphene.ObjectType):
    create_user = UserMutation.Field()
    create_role = CreateRole.Field()
    create_application = CreateApplication.Field()
    add_role_to_user = AddRoleToUser.Field()
    remove_role_from_user = RemoveRoleFromUser.Field()
    add_application_to_user = AddApplicationToUser.Field()
    delete_user = DeleteUser.Field()
    edit_user = EditUser.Field()


schema = graphene.Schema(query=Query, types=[User, Role, Application], mutation=Mutation)