lib/graphql/execution/multiplex.rb in graphql-1.6.2 vs lib/graphql/execution/multiplex.rb in graphql-1.6.3

@@ -44,43 +44,29 @@
         # @param queries [Array<GraphQL::Query>]
         # @param context [Hash]
         # @param max_complexity [Integer]
         # @return [Array<Hash>] One result per query
         def run_queries(schema, queries, context: {}, max_complexity: nil)
-          has_custom_strategy = schema.query_execution_strategy || schema.mutation_execution_strategy || schema.subscription_execution_strategy
-          if has_custom_strategy
-            if queries.length == 1
-              return [run_one_legacy(schema, queries.first)]
-            else
+          if has_custom_strategy?(schema)
+            if queries.length != 1
               raise ArgumentError, "Multiplexing doesn't support custom execution strategies, run one query at a time instead"
+            else
+              with_instrumentation(schema, queries, context: context, max_complexity: max_complexity) do
+                [run_one_legacy(schema, queries.first)]
+              end
             end
           else
-            run_as_multiplex(schema, queries, context: context, max_complexity: max_complexity)
+            with_instrumentation(schema, queries, context: context, max_complexity: max_complexity) do
+              run_as_multiplex(queries)
+            end
           end
         end
 
         private
 
-        def run_as_multiplex(schema, queries, context:, max_complexity:)
-          query_instrumenters = schema.instrumenters[:query]
-          multiplex_instrumenters = schema.instrumenters[:multiplex]
-          multiplex = self.new(schema: schema, queries: queries, context: context)
-
-          # First, run multiplex instrumentation, then query instrumentation for each query
-          multiplex_instrumenters.each { |i| i.before_multiplex(multiplex) }
-          queries.each do |query|
-            query_instrumenters.each { |i| i.before_query(query) }
-          end
-
-          multiplex_analyzers = schema.multiplex_analyzers
-          if max_complexity ||= schema.max_complexity
-            multiplex_analyzers += [GraphQL::Analysis::MaxQueryComplexity.new(max_complexity)]
-          end
-
-          GraphQL::Analysis.analyze_multiplex(multiplex, multiplex_analyzers)
-
-          # Then, do as much eager evaluation of the query as possible
+        def run_as_multiplex(queries)
+          # Do as much eager evaluation of the query as possible
           results = queries.map do |query|
             begin_query(query)
           end
 
           # Then, work through lazy results in a breadth-first way
@@ -89,17 +75,10 @@
           # Then, find all errors and assign the result to the query object
           results.each_with_index.map do |data_result, idx|
             query = queries[idx]
             finish_query(data_result, query)
           end
-        ensure
-          # Finally, run teardown instrumentation for each query + the multiplex
-          # Use `reverse_each` so instrumenters are treated like a stack
-          queries.each do |query|
-            query_instrumenters.reverse_each { |i| i.after_query(query) }
-          end
-          multiplex_instrumenters.reverse_each { |i| i.after_multiplex(multiplex) }
         end
 
         # @param query [GraphQL::Query]
         # @return [Hash] The initial result (may not be finished if there are lazy values)
         def begin_query(query)
@@ -147,23 +126,57 @@
           end
         end
 
         # use the old `query_execution_strategy` etc to run this query
         def run_one_legacy(schema, query)
-          instrumenters = schema.instrumenters[:query]
-          instrumenters.each { |i| i.before_query(query) }
           query.result = if !query.valid?
             all_errors = query.validation_errors + query.analysis_errors + query.context.errors
             if all_errors.any?
               { "errors" => all_errors.map(&:to_h) }
             else
               nil
             end
           else
             GraphQL::Query::Executor.new(query).result
           end
+        end
+
+        def has_custom_strategy?(schema)
+          schema.query_execution_strategy != GraphQL::Execution::Execute ||
+            schema.mutation_execution_strategy != GraphQL::Execution::Execute ||
+            schema.subscription_execution_strategy != GraphQL::Execution::Execute
+        end
+
+        # Apply multiplex & query instrumentation to `queries`.
+        #
+        # It yields when the queries should be executed, then runs teardown.
+        def with_instrumentation(schema, queries, context:, max_complexity:)
+          query_instrumenters = schema.instrumenters[:query]
+          multiplex_instrumenters = schema.instrumenters[:multiplex]
+          multiplex = self.new(schema: schema, queries: queries, context: context)
+
+          # First, run multiplex instrumentation, then query instrumentation for each query
+          multiplex_instrumenters.each { |i| i.before_multiplex(multiplex) }
+          queries.each do |query|
+            query_instrumenters.each { |i| i.before_query(query) }
+          end
+
+          multiplex_analyzers = schema.multiplex_analyzers
+          if max_complexity ||= schema.max_complexity
+            multiplex_analyzers += [GraphQL::Analysis::MaxQueryComplexity.new(max_complexity)]
+          end
+
+          GraphQL::Analysis.analyze_multiplex(multiplex, multiplex_analyzers)
+
+          # Let them be executed
+          yield
         ensure
-          instrumenters.reverse_each { |i| i.after_query(query) }
+          # Finally, run teardown instrumentation for each query + the multiplex
+          # Use `reverse_each` so instrumenters are treated like a stack
+          queries.each do |query|
+            query_instrumenters.reverse_each { |i| i.after_query(query) }
+          end
+          multiplex_instrumenters.reverse_each { |i| i.after_multiplex(multiplex) }
         end
       end
     end
   end
 end