From e39de8c8c45dd40273e26f80649b988d6b18b8c7 Mon Sep 17 00:00:00 2001
From: Shuhei Kitagawa <shuhei.kitagawa@c-fo.com>
Date: Mon, 30 Apr 2018 14:55:59 +0900
Subject: [PATCH] Enable to use block to define relationship

---
 lib/fast_jsonapi/object_serializer.rb  | 11 ++++++-----
 lib/fast_jsonapi/serialization_core.rb | 24 +++++++++++++++++++-----
 2 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/lib/fast_jsonapi/object_serializer.rb b/lib/fast_jsonapi/object_serializer.rb
index a3aaf36..952992d 100644
--- a/lib/fast_jsonapi/object_serializer.rb
+++ b/lib/fast_jsonapi/object_serializer.rb
@@ -181,21 +181,21 @@ module FastJsonapi
         self.relationships_to_serialize[name] = relationship
      end
 
-      def has_many(relationship_name, options = {})
+      def has_many(relationship_name, options = {}, &block)
         name = relationship_name.to_sym
-        hash = create_relationship_hash(relationship_name, :has_many, options)
+        hash = create_relationship_hash(relationship_name, :has_many, options, block)
         add_relationship(name, hash)
       end
 
-      def has_one(relationship_name, options = {})
+      def has_one(relationship_name, options = {}, &block)
         name = relationship_name.to_sym
-        hash = create_relationship_hash(relationship_name, :has_one, options)
+        hash = create_relationship_hash(relationship_name, :has_one, options, block)
         add_relationship(name, hash)
       end
 
       alias belongs_to has_one
 
-      def create_relationship_hash(base_key, relationship_type, options)
+      def create_relationship_hash(base_key, relationship_type, options, block)
         name = base_key.to_sym
         if relationship_type == :has_many
           base_serialization_key = base_key.to_s.singularize
@@ -212,6 +212,7 @@ module FastJsonapi
           id_method_name: options[:id_method_name] || "#{base_serialization_key}#{id_postfix}".to_sym,
           record_type: options[:record_type] || run_key_transform(base_key_sym),
           object_method_name: options[:object_method_name] || name,
+          object_block: block,
           serializer: compute_serializer_name(options[:serializer] || base_key_sym),
           relationship_type: relationship_type,
           cached: options[:cached] || false,
diff --git a/lib/fast_jsonapi/serialization_core.rb b/lib/fast_jsonapi/serialization_core.rb
index ab7787e..f979f17 100644
--- a/lib/fast_jsonapi/serialization_core.rb
+++ b/lib/fast_jsonapi/serialization_core.rb
@@ -48,11 +48,11 @@ module FastJsonapi
         polymorphic = relationship[:polymorphic]
 
         return ids_hash(
-          record.public_send(relationship[:id_method_name]),
+          fetch_id(record, relationship),
           relationship[:record_type]
         ) unless polymorphic
 
-        return unless associated_object = record.send(relationship[:object_method_name])
+        return unless associated_object = fetch_associated_object(record, relationship)
 
         return associated_object.map do |object|
           id_hash_from_record object, polymorphic
@@ -117,9 +117,7 @@ module FastJsonapi
 
       def get_included_records(record, includes_list, known_included_objects)
         includes_list.each_with_object([]) do |item, included_records|
-          included_objects = record.send(
-            @relationships_to_serialize[item][:object_method_name]
-          )
+          included_objects = fetch_associated_object(record, @relationships_to_serialize[item])
           next if included_objects.blank?
 
           record_type = @relationships_to_serialize[item][:record_type]
@@ -134,6 +132,22 @@ module FastJsonapi
           end
         end
       end
+
+      def fetch_associated_object(record, relationship)
+        return relationship[:object_block].call(record) unless relationship[:object_block].nil?
+        record.send(relationship[:object_method_name])
+      end
+
+      def fetch_id(record, relationship)
+        unless relationship[:object_block].nil?
+          object = relationship[:object_block].call(record)
+
+          return object.map(&:id) if object.respond_to? :map
+          return object.id
+        end
+
+        record.public_send(relationship[:id_method_name])
+      end
     end
   end
 end