Preloader for RESTful API using ActiveModelSerializers

preloader.rb

module Api
  ## Used in ActionController
  class Preloader
    # @param [Array] attributes   fields to select.
    # @param [Hash]  associations associations to include and fields to select of them.
    # @example When you select id and name from companies,
    #   preload_for(companies, [:id, :name], {})
    # @example When you includes posts and employees and employees avatar,
    #   preload_for(companies, [], {posts: {}, employees: {avatar: {}}})
    def self.preload_for(rel, attributes, associations)
      unless rel.is_a?(ActiveRecord::Relation)
        raise TypeError.new("Expected ActiveRecord::Relation, but got #{rel.class}")
      end
      attributes ||= []
      associations ||= {}

      model = rel.klass
      serializer = get_serializer(model)

      preloader = self.new(rel, model, serializer, attributes, associations)
      preloader.perform!
      preloader.rel
    end

    def self.get_serializer(klass)
      serializer = ActiveModel::Serializer.serializer_for(klass.new)
      unless serializer
        raise SerializerNotFound.new(klass)
      end
      serializer
    end

    def initialize(rel, model, serializer, attributes, associations, chain = [])
      unless serializer.is_a?(Class) && serializer <= ActiveModel::Serializer
        raise TypeError.new("Expected a descendant of ActiveModel::Serializer, but got #{serializer}")
      end

      @rel, @model, @serializer, @attributes, @associations, @chain = rel, model, serializer, attributes, associations, chain
    end

    attr_accessor :rel, :model, :serializer, :attributes, :associations

    def perform!
      return unless serializer.preload_rule
      attributes.each do |attribute|
        preload_for_attribute!(attribute)
      end

      associations.each do |association, nested|
        preload_for_association!(association, nested)
      end

      nil
    end

    private

    def preload_for_attribute!(attribute)
      rule = serializer.preload_rule.attributes[attribute]
      return unless rule

      self.rel = self.rel.includes(includes_on_chain(rule[:includes]))
    end

    def preload_for_association!(association, nested)
      rule = serializer.preload_rule.associations[association]
      return unless rule

      ar_association = rule[:includes] || association
      old = self.rel
      self.rel = self.rel.includes(includes_on_chain(ar_association))
      # binding.pry if self.rel.nil?

      if nested.present?
        nested = nested.dup
        nested_attributes = nested.delete(:only) || []
        nested_associations = nested

        association_object = model.reflections[ar_association.to_s]
        unless association_object
          raise "Could not find association(#{ar_association}) from class(#{model})"
        end
        nested_model = association_object.klass
        case rule[:serializer]
        when NilClass
          rule[:serializer] = self.class.get_serializer(nested_model)
        when String
          rule[:serializer] = const_get(rule[:serializer])
          unless rule[:serializer].is_a?(Class) && rule[:serializer] <= ActiveModel::Serializer
            raise TypeError.new("Expected a descendant of ActiveModel::Serializer, but got #{rule[:serializer]}")
          end
        end

        preloader = self.class.new(self.rel, nested_model, rule[:serializer], nested_attributes, nested_associations, (@chain + [ar_association]))
        preloader.perform!
        self.rel = preloader.rel
      end
    end

    # @example When `@chain` is [:jobs, :employee],
    #   includes_on_chain(:avatar) #=> {:jobs=>{:employees=>{:avatar=>{}}}}
    def includes_on_chain(association)
      hash = {}
      (@chain + [association]).inject(hash) { |hash, association|
        hash[association] = {}
      }
      hash
    end

    class SerializerNotFound < StandardError
      def initialize(klass)
        @klass = klass
      end

      def message
        "Serializer was not found for class: #{klass.name}"
      end
    end
  end

  ## Used in ActiveModelSerializers

  class PreloadRule
    def initialize
      @attributes, @associations = {}, {}
    end

    attr_accessor :attributes, :associations

    def attribute(name, includes:)
      unless name.is_a?(Symbol)
        raise TypeError.new("Expected Symbol, but got #{name.class}")
      end
      unless includes.is_a?(Symbol)
        raise TypeError.new("Expected Symbol, but got #{includes.class}")
      end

      attributes[name] = { includes: includes }
    end

    def association(name, includes: nil, serializer: nil)
      unless name.is_a?(Symbol)
        raise TypeError.new("Expected Symbol, but got #{name.class}")
      end
      unless includes.nil? || includes.is_a?(Symbol)
        raise TypeError.new("Expected Symbol, but got #{includes.class}")
      end
      unless serializer.nil? || serializer.is_a?(String) || (serializer.is_a?(Class) && serializer <= ActiveModel::Serializer)
        raise TypeError.new("Expected String of class name or an ancestor of ActiveModel::Serializer, but got #{serializer.class}")
      end

      # Assume association name to include
      includes ||= name

      associations[name.to_sym] = { includes: includes, serializer: serializer }
    end
  end
