graphene_sqlalchemy_mutations.py

import graphene
from sqlalchemy.inspection import inspect

from graphene.types.utils import yank_fields_from_attrs
from graphene.types.objecttype import ObjectTypeOptions
from graphene_sqlalchemy import get_session
from graphene_sqlalchemy.registry import get_global_registry
from graphene_sqlalchemy.types import construct_fields


class SQLAlchemyMutationOptions(ObjectTypeOptions):
    model = None  # type: Model
    model_inspect = None  # type: Mapper
    create = False  # type: Boolean
    arguments = None  # type: Dict[str, Argument]
    output = None  # type: Type[ObjectType]
    resolver = None  # type: Callable


class SQLAlchemyMutation(graphene.Mutation):
    @classmethod
    def __init_subclass_with_meta__(cls, model=None, create=False, registry=None, only_fields=(), exclude_fields=None,
                                    **options):
        meta = SQLAlchemyMutationOptions(cls)
        meta.create = create
        meta.model = model

        model_inspect = inspect(model)
        meta.model_inspect = model_inspect

        if not registry:
            registry = get_global_registry()

        if not isinstance(exclude_fields, list):
            if exclude_fields:
                exclude_fields = list(exclude_fields)
            else:
                exclude_fields = []

        if meta.create:
            for primary_key_column in model_inspect.primary_key:
                exclude_fields.append(primary_key_column.name)

        for relationship in model_inspect.relationships:
            exclude_fields.append(relationship.key)

        arguments = yank_fields_from_attrs(
            construct_fields(model, registry, only_fields, exclude_fields),
            _as=graphene.Argument,
        )

        super(SQLAlchemyMutation, cls).__init_subclass_with_meta__(_meta=meta, arguments=arguments, **options)

    @classmethod
    def mutate(cls, self, info, **kwargs):
        session = get_session(info.context)

        meta = cls._meta
        model = None

        if meta.create:
            model = meta.model()
            session.add(model)
        else:
            query = session.query(meta.model)
            for primary_key_column in meta.model_inspect.primary_key:
                query = query.filter(getattr(meta.model, primary_key_column.key) == kwargs[primary_key_column.name])
            model = query.first()

        for key, value in kwargs.items():
            setattr(model, key, value)

        session.commit()

        return model

    @classmethod
    def Field(cls, *args, **kwargs):
        return graphene.Field(
            cls._meta.output, args=cls._meta.arguments, resolver=cls._meta.resolver
        )

graphql-python/graphene-sqlalchemy#29