From 68b0b9ad110a7c19a7f806626cc8a1cd5b0068cc Mon Sep 17 00:00:00 2001
From: Christian Schuetz <cmschuetz@users.noreply.github.com>
Date: Fri, 11 Oct 2019 19:36:11 -0500
Subject: [PATCH] fix: Support lazy objects (#17)

* fix: Support lazy objects

- Adds field extension that supports lazy object resolution and type assignment
- Adds field extension to _entities field

* test: Scope lazy class

- Scopes lazy object class within let block
- Resolves other rubocop issues
---
 lib/apollo-federation/entities_field.rb       | 10 ++---
 .../entity_type_resolution_extension.rb       | 18 ++++++++
 spec/apollo-federation/entities_field_spec.rb | 43 +++++++++++++++++++
 3 files changed, 64 insertions(+), 7 deletions(-)
 create mode 100644 lib/apollo-federation/entity_type_resolution_extension.rb

diff --git a/lib/apollo-federation/entities_field.rb b/lib/apollo-federation/entities_field.rb
index 26fbe7f87..fcdf115e9 100644
--- a/lib/apollo-federation/entities_field.rb
+++ b/lib/apollo-federation/entities_field.rb
@@ -2,6 +2,7 @@
 
 require 'graphql'
 require 'apollo-federation/any'
+require 'apollo-federation/entity_type_resolution_extension'
 
 module ApolloFederation
   module EntitiesField
@@ -15,6 +16,7 @@ module ClassMethods
       def define_entities_field(entity_type)
         field(:_entities, [entity_type, null: true], null: false) do
           argument :representations, [Any], required: true
+          extension(EntityTypeResolutionExtension)
         end
       end
     end
@@ -38,13 +40,7 @@ def _entities(representations:)
           result = reference
         end
 
-        # TODO: This isn't 100% correct: if (for some reason) 2 different resolve_reference calls
-        # return the same object, it might not have the right type
-        # Right now, apollo-federation just adds a __typename property to the result,
-        # but I don't really like the idea of modifying the resolved object
-        context[result] = type
-        # TODO: Handle lazy objects?
-        result
+        [type, result]
       end
     end
   end
diff --git a/lib/apollo-federation/entity_type_resolution_extension.rb b/lib/apollo-federation/entity_type_resolution_extension.rb
new file mode 100644
index 000000000..62b7aec84
--- /dev/null
+++ b/lib/apollo-federation/entity_type_resolution_extension.rb
@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+class EntityTypeResolutionExtension < GraphQL::Schema::FieldExtension
+  def after_resolve(value:, context:, **_rest)
+    synced_value =
+      value.map do |type, result|
+        [type, context.query.schema.sync_lazy(result)]
+      end
+
+    # TODO: This isn't 100% correct: if (for some reason) 2 different resolve_reference calls
+    # return the same object, it might not have the right type
+    # Right now, apollo-federation just adds a __typename property to the result,
+    # but I don't really like the idea of modifying the resolved object
+    synced_value.each { |type, result| context[result] = type }
+
+    synced_value.map { |_, result| result }
+  end
+end
diff --git a/spec/apollo-federation/entities_field_spec.rb b/spec/apollo-federation/entities_field_spec.rb
index 72b231a85..23b517abe 100644
--- a/spec/apollo-federation/entities_field_spec.rb
+++ b/spec/apollo-federation/entities_field_spec.rb
@@ -221,6 +221,49 @@ def self.resolve_reference(reference, _context)
 
               it { is_expected.to match_array [{ 'id' => id.to_s, 'otherField' => 'more data' }] }
               it { expect(errors).to be_nil }
+
+              context 'when resolve_reference returns a lazy object' do
+                let(:lazy_entity) do
+                  Class.new do
+                    def initialize(data)
+                      @data = data
+                    end
+
+                    def load_entity
+                      @data
+                    end
+                  end
+                end
+
+                let(:schema) do
+                  lazy_entity_class = lazy_entity
+                  type_with_key_class = type_with_key
+                  Class.new(base_schema) do
+                    lazy_resolve(lazy_entity_class, :load_entity)
+
+                    orphan_types type_with_key_class
+                  end
+                end
+
+                let(:type_with_key) do
+                  lazy_entity_class = lazy_entity
+                  Class.new(base_object) do
+                    graphql_name 'TypeWithKey'
+                    key fields: 'id'
+                    field :id, 'ID', null: false
+                    field :other_field, 'String', null: false
+
+                    define_singleton_method :resolve_reference do |reference, _context|
+                      if reference[:id] == 123
+                        lazy_entity_class.new(id: 123, other_field: 'more data')
+                      end
+                    end
+                  end
+                end
+
+                it { is_expected.to match_array [{ 'id' => id.to_s, 'otherField' => 'more data' }] }
+                it { expect(errors).to be_nil }
+              end
             end
           end
         end