stitch federation schemas via _entities. (#65)

README.md

@@ -9,6 +9,7 @@ GraphQL stitching composes a single schema from multiple underlying GraphQL reso
 - Multiple keys per merged type.
 - Shared objects, fields, enums, and inputs across locations.
 - Combining local and remote schemas.
+- Type merging via federation `_entities` protocol.
 
 **NOT Supported:**
 - Computed fields (ie: federation-style `@requires`).
@@ -294,6 +295,66 @@ The library is configured to use a `@stitch` directive by default. You may custo
 GraphQL::Stitching.stitch_directive = "merge"
 ```
 
+#### Federation entities
+
+The [Apollo Federation specification](https://www.apollographql.com/docs/federation/subgraph-spec/) defines a standard interface for accessing merged object types across locations. Stitching can utilize a subset of this interface to facilitate basic type merging. The following spec is supported:
+
+- `@key(fields: "id")` (repeatable) specifies a key field for an object type. Keys may only select one field each.
+- `_Entity` is a union type that must contain all types that implement a `@key`.
+- `_Any` is a scalar that recieves raw JSON objects.
+- `_entities(representations: [_Any!]!): [_Entity]!` is a root query for local entity types.
+
+The composer will automatcially detect and stitch schemas with an `_entities` query, for example:
+
+```ruby
+accounts_schema = <<~GRAPHQL
+  directive @key(fields: String!) repeatable on OBJECT
+
+  type User @key(fields: "id") {
+    id: ID!
+    name: String!
+    address: String!
+  }
+
+  union _Entity = User
+  scalar _Any
+
+  type Query {
+    user(id: ID!): User
+    _entities(representations: [_Any!]!): [_Entity]!
+  }
+GRAPHQL
+
+comments_schema = <<~GRAPHQL
+  directive @key(fields: String!) repeatable on OBJECT
+
+  type User @key(fields: "id") {
+    id: ID!
+    comments: [String!]!
+  }
+
+  union _Entity = User
+  scalar _Any
+
+  type Query {
+    _entities(representations: [_Any!]!): [_Entity]!
+  }
+GRAPHQL
+
+client = GraphQL::Stitching::Client.new(locations: {
+  accounts: {
+    schema: GraphQL::Schema.from_definition(accounts_schema),
+    executable: ...,
+  },
+  comments: {
+    schema: GraphQL::Schema.from_definition(comments_schema),
+    executable: ...,
+  },
+})
+```
+
+It's perfectly fine to mix and match schemas that implement an `_entities` query with schemas that implement `@stitch` directives; the protocols achieve the same result. Note that stitching is much simpler than Apollo Federation by design, and so advanced routing features such as computed fields (ie: the `@requires` directive) will be ignored.
+
 ## Executables
 
 An executable resource performs location-specific GraphQL requests. Executables may be `GraphQL::Schema` classes, or any object that responds to `.call(location, source, variables, context)` and returns a raw GraphQL response:

lib/graphql/stitching/composer.rb

@@ -158,6 +158,26 @@ def prepare_locations_input(locations_input)
             @stitch_directives[field_path] << dir.slice(:key, :type_name)
           end
 
+          federation_entity_type = schema.types["_Entity"]
+          if federation_entity_type && federation_entity_type.kind.union? && schema.query.fields["_entities"]&.type&.unwrap == federation_entity_type
+            schema.possible_types(federation_entity_type).each do |entity_type|
+              entity_type.directives.each do |directive|
+                next unless directive.graphql_name == "key"
+
+                key = directive.arguments.keyword_arguments.fetch(:fields).strip
+                raise ComposerError, "Composite federation keys are not supported." unless /^\w+$/.match?(key)
+
+                field_path = "#{location}._entities"
+                @stitch_directives[field_path] ||= []
+                @stitch_directives[field_path] << {
+                  key: key,
+                  type_name: entity_type.graphql_name,
+                  federation: true,
+                }
+              end
+            end
+          end
+
           schemas[location.to_s] = schema
           executables[location.to_s] = input[:executable] || schema
         end
@@ -465,6 +485,7 @@ def extract_boundaries(type_name, types_by_location)
 
             boundary_kwargs.each do |kwargs|
               key = kwargs.fetch(:key)
+              impl_type_name = kwargs.fetch(:type_name, boundary_type_name)
               key_selections = GraphQL.parse("{ #{key} }").definitions[0].selections
 
               if key_selections.length != 1
@@ -487,15 +508,16 @@ def extract_boundaries(type_name, types_by_location)
                 raise ComposerError, "Mismatched input/output for #{type_name}.#{field_name}.#{argument_name} boundary. Arguments must map directly to results."
               end
 
-              @boundary_map[boundary_type_name] ||= []
-              @boundary_map[boundary_type_name] << {
+              @boundary_map[impl_type_name] ||= []
+              @boundary_map[impl_type_name] << {
                 "location" => location,
+                "type_name" => impl_type_name,
                 "key" => key_selections[0].name,
                 "field" => field_candidate.name,
                 "arg" => argument_name,
                 "list" => boundary_structure.first[:list],
-                "type_name" => boundary_type_name,
-              }
+                "federation" => kwargs[:federation],
+              }.compact
             end
           end
         end

lib/graphql/stitching/executor.rb

@@ -96,21 +96,24 @@ def build_document(origin_sets_by_operation, operation_name = nil)
           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["key"]}"
 
             if boundary["list"]
-              input = JSON.generate(origin_set.map { _1[key_selection] })
-              "_#{batch_index}_result: #{boundary["field"]}(#{boundary["arg"]}:#{input}) #{op["selections"]}"
+              input = origin_set.each_with_index.reduce(String.new) do |memo, (origin_obj, index)|
+                memo << "," if index > 0
+                memo << build_key(boundary["key"], origin_obj, federation: boundary["federation"])
+                memo
+              end
+
+              "_#{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])
+                input = build_key(boundary["key"], origin_obj, federation: boundary["federation"])
                 "_#{batch_index}_#{index}_result: #{boundary["field"]}(#{boundary["arg"]}:#{input}) #{op["selections"]}"
               end
             end
           end
 
-          doc = String.new
-          doc << "query" # << boundary fulfillment always uses query
+          doc = String.new("query") # << boundary fulfillment always uses query
 
           if operation_name
             doc << " #{operation_name}"
@@ -129,6 +132,15 @@ def build_document(origin_sets_by_operation, operation_name = nil)
           return doc, variable_defs.keys
         end
 
+        def build_key(key, origin_obj, federation: false)
+          key_value = JSON.generate(origin_obj["_STITCH_#{key}"])
+          if federation
+            "{ __typename: \"#{origin_obj["_STITCH_typename"]}\", #{key}: #{key_value} }"
+          else
+            key_value
+          end
+        end
+
         def merge_results!(origin_sets_by_operation, raw_result)
           return unless raw_result
 

test/graphql/stitching/federation_test.rb

@@ -0,0 +1,85 @@
+# frozen_string_literal: true
+
+require "test_helper"
+require_relative "../../schemas/federation"
+
+describe "GraphQL::Stitching::Federation" do
+
+  def test_federation_to_stitching
+    @supergraph = compose_definitions({
+      "federation" => Schemas::Federation::Federation1,
+      "stitching" => Schemas::Federation::Stitching,
+    })
+
+    query = %|
+      query {
+        widgets(upcs: ["1"]) {
+          megahertz
+          model
+          sprockets {
+            cogs
+            diameter
+          }
+        }
+      }
+    |
+
+    expected = {
+      "data" => {
+        "widgets" => [{
+          "megahertz" => 3,
+          "model" => "Basic",
+          "sprockets" => [
+            { "cogs" => 23, "diameter" => 77 },
+            { "cogs" => 14, "diameter" => 20 },
+          ],
+        }],
+      },
+    }
+
+    result = plan_and_execute(@supergraph, query) do |plan|
+      assert_equal ["stitching", "federation", "stitching"], plan.operations.map(&:location)
+    end
+
+    assert_equal expected, result
+  end
+
+  def test_federation_to_federation
+    @supergraph = compose_definitions({
+      "federation1" => Schemas::Federation::Federation1,
+      "federation2" => Schemas::Federation::Federation2,
+    })
+
+    query = %|
+      query {
+        widget {
+          megahertz
+          model
+          sprockets {
+            cogs
+            diameter
+          }
+        }
+      }
+    |
+
+    expected = {
+      "data" => {
+        "widget" => {
+          "megahertz" => 3,
+          "model" => "Basic",
+          "sprockets" => [
+            { "cogs" => 23, "diameter" => 77 },
+            { "cogs" => 14, "diameter" => 20 },
+          ],
+        },
+      },
+    }
+
+    result = plan_and_execute(@supergraph, query) do |plan|
+      assert_equal ["federation2", "federation1", "federation2"], plan.operations.map(&:location)
+    end
+
+    assert_equal expected, result
+  end
+end