diff --git a/graphene/core/schema.py b/graphene/core/schema.py index dd0fce9e..f8f26dc1 100644 --- a/graphene/core/schema.py +++ b/graphene/core/schema.py @@ -119,24 +119,9 @@ class Schema(object): return self._types_names def execute(self, request='', root=None, args=None, **kwargs): - executor = kwargs - executor['root'] = root - executor['args'] = args - executor['schema'] = self.schema - executor['request'] = request - contexts = [] - for plugin in self.plugins: - if not hasattr(plugin, 'context_execution'): - continue - context = plugin.context_execution(executor) - executor = context.__enter__() - contexts.append((context, executor)) - result = self.executor.execute( - **executor - ) - for context, value in contexts[::-1]: - context.__exit__(None, None, None) - return result + kwargs = dict(kwargs, request=request, root=root, args=args, schema=self.schema) + with self.plugins.context_execution(**kwargs) as execute_kwargs: + return self.executor.execute(**execute_kwargs) def introspect(self): return self.execute(introspection_query).data diff --git a/graphene/plugins/base.py b/graphene/plugins/base.py index 8ab5668a..2347beba 100644 --- a/graphene/plugins/base.py +++ b/graphene/plugins/base.py @@ -1,3 +1,4 @@ +from contextlib import contextmanager from functools import partial, reduce @@ -38,3 +39,15 @@ class PluginManager(object): def __contains__(self, name): return name in self.PLUGIN_FUNCTIONS + + @contextmanager + def context_execution(self, **executor): + contexts = [] + functions = self.get_plugin_functions('context_execution') + for f in functions: + context = f(executor) + executor = context.__enter__() + contexts.append((context, executor)) + yield executor + for context, value in contexts[::-1]: + context.__exit__(None, None, None)