active_record_depth_query.rb

# 指定したレコードの関連の探索処理
#
# @example
#   depth_query = ActiveRecordDepthQuery.new(employee, [attendance: :attendance_records])
#   depth_query.each do |relation|
#     relation.to_a #=> 1度目は従業員に紐づくAttendanceの一覧、2度目はAttendanceRecordの一覧が返ってくる
#   end
class ActiveRecordDepthQuery
  include Enumerable

  # @param record [ActiveRecord::Base, Array<ActiveRecord::Base>] 探索の起点となるレコードを指定する
  # @param associations [Array, Hash] preloadのような指定方法で読み込む関連を指定する
  def initialize(record, associations)
    @record = record
    @associations = associations
  end

  # レコードを探索する
  #
  # @yield [relation]
  # @yieldparam relation [ActiveRecord::Relation] 探索されたリレーション
  # @yieldreturn [void]
  #
  # @return [void]
  def each(&block)
    Array.wrap(@associations).flat_map do
      records = Array.wrap(@record)
      loads_on(@record.class, _1, records, &block)
    end
  end

  private

  def loads_on(klass, association, records, &block)
    case association
    when Hash
      loads_for_hash(klass, association, records, &block)
    when Symbol, String
      loads_for_one(klass, association, records, &block)
    else
      raise ArgumentError, "#{association.inspect} was not recognized for import"
    end
  end

  def loads_for_one(klass, association, records, &block)
    reflection = find_reflection(klass, association)
    loads_for_reflection(reflection, records, &block)
  end

  def loads_for_hash(klass, association, records, &block)
    association.flat_map do |parent, children|
      reflection = find_reflection(klass, parent)

      loads_for_reflection(reflection, records) do |relation|
        block.call(relation) if block_given?

        # 子関連をさらに探索する
        Array.wrap(children).flat_map do |child_association|
          loads_on(relation.klass, child_association, relation, &block)
        end
      end
    end
  end

  def loads_for_reflection(reflection, records, &block)
    relation = loads_relation_by_reflection(reflection.klass, reflection, records)
    block.call(relation) if block_given?
  end

  # 関連の種類によって読み込み方を変える
  def loads_relation_by_reflection(klass, reflection, records)
    if reflection.through_reflection?
      # has_many(:xxx, through: :yyy)
      raise NotImplementedError, 'Unsupported through reflection'
    elsif reflection.collection? || reflection.has_one?
      # has_many(:xxx) or has_one(:xxx)
      primary_ids = optimized_key_collector(records, reflection.active_record_primary_key)
      klass.where(reflection.foreign_key => primary_ids)
    elsif reflection.belongs_to?
      primary_ids = optimized_key_collector(records, reflection.foreign_key)
      klass.where(reflection.active_record_primary_key => primary_ids)
    else
      # この条件に入る処理があればテスト・実装を追加する
      raise NotImplementedError, 'not implemented yet'
    end
  end

  def find_reflection(klass, association)
    klass.reflect_on_association(association).tap do
      raise("#{klass.name} has no association #{association}") if _1.nil?
    end
  end

  # ARをインスタンス化すると重いので、必要な値のみ取得できる場合はpluckを使う
  # 指定したカラムの値を集める
  #
  # @param records [ActiveRecord::Relation, Array<ActiveRecord::Base>]
  # @param column_name [#to_s]
  #
  # @return [Array]
  def optimized_key_collector(records, column_name)
    if records.is_a?(ActiveRecord::Relation)
      # 行数が大きすぎる場合は分割してリクエストしたほうがいいかもしれない
      records.pluck(column_name)
    else
      records.map { _1.public_send(column_name) }
    end
  end
end