end

## spec

require 'rails_helper'

describe Api::Preloader do
  let!(:company) {
    company = create(:company)
    create(:user).first_or_create_job_for(company, :employee)
    company
  }
  let(:rel) { Company.all }

  describe '.get_serializer' do
    it 'gets serializer' do
      expect(Api::Preloader.get_serializer(Company)).to eq(CompanySerializer)
    end

    it 'raises not found error if serializer donesn\'t exist' do
      expect { Api::Preloader.get_serializer(Ranking) }.to raise_error(Api::Preloader::SerializerNotFound)
    end
  end

  describe '.preload_for' do
    it 'returns same relation' do
      result = Api::Preloader.preload_for(rel, [], {})
      expect(result).to be_an(ActiveRecord::Relation)
      expect(result.klass).to eq(Company)
      expect(result.count).to eq(rel.count)
    end

    it 'performs' do
      expect_any_instance_of(Api::Preloader).to receive(:perform!)
      Api::Preloader.preload_for(rel, [], {})
    end

    before do
      class CompanySerializer < ApplicationSerializer
        has_many :employees
        attributes :twitter
      end

      class EmployeeSerializer < UserSerializer
        has_one :avatar
      end
    end

    it 'preloads from an attribute rule' do
      CompanySerializer.preload do
        # nothing
      end

      loaded = Api::Preloader.preload_for(rel, [:id, :twitter], {})
      loaded.load
      expect(loaded.first.association(:links).loaded?).to eq(false)

      CompanySerializer.preload do
        attribute :twitter, includes: :links
      end

      loaded = Api::Preloader.preload_for(rel, [:id, :twitter], {})
      loaded.load

      expect(loaded.first.association(:links).loaded?).to eq(true)
    end

    it 'preloads an association' do
      CompanySerializer.preload do
        # nothing
      end

      loaded = Api::Preloader.preload_for(rel, [:id], { employees: {}})
      loaded.load
      expect(loaded.first.association(:employees).loaded?).to eq(false)

      CompanySerializer.preload do
        association :employees
      end

      loaded = Api::Preloader.preload_for(rel, [:id], { employees: {}})
      loaded.load
      expect(loaded.first.association(:employees).loaded?).to eq(true)
    end

    it 'preloads an nested association' do
      CompanySerializer.preload do
        # nothing
      end

      EmployeeSerializer.preload do
        # nothing
      end

      loaded = Api::Preloader.preload_for(rel, [:id], { employees: {avatar: {}}})
      loaded.load
      expect(loaded.first.employees.first.association(:avatar).loaded?).to eq(false)

      CompanySerializer.preload do
        association :employees
      end

      EmployeeSerializer.preload do
        association :avatar
      end

      loaded = Api::Preloader.preload_for(rel, [:id], { employees: {avatar: {}}})
      loaded.load
      expect(loaded.first.employees.first.association(:avatar).loaded?).to eq(true)
    end
  end
end

require 'rails_helper'

describe Api::PreloadRule do
  let(:rule) { Api::PreloadRule.new }

  it 'holds attributes as array' do
    expect(rule.attributes).to eq({})
  end

  it 'holds associations as hash' do
    expect(rule.associations).to eq({})
  end

  describe '#attribute' do
    it 'add attributes' do
      rule.attribute :facebook, includes: :links
      expect(rule.attributes[:facebook]).to eq({ includes: :links })

      rule.attribute :twitter, includes: :links
      expect(rule.attributes[:twitter]).to eq({ includes: :links })
    end
  end

  describe '#association' do
    it 'add associations' do
      rule.association :employees
      expect(rule.associations[:employees]).to eq({ includes: :employees, serializer: nil })
    end

    it 'add associations with includes' do
      rule.association :employees, includes: :users
      expect(rule.associations[:employees]).to eq({ includes: :users, serializer: nil })
    end

    it 'add associations with serializer name' do
      rule.association :employees, serializer: 'EmployeeSerializer'
      expect(rule.associations[:employees]).to eq({ includes: :employees, serializer: 'EmployeeSerializer' })
    end

    it 'add associations with serializer class' do
      rule.association :employees, serializer: EmployeeSerializer
      expect(rule.associations[:employees]).to eq({ includes: :employees, serializer: EmployeeSerializer })
    end
  end
end