graphql-stitching 0.0.1

data/lib/graphql/stitching/executor.rb

@@ -0,0 +1,254 @@
+# frozen_string_literal: true
+
+require "json"
+
+module GraphQL
+  module Stitching
+    class Executor
+
+      class RootSource < GraphQL::Dataloader::Source
+        def initialize(executor)
+          @executor = executor
+        end
+
+        def fetch(ops)
+          op = ops.first # There should only ever be one per location at a time
+
+          query_document = build_query(op)
+          query_variables = @executor.variables.slice(*op["variables"].keys)
+          result = @executor.supergraph.execute_at_location(op["location"], query_document, query_variables)
+          @executor.query_count += 1
+
+          @executor.data.merge!(result["data"]) if result["data"]
+          if result["errors"]&.any?
+            result["errors"].each { _1.delete("locations") }
+            @executor.errors.concat(result["errors"])
+          end
+          op["key"]
+        end
+
+        def build_query(op)
+          if op["variables"].any?
+            variable_defs = op["variables"].map { |k, v| "$#{k}:#{v}" }.join(",")
+            "#{op["operation_type"]}(#{variable_defs})#{op["selections"]}"
+          else
+            "#{op["operation_type"]}#{op["selections"]}"
+          end
+        end
+      end
+
+      class BoundarySource < GraphQL::Dataloader::Source
+        def initialize(executor, location)
+          @executor = executor
+          @location = location
+        end
+
+        def fetch(ops)
+          origin_sets_by_operation = ops.each_with_object({}) do |op, memo|
+            origin_set = op["insertion_path"].reduce([@executor.data]) do |set, path_segment|
+              mapped = set.flat_map { |obj| obj && obj[path_segment] }
+              mapped.compact!
+              mapped
+            end
+
+            if op["type_condition"]
+              # operations planned around unused fragment conditions should not trigger requests
+              origin_set.select! { _1["_STITCH_typename"] == op["type_condition"] }
+            end
+
+            memo[op] = origin_set if origin_set.any?
+          end
+
+          if origin_sets_by_operation.any?
+            query_document, variable_names = build_query(origin_sets_by_operation)
+            variables = @executor.variables.slice(*variable_names)
+            raw_result = @executor.supergraph.execute_at_location(@location, query_document, variables)
+            @executor.query_count += 1
+
+            merge_results!(origin_sets_by_operation, raw_result.dig("data"))
+
+            errors = raw_result.dig("errors")
+            @executor.errors.concat(extract_errors!(origin_sets_by_operation, errors)) if errors&.any?
+          end
+
+          ops.map { origin_sets_by_operation[_1] ? _1["key"] : nil }
+        end
+
+        def build_query(origin_sets_by_operation)
+          variable_defs = {}
+          query_fields = origin_sets_by_operation.map.with_index do |(op, origin_set), batch_index|
+            variable_defs.merge!(op["variables"])
+            boundary = op["boundary"]
+            key_selection = "_STITCH_#{boundary["selection"]}"
+
+            if boundary["list"]
+              input = JSON.generate(origin_set.map { _1[key_selection] })
+              "_#{batch_index}_result: #{boundary["field"]}(#{boundary["arg"]}:#{input}) #{op["selections"]}"
+            else
+              origin_set.map.with_index do |origin_obj, index|
+                input = JSON.generate(origin_obj[key_selection])
+                "_#{batch_index}_#{index}_result: #{boundary["field"]}(#{boundary["arg"]}:#{input}) #{op["selections"]}"
+              end
+            end
+          end
+
+          query_document = if variable_defs.any?
+            query_variables = variable_defs.map { |k, v| "$#{k}:#{v}" }.join(",")
+            "query(#{query_variables}){ #{query_fields.join(" ")} }"
+          else
+            "query{ #{query_fields.join(" ")} }"
+          end
+
+          return query_document, variable_defs.keys
+        end
+
+        def merge_results!(origin_sets_by_operation, raw_result)
+          return unless raw_result
+
+          origin_sets_by_operation.each_with_index do |(op, origin_set), batch_index|
+            results = if op.dig("boundary", "list")
+              raw_result["_#{batch_index}_result"]
+            else
+              origin_set.map.with_index { |_, index| raw_result["_#{batch_index}_#{index}_result"] }
+            end
+
+            next unless results&.any?
+
+            origin_set.each_with_index do |origin_obj, index|
+              origin_obj.merge!(results[index]) if results[index]
+            end
+          end
+        end
+
+        # https://spec.graphql.org/June2018/#sec-Errors
+        def extract_errors!(origin_sets_by_operation, errors)
+          ops = origin_sets_by_operation.keys
+          origin_sets = origin_sets_by_operation.values
+          pathed_errors_by_op_index_and_object_id = {}
+
+          errors_result = errors.each_with_object([]) do |err, memo|
+            err.delete("locations")
+            path = err["path"]
+
+            if path && path.length > 0
+              result_alias = /^_(\d+)(?:_(\d+))?_result$/.match(path.first.to_s)
+
+              if result_alias
+                path = err["path"] = path[1..-1]
+
+                origin_obj = if result_alias[2]
+                  origin_sets.dig(result_alias[1].to_i, result_alias[2].to_i)
+                elsif path[0].is_a?(Integer) || /\d+/.match?(path[0].to_s)
+                  origin_sets.dig(result_alias[1].to_i, path.shift.to_i)
+                end
+
+                if origin_obj
+                  by_op_index = pathed_errors_by_op_index_and_object_id[result_alias[1].to_i] ||= {}
+                  by_object_id = by_op_index[origin_obj.object_id] ||= []
+                  by_object_id << err
+                  next
+                end
+              end
+            end
+
+            memo << err
+          end
+
+          if pathed_errors_by_op_index_and_object_id.any?
+            pathed_errors_by_op_index_and_object_id.each do |op_index, pathed_errors_by_object_id|
+              repath_errors!(pathed_errors_by_object_id, ops.dig(op_index, "insertion_path"))
+              errors_result.concat(pathed_errors_by_object_id.values)
+            end
+          end
+          errors_result.flatten!
+        end
+
+        private
+
+        # traverses forward through origin data, expanding arrays to follow all paths
+        # any errors found for an origin object_id have their path prefixed by the object path
+        def repath_errors!(pathed_errors_by_object_id, forward_path, current_path=[], root=@executor.data)
+          current_path << forward_path.first
+          forward_path = forward_path[1..-1]
+          scope = root[current_path.last]
+
+          if forward_path.any? && scope.is_a?(Array)
+            scope.each_with_index do |element, index|
+              inner_elements = element.is_a?(Array) ? element.flatten : [element]
+              inner_elements.each do |inner_element|
+                repath_errors!(pathed_errors_by_object_id, forward_path, [*current_path, index], inner_element)
+              end
+            end
+
+          elsif forward_path.any?
+            repath_errors!(pathed_errors_by_object_id, forward_path, [*current_path, index], scope)
+
+          elsif scope.is_a?(Array)
+            scope.each_with_index do |element, index|
+              inner_elements = element.is_a?(Array) ? element.flatten : [element]
+              inner_elements.each do |inner_element|
+                errors = pathed_errors_by_object_id[inner_element.object_id]
+                errors.each { _1["path"] = [*current_path, index, *_1["path"]] } if errors
+              end
+            end
+
+          else
+            errors = pathed_errors_by_object_id[scope.object_id]
+            errors.each { _1["path"] = [*current_path, *_1["path"]] } if errors
+          end
+        end
+      end
+
+      attr_reader :supergraph, :data, :errors, :variables
+      attr_accessor :query_count
+
+      def initialize(supergraph:, plan:, variables: {}, nonblocking: false)
+        @supergraph = supergraph
+        @variables = variables
+        @queue = plan["ops"]
+        @data = {}
+        @errors = []
+        @query_count = 0
+        @dataloader = GraphQL::Dataloader.new(nonblocking: nonblocking)
+      end
+
+      def perform(document=nil)
+        exec!
+
+        result = {}
+        result["data"] = @data if @data && @data.length > 0
+        result["errors"] = @errors if @errors.length > 0
+
+        result if document.nil?
+
+        GraphQL::Stitching::Shaper.new(supergraph: @supergraph, document: document, raw: result).perform!
+      end
+
+      private
+
+      def exec!(after_keys = [0])
+        @dataloader.append_job do
+          requests = @queue
+            .select { after_keys.include?(_1["after_key"]) }
+            .group_by { _1["location"] }
+            .map do |location, ops|
+              if ops.first["after_key"].zero?
+                @dataloader.with(RootSource, self).request_all(ops)
+              else
+                @dataloader.with(BoundarySource, self, location).request_all(ops)
+              end
+            end
+
+          requests.each(&method(:exec_request))
+        end
+        @dataloader.run
+      end
+
+      def exec_request(request)
+        next_keys = request.load
+        next_keys.compact!
+        exec!(next_keys) if next_keys.any?
+      end
+    end
+  end
+end

