Restore relation connection with original nodes

Author: ssnickolay

lib/graphql/cache/resolvers/connection_resolver.rb

@@ -6,12 +6,17 @@ module Resolvers
       # Pass cache write method into GraphQL::Relay::BaseConnection
       # and wrap them original Connection methods
       class ConnectionResolver < BaseResolver
-        class ConnectionCache < Module
+        NodesCache = Struct.new(:nodes, :paged_nodes)
+
+        class RelationConnectionOverload < Module
           module WrappedMethods
             def paged_nodes
               cache_write = instance_variable_get(:@__cache_write)
 
-              cache_write.call { super }
+              super.tap do |result|
+                # save original relation (aka @nodes) and loaded records
+                cache_write.call { NodesCache.new(@nodes, result) }
+              end
             end
           end
 
@@ -27,26 +32,43 @@ def extended(base)
 
         def call(args:, field:, parent:, context:, force_cache:)
           if force_cache || (cached = read).nil?
-            define_connection_cache(resolve_proc.call)
+            define_relation_cache(resolve_proc.call)
           else
             wrap_connection(cached, args, field, parent: parent, context: context)
           end
         end
 
         private
 
-        def wrap_connection(value, args, field, **kwargs)
-          GraphQL::Relay::BaseConnection.connection_for_nodes(value).new(
-            value,
+        def wrap_connection(cached, args, field, **kwargs)
+          nodes, paged_nodes = parse(cached)
+
+          GraphQL::Relay::BaseConnection.connection_for_nodes(nodes).new(
+            nodes,
             args,
             field: field,
             parent: kwargs[:parent],
             context: kwargs[:context]
-          )
+          ).tap do |connection|
+            # restore cached paged_nodes
+            connection.instance_variable_set(:@paged_nodes, paged_nodes) if paged_nodes
+          end
+        end
+
+        def define_relation_cache(connection)
+          if connection.is_a?(GraphQL::Relay::RelationConnection)
+            # inject cached logic into the relation connection
+            connection.extend(RelationConnectionOverload.new(method(:write)))
+          else
+            # cache loaded connection (works for ArrayConnection)
+            write { connection }
+          end
         end
 
-        def define_connection_cache(connection)
-          connection.extend(ConnectionCache.new(method(:write)))
+        def parse(cached)
+          return [cached, nil] unless cached.is_a?(NodesCache)
+
+          [cached.nodes, cached.paged_nodes]
         end
       end
     end

spec/features/connections_spec.rb

@@ -105,7 +105,7 @@ def execute(query, context = {})
         <<~SQL
           SELECT \"customers\".* FROM \"customers\" ORDER BY \"customers\".\"id\" DESC LIMIT ?\e[0m  [[\"LIMIT\", 1]]
           SELECT \"customers\".* FROM \"customers\" WHERE \"customers\".\"id\" = ? LIMIT ?\e[0m  [[\"id\", 1], [\"LIMIT\", 1]]
-          SELECT \"orders\".* FROM \"orders\" WHERE \"orders\".\"customer_id\" = ?\e[0m  [[\"customer_id\", 1]]
+          SELECT \"orders\".* FROM \"orders\" WHERE \"orders\".\"customer_id\" = ? LIMIT ?\e[0m  [[\"customer_id\", 1], [\"LIMIT\", 50]]
         SQL
       )
     end

spec/spec_helper.rb

@@ -1,4 +1,7 @@
-require 'active_record' # should be required before graphql-ruby
+# ORMs should be required before graphql-ruby
+require 'active_record'
+require 'sequel'
+
 require 'bundler/setup'
 require 'pry'
 

test_schema/active_record/graphql_schema.rb

@@ -16,6 +16,8 @@ class CacheSchema < ::CacheSchema
     query AR::QueryType
     use GraphQL::Cache
 
+    default_max_page_size 50
+
     def self.resolve_type(_type, obj, _ctx)
       "AR::#{obj.class.name}Type"
     end

test_schema/sequel/graphql_schema.rb

@@ -31,6 +31,8 @@ class CacheSchema < GraphQL::Schema
 
   use GraphQL::Cache
 
+  default_max_page_size 50
+
   def self.resolve_type(_type, obj, _ctx)
     "#{obj.class.name}Type"
   end