From 957115d7a06511cc0a30edd584c3f80fe5ce7a7b Mon Sep 17 00:00:00 2001
From: Joel Turkel <Jturkel@salsify.com>
Date: Wed, 19 Apr 2023 18:22:07 -0400
Subject: [PATCH] fix: inherit Query type from schema superclass (#241)

---
 lib/apollo-federation/schema.rb               |  6 ++-
 .../service_field_v1_spec.rb                  | 33 +++++++++++++++++
 .../service_field_v2_spec.rb                  | 37 +++++++++++++++++++
 3 files changed, 75 insertions(+), 1 deletion(-)

diff --git a/lib/apollo-federation/schema.rb b/lib/apollo-federation/schema.rb
index c8a591744..ab9462c59 100644
--- a/lib/apollo-federation/schema.rb
+++ b/lib/apollo-federation/schema.rb
@@ -46,7 +46,7 @@ def query(new_query_object = nil)
           @orig_query_object = new_query_object
         else
           if !@federation_query_object
-            @federation_query_object = federation_query(@orig_query_object)
+            @federation_query_object = federation_query(original_query)
             @federation_query_object.define_entities_field(schema_entities)
 
             super(@federation_query_object)
@@ -58,6 +58,10 @@ def query(new_query_object = nil)
 
       private
 
+      def original_query
+        @orig_query_object || find_inherited_value(:original_query)
+      end
+
       def federation_preamble
         federation_namespace = ", as: \"#{link_namespace}\"" if link_namespace != DEFAULT_LINK_NAMESPACE
 
diff --git a/spec/apollo-federation/service_field_v1_spec.rb b/spec/apollo-federation/service_field_v1_spec.rb
index dc3192a7d..437ad92f0 100644
--- a/spec/apollo-federation/service_field_v1_spec.rb
+++ b/spec/apollo-federation/service_field_v1_spec.rb
@@ -479,6 +479,39 @@ def self.visible?(context)
       )
     end
 
+    it 'returns SDL that inherits the query type' do
+      product = Class.new(base_object) do
+        graphql_name 'Product'
+        extend_type
+
+        field :upc, String, null: false
+      end
+
+      query_obj = Class.new(base_object) do
+        graphql_name 'Query'
+
+        field :product, product, null: true
+      end
+
+      new_base_schema = Class.new(base_schema) do
+        query query_obj
+      end
+
+      schema = Class.new(new_base_schema)
+
+      expect(execute_sdl(schema)).to match_sdl(
+        <<~GRAPHQL,
+          type Product @extends {
+            upc: String!
+          }
+
+          type Query {
+            product: Product
+          }
+        GRAPHQL
+      )
+    end
+
     context 'with context in schema generation' do
       let(:schema) do
         product = Class.new(base_object) do
diff --git a/spec/apollo-federation/service_field_v2_spec.rb b/spec/apollo-federation/service_field_v2_spec.rb
index ca5347d07..60e0b89af 100644
--- a/spec/apollo-federation/service_field_v2_spec.rb
+++ b/spec/apollo-federation/service_field_v2_spec.rb
@@ -1946,6 +1946,43 @@ def self.visible?(context)
       )
     end
 
+    it 'returns SDL that inherits the query type' do
+      product = Class.new(base_object) do
+        graphql_name 'Product'
+        extend_type
+
+        field :upc, String, null: false
+      end
+
+      query_obj = Class.new(base_object) do
+        graphql_name 'Query'
+
+        field :product, product, null: true
+      end
+
+      new_base_schema = Class.new(base_schema) do
+        federation version: '2.0', link: { as: 'fed2' }
+        query query_obj
+      end
+
+      schema = Class.new(new_base_schema)
+
+      expect(execute_sdl(schema)).to match_sdl(
+        <<~GRAPHQL,
+          extend schema
+            @link(url: "https://specs.apollo.dev/federation/v2.3", as: "fed2", import: ["@inaccessible", "@tag"])
+
+          type Product @fed2__extends {
+            upc: String!
+          }
+
+          type Query {
+            product: Product
+          }
+        GRAPHQL
+      )
+    end
+
     context 'with context in schema generation' do
       let(:schema) do
         product = Class.new(base_object) do