data/lib/graphql/stitching/gateway.rb

@@ -0,0 +1,120 @@
+# frozen_string_literal: true
+
+require "json"
+
+module GraphQL
+  module Stitching
+    class Gateway
+      class GatewayError < StitchingError; end
+
+      EMPTY_CONTEXT = {}.freeze
+
+      attr_reader :supergraph
+
+      def initialize(locations: nil, supergraph: nil)
+        @supergraph = if locations && supergraph
+          raise GatewayError, "Cannot provide both locations and a supergraph."
+        elsif supergraph && !supergraph.is_a?(Supergraph)
+          raise GatewayError, "Provided supergraph must be a GraphQL::Stitching::Supergraph instance."
+        elsif supergraph
+          supergraph
+        elsif locations
+          build_supergraph_from_locations_config(locations)
+        else
+          raise GatewayError, "No locations or supergraph provided."
+        end
+      end
+
+      def execute(query:, variables: nil, operation_name: nil, context: EMPTY_CONTEXT, validate: true)
+        document = GraphQL::Stitching::Document.new(query, operation_name: operation_name)
+
+        if validate
+          validation_errors = @supergraph.schema.validate(document.ast)
+          return error_result(validation_errors) if validation_errors.any?
+        end
+
+        begin
+          plan = fetch_plan(document, context) do
+            GraphQL::Stitching::Planner.new(
+              supergraph: @supergraph,
+              document: document,
+            ).perform.to_h
+          end
+
+          GraphQL::Stitching::Executor.new(
+            supergraph: @supergraph,
+            plan: plan,
+            variables: variables || {},
+          ).perform(document)
+        rescue StandardError => e
+          custom_message = @on_error.call(e, context) if @on_error
+          error_result([{ "message" => custom_message || "An unexpected error occured." }])
+        end
+      end
+
+      def on_cache_read(&block)
+        raise GatewayError, "A cache read block is required." unless block_given?
+        @on_cache_read = block
+      end
+
+      def on_cache_write(&block)
+        raise GatewayError, "A cache write block is required." unless block_given?
+        @on_cache_write = block
+      end
+
+      def on_error(&block)
+        raise GatewayError, "An error handler block is required." unless block_given?
+        @on_error = block
+      end
+
+      private
+
+      def build_supergraph_from_locations_config(locations)
+        schemas = locations.each_with_object({}) do |(location, config), memo|
+          schema = config[:schema]
+          if schema.nil?
+            raise GatewayError, "A schema is required for `#{location}` location."
+          elsif !(schema.is_a?(Class) && schema <= GraphQL::Schema)
+            raise GatewayError, "The schema for `#{location}` location must be a GraphQL::Schema class."
+          else
+            memo[location.to_s] = schema
+          end
+        end
+
+        supergraph = GraphQL::Stitching::Composer.new(schemas: schemas).perform
+
+        locations.each do |location, config|
+          executable = config[:executable]
+          supergraph.assign_executable(location.to_s, executable) if executable
+        end
+
+        supergraph
+      end
+
+      def fetch_plan(document, context)
+        if @on_cache_read
+          cached_plan = @on_cache_read.call(document.digest, context)
+          return JSON.parse(cached_plan) if cached_plan
+        end
+
+        plan_json = yield
+
+        if @on_cache_write
+          @on_cache_write.call(document.digest, JSON.generate(plan_json), context)
+        end
+
+        plan_json
+      end
+
+      def error_result(errors)
+        public_errors = errors.map do |e|
+          public_error = e.is_a?(Hash) ? e : e.to_h
+          public_error["path"] ||= []
+          public_error
+        end
+
+        { "errors" => public_errors }
+      end
+    end
+  end
+end

