ActiveRecord association preloading that works around some of the ActiveRecord::Associations::Preloader limitations

preloader.rb

# Preloads ActiveRecord model associations with shared object references.
# This internally uses ActiveRecord::Associations::Preloader but works around some of its limitations.
# It works around a bug with associations containing loaded/unloaded models. https://github.com/rails/rails/issues/32140
#
# Usage:
#   Preloader.new(models, association_spec).preload
# Also supports a third `scope` argument that is used when querying.
#
# There is an additional class utility method `uniquify_associations` to share object references that can be useful in preloading.
class Preloader
  attr_reader :models, :association_spec, :scope

  class << self
    # Utility to share object references among multiple models.
    #  This may not be possible from associations from different classes.
    # Ex: uniquify_associations(books, :authors, authors)
    def uniquify_associations(source_models, association, associated_models)
      raise ArgumentError.new('associated_models must all be unique') if associated_models != associated_models.uniq

      # Note: Rails models of the same class/ID have the same hash key so can be looked up by the AR model
      associated_models_map = associated_models.index_by(&:itself)

      Array.wrap(source_models).each do |source_model|
        uniquify_association(source_model, association, associated_models_map)
      end
    end

    private

    def uniquify_association(source_model, association, associated_models_map)
      association_value = source_model.send(association)

      new_value = begin
        if association_value.is_a?(ActiveRecord::Base)
          associated_models_map.fetch(association_value)
        elsif association_value.present?
          association_value.map do |model|
            associated_models_map.fetch(model)
          end
        else
          association_value
        end
      end

      source_model.send("#{association}=", new_value)
    end
  end

  def initialize(models, association_spec, scope=nil)
    @original_models = Array.wrap(models)
    @root_associations = []
    @models = @original_models.uniq
    @association_spec = association_spec
    @scope = scope
  end

  def preload
    traverse_and_preload(models, association_spec, true)

    if @original_models.size != models.size
      models_map = models.index_by(&:itself)
      @original_models.each do |model|
        preloaded_model = models_map[model]

        next if preloaded_model.equal?(model)

        copy_associations(preloaded_model, model)
      end
    end

    nil
  end

  private

  def traverse_and_preload(models, association_spec, root_association)
    case association_spec
    when Symbol
      perform_preload(models, association_spec, root_association)
    when Array
      association_spec.each do |association_spec_element|
        traverse_and_preload(models, association_spec_element, root_association)
      end
    when Hash
      association_spec.each do |association_key, association_spec_element|
        perform_preload(models, association_key, root_association)

        traverse_and_preload(association_array(models, association_key), association_spec_element, false)
      end
    else
      raise ArgumentError.new("Unknown association specification: #{association_spec}")
    end
  end

  def association_array(models, association)
    models.each_with_object([]) do |model, result|
      association_value = model.send(association)
      case association_value
      when Enumerable
        result.concat(association_value)
      when NilClass
        nil
      else
        result.push(association_value)
      end
    end.uniq.tap do |associated_models|
      models.each do |model|
        self.class.uniquify_associations(model, association, associated_models)
      end
    end
  end

  def perform_preload(models, association, root_association)
    raise ArgumentError unless association.is_a?(Symbol)

    @root_associations << association if root_association

    ActiveRecord::Associations::Preloader.new.preload(
      models.reject { |m| m.association(association).loaded? },
      association,
      scope,
    )
  end

  def copy_associations(source_model, destination_model)
    @root_associations.each do |association|
      destination_model.send("#{association}=", source_model.send(association))
    end
  end
end