Merge pull request rmosolgo#354 from rmosolgo/instrumentation

Instrumentation

Author: rmosolgo

guides/schema/instrumentation.md

@@ -0,0 +1,71 @@
+---
+title: Schema — Instrumentation
+---
+
+Instrumentation provides hooks for inserting custom code around field resolution and query execution.
+
+## Field Instrumentation
+
+Field instrumentation can be attached during schema definition:
+
+```ruby
+MySchema = GraphQL::Schema.define do
+  instrument(:field, FieldTimerInstrumentation.new)
+end
+```
+
+The instrumenter is an object which responds to `#instrument(type, field)`. `#instrument` should return a `GraphQL::Field` instance which will be used during query execution. `#instrument` is called with each type-field pair for _all_ Object types and Interface types in your schema.
+
+Here's an example field instrumenter:
+
+```ruby
+class FieldTimerInstrumentation
+  # If a field was flagged to be timed,
+  # wrap its resolve proc with a timer.
+  def instrument(type, field)
+    if field.metadata[:timed]
+      old_resolve_proc = field.resolve_proc
+      new_resolve_proc = ->(obj, args, ctx) {
+        Rails.logger.info("#{type.name}.#{field.name} START: #{Time.now.to_i}")
+        old_resolve_proc.call(obj, args, ctx)
+        Rails.logger.info("#{type.name}.#{field.name} END: #{Time.now.to_i}")
+      }
+
+      # Return a copy of `field`, with a new resolve proc
+      field.redefine do
+        resolve(new_resolve_proc)
+      end
+    end
+  end
+end
+```
+
+It can be attached as shown above. You can use `redefine { ... }` to make a shallow copy of the  {{ "GraphQL::Field" | api_doc }} and extend its definition.
+
+## Query Instrumentation
+
+
+Query instrumentation can be attached during schema definition:
+
+```ruby
+MySchema = GraphQL::Schema.define do
+  instrument(:query, QueryTimerInstrumentation)
+end
+```
+
+The instrumenter must implement `#before_query(query)` and `#after_query(query)`. The return values of these methods are not used. They receive the {{ "GraphQL::Query" | api_doc }} instance.
+
+```ruby
+module MyQueryInstrumentation
+  module_function
+
+  # Log the time of the query
+  def before_query(query)
+    Rails.logger.info("Query begin: #{Time.now.to_i}")
+  end
+
+  def after_query(query)
+    Rails.logger.info("Query end: #{Time.now.to_i}")
+  end
+end
+```

lib/graphql/query.rb