data/lib/graphql/stitching/planner.rb

@@ -0,0 +1,323 @@
+# frozen_string_literal: true
+
+module GraphQL
+  module Stitching
+    class Planner
+      SUPERGRAPH_LOCATIONS = [Supergraph::LOCATION].freeze
+      TYPENAME_NODE = GraphQL::Language::Nodes::Field.new(alias: "_STITCH_typename", name: "__typename")
+
+      def initialize(supergraph:, document:)
+        @supergraph = supergraph
+        @document = document
+        @sequence_key = 0
+        @operations_by_grouping = {}
+      end
+
+      def perform
+        build_root_operations
+        expand_abstract_boundaries
+        self
+      end
+
+      def operations
+        ops = @operations_by_grouping.values
+        ops.sort_by!(&:key)
+        ops
+      end
+
+      def to_h
+        { "ops" => operations.map(&:to_h) }
+      end
+
+      private
+
+      def add_operation(location:, parent_type:, selections: nil, insertion_path: [], operation_type: "query", after_key: 0, boundary: nil)
+        parent_key = @sequence_key += 1
+        selection_set, variables = if selections&.any?
+          extract_locale_selections(location, parent_type, selections, insertion_path, parent_key)
+        end
+
+        grouping = [after_key, location, parent_type.graphql_name, *insertion_path].join("/")
+
+        if op = @operations_by_grouping[grouping]
+          op.selections += selection_set if selection_set
+          op.variables.merge!(variables) if variables
+          return op
+        end
+
+        type_conditional = !parent_type.kind.abstract? && parent_type != @supergraph.schema.query && parent_type != @supergraph.schema.mutation
+
+        @operations_by_grouping[grouping] = PlannerOperation.new(
+          key: parent_key,
+          after_key: after_key,
+          location: location,
+          parent_type: parent_type,
+          operation_type: operation_type,
+          insertion_path: insertion_path,
+          type_condition: type_conditional ? parent_type.graphql_name : nil,
+          selections: selection_set || [],
+          variables: variables || {},
+          boundary: boundary,
+        )
+      end
+
+      def build_root_operations
+        case @document.operation.operation_type
+        when "query"
+          # plan steps grouping all fields by location for async execution
+          parent_type = @supergraph.schema.query
+
+          selections_by_location = @document.operation.selections.each_with_object({}) do |node, memo|
+            locations = @supergraph.locations_by_type_and_field[parent_type.graphql_name][node.name] || SUPERGRAPH_LOCATIONS
+            memo[locations.last] ||= []
+            memo[locations.last] << node
+          end
+
+          selections_by_location.each do |location, selections|
+            add_operation(location: location, parent_type: parent_type, selections: selections)
+          end
+
+        when "mutation"
+          # plan steps grouping sequential fields by location for serial execution
+          parent_type = @supergraph.schema.mutation
+          location_groups = []
+
+          @document.operation.selections.reduce(nil) do |last_location, node|
+            location = @supergraph.locations_by_type_and_field[parent_type.graphql_name][node.name].last
+            if location != last_location
+              location_groups << {
+                location: location,
+                selections: [],
+              }
+            end
+            location_groups.last[:selections] << node
+            location
+          end
+
+          location_groups.reduce(0) do |after_key, group|
+            add_operation(
+              location: group[:location],
+              selections: group[:selections],
+              operation_type: "mutation",
+              parent_type: parent_type,
+              after_key: after_key
+            ).key
+          end
+
+        else
+          raise "Invalid operation type."
+        end
+      end
+
+      def extract_locale_selections(current_location, parent_type, input_selections, insertion_path, after_key)
+        remote_selections = []
+        selections_result = []
+        variables_result = {}
+        implements_fragments = false
+
+        if parent_type.kind.interface?
+          # fields of a merged interface may not belong to the interface at the local level,
+          # so these non-local interface fields get expanded into typed fragments for planning
+          local_interface_fields = @supergraph.fields_by_type_and_location[parent_type.graphql_name][current_location]
+          extended_selections = []
+
+          input_selections.reject! do |node|
+            if node.is_a?(GraphQL::Language::Nodes::Field) && !local_interface_fields.include?(node.name)
+              extended_selections << node
+              true
+            end
+          end
+
+          if extended_selections.any?
+            possible_types = Util.get_possible_types(@supergraph.schema, parent_type)
+            possible_types.each do |possible_type|
+              next if possible_type.kind.abstract? # ignore child interfaces
+              next unless @supergraph.locations_by_type[possible_type.graphql_name].include?(current_location)
+
+              type_name = GraphQL::Language::Nodes::TypeName.new(name: possible_type.graphql_name)
+              input_selections << GraphQL::Language::Nodes::InlineFragment.new(type: type_name, selections: extended_selections)
+            end
+          end
+        end
+
+        input_selections.each do |node|
+          case node
+          when GraphQL::Language::Nodes::Field
+            if node.name == "__typename"
+              selections_result << node
+              next
+            end
+
+            possible_locations = @supergraph.locations_by_type_and_field[parent_type.graphql_name][node.name] || SUPERGRAPH_LOCATIONS
+            unless possible_locations.include?(current_location)
+              remote_selections << node
+              next
+            end
+
+            field_type = Util.get_named_type_for_field_node(@supergraph.schema, parent_type, node)
+
+            extract_node_variables!(node, variables_result)
+
+            if Util.is_leaf_type?(field_type)
+              selections_result << node
+            else
+              expanded_path = [*insertion_path, node.alias || node.name]
+              selection_set, variables = extract_locale_selections(current_location, field_type, node.selections, expanded_path, after_key)
+              selections_result << node.merge(selections: selection_set)
+              variables_result.merge!(variables)
+            end
+
+          when GraphQL::Language::Nodes::InlineFragment
+            next unless @supergraph.locations_by_type[node.type.name].include?(current_location)
+
+            fragment_type = @supergraph.schema.types[node.type.name]
+            selection_set, variables = extract_locale_selections(current_location, fragment_type, node.selections, insertion_path, after_key)
+            selections_result << node.merge(selections: selection_set)
+            variables_result.merge!(variables)
+            implements_fragments = true
+
+          when GraphQL::Language::Nodes::FragmentSpread
+            fragment = @document.fragment_definitions[node.name]
+            next unless @supergraph.locations_by_type[fragment.type.name].include?(current_location)
+
+            fragment_type = @supergraph.schema.types[fragment.type.name]
+            selection_set, variables = extract_locale_selections(current_location, fragment_type, fragment.selections, insertion_path, after_key)
+            selections_result << GraphQL::Language::Nodes::InlineFragment.new(type: fragment.type, selections: selection_set)
+            variables_result.merge!(variables)
+            implements_fragments = true
+
+          else
+            raise "Unexpected node of type #{node.class.name} in selection set."
+          end
+        end
+
+        if remote_selections.any?
+          selection_set = build_child_operations(current_location, parent_type, remote_selections, insertion_path, after_key)
+          selections_result.concat(selection_set)
+        end
+
+        if parent_type.kind.abstract? || implements_fragments
+          selections_result << TYPENAME_NODE
+        end
+
+        return selections_result, variables_result
+      end
+
+      def build_child_operations(current_location, parent_type, input_selections, insertion_path, after_key)
+        parent_selections_result = []
+        selections_by_location = {}
+
+        # distribute unique fields among required locations
+        input_selections.reject! do |node|
+          possible_locations = @supergraph.locations_by_type_and_field[parent_type.graphql_name][node.name]
+          if possible_locations.length == 1
+            selections_by_location[possible_locations.first] ||= []
+            selections_by_location[possible_locations.first] << node
+            true
+          end
+        end
+
+        # distribute non-unique fields among available locations, preferring used locations
+        if input_selections.any?
+          # weight locations by number of needed fields available, prefer greater availability
+          location_weights = input_selections.each_with_object({}) do |node, memo|
+            possible_locations = @supergraph.locations_by_type_and_field[parent_type.graphql_name][node.name]
+            possible_locations.each do |location|
+              memo[location] ||= 0
+              memo[location] += 1
+            end
+          end
+
+          input_selections.each do |node|
+            possible_locations = @supergraph.locations_by_type_and_field[parent_type.graphql_name][node.name]
+
+            perfect_location_score = input_selections.length
+            preferred_location_score = 0
+            preferred_location = possible_locations.reduce(possible_locations.first) do |current_loc, candidate_loc|
+              score = selections_by_location[location] ? perfect_location_score : 0
+              score += location_weights.fetch(candidate_loc, 0)
+
+              if score > preferred_location_score
+                preferred_location_score = score
+                candidate_loc
+              else
+                current_loc
+              end
+            end
+
+            selections_by_location[preferred_location] ||= []
+            selections_by_location[preferred_location] << node
+          end
+        end
+
+        routes = @supergraph.route_type_to_locations(parent_type.graphql_name, current_location, selections_by_location.keys)
+        routes.values.each_with_object({}) do |route, memo|
+          route.reduce(nil) do |parent_op, boundary|
+            location = boundary["location"]
+            next memo[location] if memo[location]
+
+            child_op = memo[location] = add_operation(
+              location: location,
+              selections: selections_by_location[location],
+              parent_type: parent_type,
+              insertion_path: insertion_path,
+              boundary: boundary,
+              after_key: after_key,
+            )
+
+            foreign_key_node = GraphQL::Language::Nodes::Field.new(
+              alias: "_STITCH_#{boundary["selection"]}",
+              name: boundary["selection"]
+            )
+
+            if parent_op
+              parent_op.selections << foreign_key_node << TYPENAME_NODE
+            else
+              parent_selections_result << foreign_key_node << TYPENAME_NODE
+            end
+
+            child_op
+          end
+        end
+
+        parent_selections_result
+      end
+
+      def extract_node_variables!(node_with_args, variables={})
+        node_with_args.arguments.each_with_object(variables) do |argument, memo|
+          case argument.value
+          when GraphQL::Language::Nodes::InputObject
+            extract_node_variables!(argument.value, memo)
+          when GraphQL::Language::Nodes::VariableIdentifier
+            memo[argument.value.name] ||= @document.variable_definitions[argument.value.name]
+          end
+        end
+      end
+
+      # expand concrete type selections into typed fragments when sending to abstract boundaries
+      def expand_abstract_boundaries
+        @operations_by_grouping.each do |_grouping, op|
+          next unless op.boundary
+
+          boundary_type = @supergraph.schema.get_type(op.boundary["type_name"])
+          next unless boundary_type.kind.abstract?
+
+          unless op.parent_type == boundary_type
+            to_typed_selections = []
+            op.selections.reject! do |node|
+              if node.is_a?(GraphQL::Language::Nodes::Field)
+                to_typed_selections << node
+                true
+              end
+            end
+
+            if to_typed_selections.any?
+              type_name = GraphQL::Language::Nodes::TypeName.new(name: op.parent_type.graphql_name)
+              op.selections << GraphQL::Language::Nodes::InlineFragment.new(type: type_name, selections: to_typed_selections)
+            end
+          end
+        end
+      end
+    end
+  end
+end