Fix connection wrapper inheritance

Author: rmosolgo

lib/graphql/pagination/connections.rb

@@ -30,16 +30,17 @@ class ImplementationMissingError < GraphQL::Error
 
       def self.use(schema_defn)
         if schema_defn.is_a?(Class)
-          schema_defn.connections = self.new
+          schema_defn.connections = self.new(schema: schema_defn)
         else
           # Unwrap a `.define` object
           schema_defn = schema_defn.target
-          schema_defn.connections = self.new
+          schema_defn.connections = self.new(schema: schema_defn)
           schema_defn.class.connections = schema_defn.connections
         end
       end
 
-      def initialize
+      def initialize(schema:)
+        @schema = schema
         @wrappers = {}
         add_default
       end
@@ -52,12 +53,23 @@ def delete(nodes_class)
         @wrappers.delete(nodes_class)
       end
 
+      def all_wrappers
+        all_wrappers = {}
+        @schema.ancestors.reverse_each do |schema_class|
+          if schema_class.respond_to?(:connections) && (c = schema_class.connections)
+            all_wrappers.merge!(c.wrappers)
+          end
+        end
+        all_wrappers
+      end
+
       # Used by the runtime to wrap values in connection wrappers.
       # @api Private
-      def wrap(field, object, arguments, context)
+      def wrap(field, object, arguments, context, wrappers: all_wrappers)
         impl = nil
+
         object.class.ancestors.each { |cls|
-          impl = @wrappers[cls]
+          impl = wrappers[cls]
           break if impl
         }
 
@@ -76,6 +88,10 @@ def wrap(field, object, arguments, context)
         )
       end
 
+      protected
+
+      attr_reader :wrappers
+
       private
 
       def add_default

lib/graphql/schema.rb

@@ -942,8 +942,24 @@ def get_type(type_name)
           find_inherited_value(:types, EMPTY_HASH)[type_name]
       end
 
+      # @api private
+      attr_writer :connections
+
       # @return [GraphQL::Pagination::Connections] if installed
-      attr_accessor :connections
+      def connections
+        if defined?(@connections)
+          @connections
+        else
+          inherited_connections = find_inherited_value(:connections, nil)
+          # This schema is part of an inheritance chain which is using new connections,
+          # make a new instance, so we don't pollute the upstream one.
+          if inherited_connections
+            @connections = Pagination::Connections.new(schema: self)
+          else
+            nil
+          end
+        end
+      end
 
       def new_connections?
         !!connections

lib/graphql/schema/field/connection_extension.rb

@@ -38,7 +38,8 @@ def after_resolve(value:, object:, arguments:, context:, memo:)
             value.max_page_size ||= field.max_page_size
             value
           elsif context.schema.new_connections?
-            context.schema.connections.wrap(field, value, arguments, context)
+            wrappers = context.namespace(:connections)[:all_wrappers] ||= context.schema.connections.all_wrappers
+            context.schema.connections.wrap(field, value, arguments, context, wrappers: wrappers)
           else
             if object.is_a?(GraphQL::Schema::Object)
               object = object.object

spec/graphql/pagination/connections_spec.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+require "spec_helper"
+
+describe GraphQL::Pagination::Connections do
+  ITEMS = ConnectionAssertions::NAMES.map { |n| { name: n } }
+
+  class ArrayConnectionWithTotalCount < GraphQL::Pagination::ArrayConnection
+    def total_count
+      items.size
+    end
+  end
+
+  let(:base_schema) {
+    ConnectionAssertions.build_schema(
+      connection_class: GraphQL::Pagination::ArrayConnection,
+      total_count_connection_class: ArrayConnectionWithTotalCount,
+      get_items: -> { ITEMS }
+    )
+  }
+
+  # These wouldn't _work_, I just need to test `.wrap`
+  class SetConnection < GraphQL::Pagination::ArrayConnection; end
+  class HashConnection < GraphQL::Pagination::ArrayConnection; end
+  class OtherArrayConnection < GraphQL::Pagination::ArrayConnection; end
+
+  let(:schema) do
+    other_base_schema = Class.new(base_schema) do
+      connections.add(Set, SetConnection)
+    end
+
+    Class.new(other_base_schema) do
+      connections.add(Hash, HashConnection)
+      connections.add(Array, OtherArrayConnection)
+    end
+  end
+
+  it "returns connections by class, using inherited mappings and local overrides" do
+    field_defn = OpenStruct.new(max_page_size: 10)
+
+    set_wrapper = schema.connections.wrap(field_defn, Set.new([1,2,3]), {}, nil)
+    assert_instance_of SetConnection, set_wrapper
+
+    hash_wrapper = schema.connections.wrap(field_defn, {1 => :a, 2 => :b}, {}, nil)
+    assert_instance_of HashConnection, hash_wrapper
+
+    array_wrapper = schema.connections.wrap(field_defn, [1,2,3], {}, nil)
+    assert_instance_of OtherArrayConnection, array_wrapper
+  end
+
+  it "uses passed-in wrappers" do
+    field_defn = OpenStruct.new(max_page_size: 10)
+
+    assert_raises GraphQL::Pagination::Connections::ImplementationMissingError do
+      schema.connections.wrap(field_defn, Set.new([1,2,3]), {}, nil, wrappers: {})
+    end
+  end
+end

spec/support/connection_assertions.rb

@@ -20,8 +20,11 @@ module ConnectionAssertions
   ]
 
   def self.build_schema(get_items:, connection_class:, total_count_connection_class:)
-    Class.new(GraphQL::Schema) do
+    base_schema = Class.new(GraphQL::Schema) do
       use GraphQL::Pagination::Connections
+    end
+
+    Class.new(base_schema) do
       use GraphQL::Execution::Interpreter
 
       default_max_page_size ConnectionAssertions::MAX_PAGE_SIZE