ActiveRecord で関連レコードの size をプリロードする その2

前回のやつをその後いじってた。
動作確認は ActiveRecord-2.1.2
2.1.1 では動きません。(差分があります)

使い方

records = Record.find(:all, :include => :subrecords_count)
records[0].subrecords.size # カウント SQL が実行されない

active_record_sizes_preload.rb

require 'active_record'

module ActiveRecord
  module AssociationPreload::ClassMethods
    def preload_count_associations(records, associations, preload_options={})
      preload_associations(records, associations, preload_options.merge({:only_count => true}))
    end

  protected
    def preload_associations_with_count(records, associations, preload_options={})
      records = [records].flatten.compact.uniq
      return if records.empty?
      case associations
        when Symbol, String then
          if associations.to_s =~ /(.*)_count$/
            return preload_sizes(records, $1.to_sym, preload_options)
          end
          return if preload_options[:only_count]
      end
      preload_associations_without_count(records, associations, preload_options)
    end
    alias_method_chain :preload_associations, :count

  private
    def preload_sizes(records, association, preload_options)
      class_to_reflection = {}
      records.group_by {|record| class_to_reflection[record.class] ||= record.class.reflections[association]}.each do |reflection, records|
        raise ConfigurationError, "Association named '#{ association }' was not found; perhaps you misspelled it?" unless reflection

        sizes = {}
        ids = records.map {|record| record.id}
        conditions = "#{reflection.primary_key_name} IN (?)"
        conditions << append_conditions(reflection, preload_options)
        reflection.klass.count(:conditions => [conditions, ids],
                               :group => reflection.primary_key_name
                              ).map {|record| sizes[record[0]] = record[1]}

        records.each do |record|
          unless sizes[record.id]
            record.send(association).loaded
          else
            write_attribute = defined?(record.write_attribute_without_dirty) ? record.method(:write_attribute_without_dirty) : record.method(:write_attribtue)
            write_attribute.call("#{reflection.name}_count", sizes[record.id])
          end
        end
      end
    end

  end

  class Associations::ClassMethods::JoinDependency
    def instantiate_with_count(row)
      records = instantiate_without_count(row)
      records.first.class.preload_count_associations(records, @associations) unless records.empty?
      records
    end
    alias_method_chain :instantiate, :count

  protected
    def build_with_count(associations, parent = nil)
      case associations
        when Symbol, String
          return if associations.to_s =~ /_count$/
      end
      build_without_count(associations, parent)
    end
    alias_method_chain :build, :count

    def construct_with_count(parent, associations, joins, row)
      case associations
        when Symbol, String
          return if associations.to_s =~ /_count$/
      end
      construct_without_count(parent, associations, joins, row)
    end
    alias_method_chain :construct, :count

  end

end

ActiveRecord-2.1.1 以前用 patch。

--- lib/active_record_sizes_preload.rb
+++ lib/active_record_sizes_preload.rb
@@ -30,7 +30,7 @@
         sizes = {}
         ids = records.map {|record| record.id}
         conditions = "#{reflection.primary_key_name} IN (?)"
-        conditions << append_conditions(reflection.options, preload_options)
+        conditions << append_conditions(reflection, preload_options)
         reflection.klass.count(:conditions => [conditions, ids],
                                :group => reflection.primary_key_name
                               ).map {|record| sizes[record[0]] = record[1]}