From 665d65d755e00e3a9363d7a6b22416543c9e3cee Mon Sep 17 00:00:00 2001
From: Christian Joudrey <cmallette@gmail.com>
Date: Sun, 12 Feb 2017 15:38:14 -0500
Subject: [PATCH] Batch load :homeworld associations

---
 app/models/graph.rb                    |  1 +
 app/models/graph/association_loader.rb | 49 ++++++++++++++++++++++++++
 app/models/graph/types/person.rb       |  8 +++--
 app/models/graph/types/species.rb      |  6 +++-
 4 files changed, 61 insertions(+), 3 deletions(-)
 create mode 100644 app/models/graph/association_loader.rb

diff --git a/app/models/graph.rb b/app/models/graph.rb
index 5579df0..d8ca5e0 100644
--- a/app/models/graph.rb
+++ b/app/models/graph.rb
@@ -7,6 +7,7 @@ def find_by_id_field(type, model)
         resolve ->(_, args, _) do
           gid = GlobalID.parse(args[:id])
 
+          return unless gid
           return unless gid.model_name == type.name
 
           Graph::FindLoader.for(model).load(gid.model_id.to_i)
diff --git a/app/models/graph/association_loader.rb b/app/models/graph/association_loader.rb
new file mode 100644
index 0000000..6a195a7
--- /dev/null
+++ b/app/models/graph/association_loader.rb
@@ -0,0 +1,49 @@
+# From: https://github.com/Shopify/graphql-batch/issues/24#issuecomment-277331157
+class Graph::AssociationLoader < GraphQL::Batch::Loader
+  def self.validate(model, association_name)
+    new(model, association_name)
+    nil
+  end
+
+  def initialize(model, association_name)
+    @model = model
+    @association_name = association_name
+    validate
+  end
+
+  def load(record)
+    raise TypeError, "#{@model} loader can't load association for #{record.class}" unless record.is_a?(@model)
+    return Promise.resolve(read_association(record)) if association_loaded?(record)
+    super
+  end
+
+  # We want to load the associations on all records, even if they have the same id
+  def cache_key(record)
+    record.object_id
+  end
+
+  def perform(records)
+    preload_association(records)
+    records.each { |record| fulfill(record, read_association(record)) }
+  end
+
+  private
+
+  def validate
+    unless @model.reflect_on_association(@association_name)
+      raise ArgumentError, "No association #{@association_name} on #{@model}"
+    end
+  end
+
+  def preload_association(records)
+    ::ActiveRecord::Associations::Preloader.new.preload(records, @association_name)
+  end
+
+  def read_association(record)
+    record.public_send(@association_name)
+  end
+
+  def association_loaded?(record)
+    record.association(@association_name).loaded?
+  end
+end
diff --git a/app/models/graph/types/person.rb b/app/models/graph/types/person.rb
index 567f146..0b77a6a 100644
--- a/app/models/graph/types/person.rb
+++ b/app/models/graph/types/person.rb
@@ -22,13 +22,17 @@ module Types
         "The eye color of this person. Will be \"unknown\" if not known or \"n/a\" if the person does not have an eye.",
         property: :eye_color
 
-      field :gender, GenderEnum, "​The gender of this person."
+      field :gender, GenderEnum, "The gender of this person."
       field :hairColor, types.String,
         "The hair color of this person. Will be \"unknown\" if not known or \"n/a\" if the person does not have hair.",
         property: :hair_color
 
       field :height, types.Int, "The height of the person in centimeters."
-      field :homeworld, Planet, "A planet that this person was born on or inhabits."
+      field :homeworld, Planet, "A planet that this person was born on or inhabits." do
+        resolve -> (person, _, _) do
+          Graph::AssociationLoader.for(::Person, :homeworld).load(person)
+        end
+      end
       field :mass, types.Int, "The mass of the person in kilograms."
       field :name, !types.String, "The name of this person."
       field :skinColor, types.String, "The skin color of this person.", property: :skin_color
diff --git a/app/models/graph/types/species.rb b/app/models/graph/types/species.rb
index ce87fe4..c86ef04 100644
--- a/app/models/graph/types/species.rb
+++ b/app/models/graph/types/species.rb
@@ -50,7 +50,11 @@ module Types
       end
 
       field :language, types.String, "The language commonly spoken by this species."
-      field :homeworld, Graph::Types::Planet, "A planet that this species originates from type."
+      field :homeworld, Planet, "A planet that this species originates from type." do
+        resolve -> (species, _, _) do
+          Graph::AssociationLoader.for(::Species, :homeworld).load(species)
+        end
+      end
     end
   end
 end