Use class name provided by association AND support self-referential counts

lib/preload_counts/ar.rb

@@ -1,9 +1,9 @@
-# This adds a scope to preload the counts of an association in one SQL query. 
+# This adds a scope to preload the counts of an association in one SQL query.
 #
 # Consider the following code:
 # Service.all.each{|s| puts s.incidents.acknowledged.count}
 #
-# Each time count is called, a db query is made to fetch the count. 
+# Each time count is called, a db query is made to fetch the count.
 #
 # Adding this to the Service class:
 #
@@ -43,7 +43,7 @@ def preload_counts(options)
 
       end
     end
-    
+
     private
     def scopes_to_select(association, scopes)
       scopes.map do |scope|
@@ -52,7 +52,10 @@ def scopes_to_select(association, scopes)
     end
 
     def scope_to_select(association, scope)
-      resolved_association = association.to_s.singularize.camelize.constantize
+      association_reflection = reflect_on_association(association)
+      resolved_association = association_reflection.class_name.constantize
+      association_table_name = association_reflection.table_name
+
       conditions = []
 
       if scope
@@ -68,14 +71,26 @@ def scope_to_select(association, scope)
       association_condition = self.reflections[association].options[:conditions]
       conditions << association_condition if association_condition
 
+      foreign_key = association_reflection.foreign_key
+
+      # Get a unique table alias, which we need for preloading counts of a
+      # self-referential relationship
+      @count ||= 0
+      association_table_alias = "#{association_table_name}_#{@count}"
+
+      join_column = "#{association_table_alias}.#{foreign_key}"
+
       # FIXME This is a really hacking way of getting the named_scope condition.
-      # In Rails 3 we would have AREL to get to it. 
+      # In Rails 3 we would have AREL to get to it.
       sql = <<-SQL
-      (SELECT count(*) 
-       FROM #{association} 
-       WHERE #{association}.#{table_name.singularize}_id = #{table_name}.id AND 
-       #{conditions_to_sql conditions}) as #{find_accessor_name(association, scope)} 
+      (SELECT count(*)
+       FROM #{association_table_name} AS #{association_table_alias}
+       WHERE #{join_column} = #{table_name}.id AND
+       #{conditions_to_sql conditions}) as #{find_accessor_name(association, scope)}
       SQL
+
+      @count += 1
+      sql
     end
 
     def find_accessor_name(association, scope)
@@ -100,4 +115,3 @@ def self.included(receiver)
 end
 
 ActiveRecord::Base.class_eval { include PreloadCounts }
-

spec/active_record_spec.rb

@@ -11,28 +11,42 @@
 def setup_db
   ActiveRecord::Schema.define(:version => 1) do
     create_table :posts do |t|
+      t.integer :parent_id
     end
 
     create_table :comments do |t|
-      t.integer :post_id, :null => false 
+      t.integer :post_id, :null => false
       t.datetime :deleted_at
     end
+
+    create_table :moduled_other_comments do |t|
+      t.integer :post_id, :null => false
+    end
   end
 end
 
-setup_db 
+setup_db
 
 class Post < ActiveRecord::Base
+  has_many :children, :class_name => 'Post', :foreign_key => 'parent_id'
+
   has_many :comments
-  has_many :active_comments, :conditions => "deleted_at IS NULL", :class_name => 'Comment'
+  has_many :active_comments, :conditions => "deleted_at IS NULL",
+    :class_name => 'Comment'
+
+  has_many :moduled_other_comments, :class_name => 'Moduled::OtherComment'
+
+  preload_counts :children
   preload_counts :comments => [:with_even_id]
   preload_counts :active_comments
+  preload_counts :moduled_other_comments
 end
 
 class PostWithActiveComments < ActiveRecord::Base
   set_table_name :posts
 
-  has_many :comments, :conditions => "deleted_at IS NULL"
+  has_many :comments, :conditions => "deleted_at IS NULL",
+    :foreign_key => 'post_id'
   preload_counts :comments
 end
 
@@ -46,37 +60,51 @@ class Comment < ActiveRecord::Base
   end
 end
 
+module Moduled
+end
+class Moduled::OtherComment < ActiveRecord::Base
+  set_table_name 'moduled_other_comments'
+  belongs_to :post
+end
+
 def create_data
-  post = Post.create 
+  post = Post.create
+  3.times { post.children.create }
   5.times { post.comments.create }
   5.times { post.comments.create :deleted_at => Time.now }
+
+  2.times { post.moduled_other_comments.create }
 end
 
 create_data
 
 describe Post do
   it "should have a preload_comment_counts scope" do
-    Post.should respond_to(:preload_comment_counts) 
+    Post.should respond_to(:preload_comment_counts)
+  end
+
+  it "should have a preload_moduled_comment_counts scope" do
+    Post.should respond_to(:preload_moduled_other_comment_counts)
   end
 
   describe 'instance' do
-    let(:post) { Post.first } 
+    let(:post) { Post.first }
 
     it "should have a comment_count accessor" do
-      post.should respond_to(:comments_count) 
+      post.should respond_to(:comments_count)
     end
 
     it "should be able to get count without preloading them" do
       post.comments_count.should equal(10)
     end
 
     it "should have an active_comments_count accessor" do
-      post.should respond_to(:comments_count) 
+      post.should respond_to(:comments_count)
     end
   end
 
   describe 'instance with preloaded count' do
-    let(:post) { Post.preload_comment_counts.first } 
+    let(:post) { Post.preload_comment_counts.first }
 
     it "should be able to get the association count" do
       post.comments_count.should equal(10)
@@ -86,11 +114,27 @@ def create_data
       post.with_even_id_comments_count.should equal(5)
     end
   end
+
+  describe 'instance with preloaded moduled comment count' do
+    let(:post) { Post.preload_moduled_other_comment_counts.first }
+
+    it "should be able to get the moduled association count" do
+      post.moduled_other_comments_count.should equal(2)
+    end
+  end
+
+  describe 'instance with preloaded self-referential count' do
+    let(:post) { Post.preload_child_counts.first }
+
+    it "should be able to get the self-referential association count" do
+      post.children_count.should equal(3)
+    end
+  end
 end
 
 describe PostWithActiveComments do
   describe 'instance with preloaded count' do
-    let(:post) { PostWithActiveComments.preload_comment_counts.first } 
+    let(:post) { PostWithActiveComments.preload_comment_counts.first }
 
     it "should be able to get the association count" do
       post.comments_count.should equal(5)