@@ -91,16 +91,9 @@ def initialize(schema, query_string = nil, document: nil, context: nil, variable
     # Get the result for this query, executing it once
     def result
       @result ||= begin
-        if !valid?
-          all_errors = validation_errors + analysis_errors
-          if all_errors.any?
-            { "errors" => all_errors.map(&:to_h) }
-          else
-            nil
-          end
-        else
-          Executor.new(self).result
-        end
+        instrumenters = @schema.instrumenters[:query]
+        execution_call = ExecutionCall.new(self, instrumenters)
+        execution_call.call
       end
     end
 
@@ -209,5 +202,37 @@ def find_operation(operations, operation_name)
         operations[operation_name]
       end
     end
+
+    class ExecutionCall
+      def initialize(query, instrumenters)
+        @query = query
+        @instrumenters = instrumenters
+      end
+
+      # Check if the query is valid, and if it is,
+      # execute it, calling instrumenters along the way
+      # @return [Hash] The GraphQL response
+      def call
+        @instrumenters.each { |i| i.before_query(@query) }
+        result = get_result
+        @instrumenters.each { |i| i.after_query(@query) }
+        result
+      end
+
+      private
+
+      def get_result
+        if !@query.valid?
+          all_errors = @query.validation_errors + @query.analysis_errors
+          if all_errors.any?
+            { "errors" => all_errors.map(&:to_h) }
+          else
+            nil
+          end
+        else
+          Executor.new(@query).result
+        end
+      end
+    end
   end
 end

lib/graphql/query/variables.rb

@@ -2,6 +2,8 @@ module GraphQL
   class Query
     # Read-only access to query variables, applying default values if needed.
     class Variables
+      extend Forwardable
+
       # @return [Array<GraphQL::Query::VariableValidationError>]  Any errors encountered when parsing the provided variables and literal values
       attr_reader :errors
 
@@ -20,6 +22,8 @@ def [](key)
         @storage.fetch(key)
       end
 
+      def_delegators :@storage, :length
+
       private
 
       # Find the right value for this variable:

lib/graphql/schema.rb

@@ -1,5 +1,6 @@
 require "graphql/schema/catchall_middleware"
 require "graphql/schema/invalid_type_error"
+require "graphql/schema/instrumented_field_map"
 require "graphql/schema/middleware_chain"
 require "graphql/schema/possible_types"
 require "graphql/schema/rescue_middleware"
@@ -55,6 +56,7 @@ class Schema
       :orphan_types, :resolve_type,
       :object_from_id, :id_from_object,
       directives: ->(schema, directives) { schema.directives = directives.reduce({}) { |m, d| m[d.name] = d; m  }},
+      instrument: -> (schema, type, instrumenter) { schema.instrumenters[type] << instrumenter },
       query_analyzer: ->(schema, analyzer) { schema.query_analyzers << analyzer },
       middleware: ->(schema, middleware) { schema.middleware << middleware },
       rescue_from: ->(schema, err_class, &block) { schema.rescue_from(err_class, &block)}
@@ -64,7 +66,7 @@ class Schema
       :query_execution_strategy, :mutation_execution_strategy, :subscription_execution_strategy,
       :max_depth, :max_complexity,
       :orphan_types, :directives,
-      :query_analyzers, :middleware
+      :query_analyzers, :middleware, :instrumenters
 
     BUILT_IN_TYPES = Hash[[INT_TYPE, STRING_TYPE, FLOAT_TYPE, BOOLEAN_TYPE, ID_TYPE].map{ |type| [type.name, type] }]
     DIRECTIVES = [GraphQL::Directive::IncludeDirective, GraphQL::Directive::SkipDirective, GraphQL::Directive::DeprecatedDirective]
@@ -89,6 +91,7 @@ def initialize
       @resolve_type_proc = nil
       @object_from_id_proc = nil
       @id_from_object_proc = nil
+      @instrumenters = Hash.new { |h, k| h[k] = [] }
       # Default to the built-in execution strategy:
       @query_execution_strategy = GraphQL::Query::SerialExecution
       @mutation_execution_strategy = GraphQL::Query::SerialExecution
@@ -107,7 +110,23 @@ def remove_handler(*args, &block)
 
     def define(**kwargs, &block)
       super
-      types
+      all_types = orphan_types + [query, mutation, subscription, GraphQL::Introspection::SchemaType]
+      @types = GraphQL::Schema::ReduceTypes.reduce(all_types.compact)
+
+      @instrumented_field_map = InstrumentedFieldMap.new(self)
+      field_instrumenters = @instrumenters[:field]
+      types.each do |type_name, type|
+        if type.kind.fields?
+          type.all_fields.each do |field_defn|
+
+            instrumented_field_defn = field_instrumenters.reduce(field_defn) do |defn, inst|
+              inst.instrument(type, defn)
+            end
+
+            @instrumented_field_map.set(type.name, field_defn.name, instrumented_field_defn)
+          end
+        end
+      end
       # Assert that all necessary configs are present:
       validation_error = Validation.validate(self)
       validation_error && raise(NotImplementedError, validation_error)
@@ -117,13 +136,7 @@ def define(**kwargs, &block)
 
     # @see [GraphQL::Schema::Warden] Restricted access to members of a schema
     # @return [GraphQL::Schema::TypeMap] `{ name => type }` pairs of types in this schema
-    def types
-      @types ||= begin
-        ensure_defined
-        all_types = orphan_types + [query, mutation, subscription, GraphQL::Introspection::SchemaType]
-        GraphQL::Schema::ReduceTypes.reduce(all_types.compact)
-      end
-    end
+    attr_reader :types
 
     # Execute a query on itself.
     # See {Query#initialize} for arguments.
@@ -139,8 +152,7 @@ def execute(*args)
     # @return [GraphQL::Field, nil] The field named `field_name` on `parent_type`
     def get_field(parent_type, field_name)
       ensure_defined
-
-      defined_field = parent_type.get_field(field_name)
+      defined_field = @instrumented_field_map.get(parent_type.name, field_name)
       if defined_field
         defined_field
       elsif field_name == "__typename"

lib/graphql/schema/instrumented_field_map.rb

@@ -0,0 +1,23 @@
+module GraphQL
+  class Schema
+    # A two-level map with fields as the last values.
+    # The first level is type names, which point to a second map.
+    # The second level is field names, which point to fields.
+    #
+    # The catch is, the fields in this map _may_ have been modified by being instrumented.
+    class InstrumentedFieldMap
+      def initialize(schema)
+        @storage = Hash.new { |h, k| h[k] = {} }
+      end
+
+      def set(type_name, field_name, field)
+        @storage[type_name][field_name] = field
+      end
+
+      def get(type_name, field_name)
+        type = @storage[type_name]
+        type && type[field_name]
+      end
+    end
+  end
+end

spec/graphql/schema_spec.rb

@@ -182,4 +182,75 @@
       assert_equal GraphQL::Schema::Printer.print_schema(schema), GraphQL::Schema::Printer.print_schema(built_schema)
     end
   end
+
+  describe "#instrument" do
+    class MultiplyInstrumenter
+      def initialize(multiplier)
+        @multiplier = multiplier
+      end
+
+      def instrument(type_defn, field_defn)
+        if type_defn.name == "Query" && field_defn.name == "int"
+          prev_proc = field_defn.resolve_proc
+          new_resolve_proc = ->(obj, args, ctx) {
+            inner_value = prev_proc.call(obj, args, ctx)
+            inner_value * @multiplier
+          }
+
+          field_defn.redefine do
+            resolve(new_resolve_proc)
+          end
+        else
+          field_defn
+        end
+      end
+    end
+
+    class VariableCountInstrumenter
+      attr_reader :counts
+      def initialize
+        @counts = []
+      end
+
+      def before_query(query)
+        @counts << query.variables.length
+      end
+
+      def after_query(query)
+      end
+    end
+
+    let(:variable_counter) {
+      VariableCountInstrumenter.new
+    }
+    let(:query_type) {
+      GraphQL::ObjectType.define do
+        name "Query"
+        field :int, types.Int do
+          argument :value, types.Int
+          resolve -> (obj, args, ctx) { args[:value] }
+        end
+      end
+    }
+
+    let(:schema) {
+      spec = self
+      GraphQL::Schema.define do
+        query(spec.query_type)
+        instrument(:field, MultiplyInstrumenter.new(3))
+        instrument(:query, spec.variable_counter)
+      end
+    }
+
+    it "can modify field definitions" do
+      res = schema.execute(" { int(value: 2) } ")
+      assert_equal 6, res["data"]["int"]
+    end
+
+    it "can wrap query execution" do
+      schema.execute("query getInt($val: Int = 5){ int(value: $val) } ")
+      schema.execute("query getInt($val: Int = 5, $val2: Int = 3){ int(value: $val) int2: int(value: $val2) } ")
+      assert_equal [1, 2], variable_counter.counts
+    end
+  end
 end