lib/graphql/execution/multiplex.rb in graphql-1.7.14 vs lib/graphql/execution/multiplex.rb in graphql-1.8.0.pre1

@@ -58,16 +58,16 @@
           multiplex.trace("execute_multiplex", { multiplex: multiplex }) do
             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
-                instrument_and_analyze(multiplex, max_complexity: max_complexity) do
+                with_instrumentation(multiplex, max_complexity: max_complexity) do
                   [run_one_legacy(schema, queries.first)]
                 end
               end
             else
-              instrument_and_analyze(multiplex, max_complexity: max_complexity) do
+              with_instrumentation(multiplex, max_complexity: max_complexity) do
                 run_as_multiplex(multiplex)
               end
             end
           end
         end
@@ -91,11 +91,11 @@
             # Get the Query::Result, not the Hash
             query.result
           end
         rescue StandardError
           # Assign values here so that the query's `@executed` becomes true
-          queries.map { |q| q.result_values ||= {} }
+          queries.map { |q| q.result_values ||= nil }
           raise
         end
 
         # @param query [GraphQL::Query]
         # @return [Hash] The initial result (may not be finished if there are lazy values)
@@ -160,20 +160,37 @@
         end
 
         # Apply multiplex & query instrumentation to `queries`.
         #
         # It yields when the queries should be executed, then runs teardown.
-        def instrument_and_analyze(multiplex, max_complexity:)
-          GraphQL::Execution::Instrumentation.apply_instrumenters(multiplex) do
-            schema = multiplex.schema
-            multiplex_analyzers = schema.multiplex_analyzers
-            if max_complexity
-              multiplex_analyzers += [GraphQL::Analysis::MaxQueryComplexity.new(max_complexity)]
-            end
+        def with_instrumentation(multiplex, max_complexity:)
+          schema = multiplex.schema
+          queries = multiplex.queries
+          query_instrumenters = schema.instrumenters[:query]
+          multiplex_instrumenters = schema.instrumenters[:multiplex]
 
-            GraphQL::Analysis.analyze_multiplex(multiplex, multiplex_analyzers)
-            yield
+          # 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
+            multiplex_analyzers += [GraphQL::Analysis::MaxQueryComplexity.new(max_complexity)]
+          end
+
+          GraphQL::Analysis.analyze_multiplex(multiplex, multiplex_analyzers)
+
+          # Let them be executed
+          yield
+        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
       end
     end
   end
 end