Add counter caching.

lib/mongo_mapper.rb

@@ -36,6 +36,7 @@ module Plugins
     autoload :Callbacks,          'mongo_mapper/plugins/callbacks'
     autoload :Caching,            'mongo_mapper/plugins/caching'
     autoload :Clone,              'mongo_mapper/plugins/clone'
+    autoload :CounterCache,       'mongo_mapper/plugins/counter_cache'
     autoload :Dirty,              'mongo_mapper/plugins/dirty'
     autoload :Document,           'mongo_mapper/plugins/document'
     autoload :DynamicQuerying,    'mongo_mapper/plugins/dynamic_querying'

lib/mongo_mapper/document.rb

@@ -36,6 +36,7 @@ module Document
     include Plugins::EmbeddedCallbacks
     include Plugins::Callbacks # for now callbacks needs to be after validations
     include Plugins::IdentityMap
+    include Plugins::CounterCache
 
     included do
       extend Plugins

lib/mongo_mapper/plugins/counter_cache.rb

@@ -0,0 +1,65 @@
+module MongoMapper
+  module Plugins
+    # Counter Caching for MongoMapper::Document
+    #
+    # Examples:
+    #
+    #   class Post
+    #     belongs_to :user
+    #     counter_cache :user
+    #   end
+    #
+    #   or:
+    #
+    #   class Post
+    #     belongs_to :user
+    #     counter_cache :user, :custom_posts_count
+    #   end
+    #
+    # Field names follow rails conventions, so counter_cache :user will increment the Integer field `posts_count' on User
+    module CounterCache
+      class InvalidCounterCacheError < StandardError; end
+
+      extend ActiveSupport::Concern
+
+      module ClassMethods
+        def counter_cache(association_name, options = {})
+          options.symbolize_keys!
+
+          field = options[:field] ?
+            options[:field] :
+            "#{self.collection_name.gsub(/.*\./, '')}_count"
+
+          association = associations[association_name]
+
+          if !association
+            raise InvalidCounterCacheError, "You must define an association with name `#{association_name}' on model #{self}"
+          end
+
+          association_class = association.klass
+          key_names = association_class.keys.keys
+
+          if !key_names.include?(field.to_s)
+            raise InvalidCounterCacheError, "Missing `key #{field.to_sym.inspect}, Integer, :default => 0' on model #{association_class}"
+          end
+
+          after_create do
+            if obj = self.send(association_name)
+              obj.increment(field => 1)
+              obj.write_attribute(field, obj.read_attribute(field) + 1)
+            end
+            true
+          end
+
+          after_destroy do
+            if obj = self.send(association_name)
+              obj.decrement(field => 1)
+              obj.write_attribute(field, obj.read_attribute(field) - 1)
+            end
+            true
+          end
+        end
+      end
+    end
+  end
+end

spec/functional/counter_cache_spec.rb

@@ -0,0 +1,149 @@
+require 'spec_helper'
+
+module CounterCacheFixtureModels
+  class User
+    include MongoMapper::Document
+
+    key :posts_count, Integer, :default => 0
+
+    has_many :posts,
+             :class_name => "CounterCacheFixtureModels::Post"
+  end
+
+  class Post
+    include MongoMapper::Document
+
+    key :comments_count, Integer, :default => 0
+    key :some_custom_comments_count, Integer, :default => 0
+
+    has_many :comments,
+             :class_name => "CounterCacheFixtureModels::Comment"
+
+    belongs_to :user,
+               :class_name => "CounterCacheFixtureModels::User"
+
+    counter_cache :user
+  end
+
+  class Comment
+    include MongoMapper::Document
+
+    belongs_to :post,
+               :class_name => "CounterCacheFixtureModels::Post"
+
+    counter_cache :post
+  end
+
+  class CustomComment
+    include MongoMapper::Document
+
+    belongs_to :post,
+               :class_name => "CounterCacheFixtureModels::Post"
+
+    counter_cache :post, :field => :some_custom_comments_count
+  end
+end
+
+describe MongoMapper::Plugins::CounterCache do
+  before do
+    @post_class = CounterCacheFixtureModels::Post
+    @comment_class = CounterCacheFixtureModels::Comment
+    @user_class = CounterCacheFixtureModels::User
+    @custom_comment_class = CounterCacheFixtureModels::CustomComment
+  end
+
+  it "should have a key with posts_count defaulting to 0" do
+    @post_class.new.comments_count.should == 0
+  end
+
+  it "should update the count when a new object is created" do
+    post = @post_class.new
+    comment = @comment_class.new
+
+    post.save!
+
+    comment.post = post
+    comment.save!
+
+    post.reload
+    post.comments_count.should == 1
+
+    second_comment = @comment_class.new
+    second_comment.post = post
+    second_comment.save!
+
+    post.reload
+    post.comments_count.should == 2
+  end
+
+  it "should decrease the count by one when an object is destroyed" do
+    post = @post_class.new
+    comment = @comment_class.new
+
+    post.save!
+
+    comment.post = post
+    comment.save!
+
+    post.reload
+    post.comments_count.should == 1
+
+    comment.destroy
+    post.reload
+    post.comments_count.should == 0
+  end
+
+  it "should use the correct association name" do
+    @user = @user_class.new
+    @post = @post_class.new
+
+    @user.save!
+    @post.user = @user
+    @post.save!
+
+    @user.reload
+    @user.posts_count.should == 1
+  end
+
+  it "should be able to use a custom field name" do
+    @post = @post_class.new
+    @custom_comment = @custom_comment_class.new
+
+    @post.save!
+    @custom_comment.post = @post
+    @custom_comment.save!
+
+    @post.reload
+    @post.some_custom_comments_count.should == 1
+  end
+
+  it "should thrown an error if there is no association" do
+    lambda {
+      CounterCacheFixtureModels.module_eval do
+        class CommentWithInvalidAssociation
+          include MongoMapper::Document
+
+          belongs_to :post,
+                     :class_name => "CounterCacheFixtureModels::Post"
+
+          counter_cache :foo
+        end
+      end
+    }.should raise_error(MongoMapper::Plugins::CounterCache::InvalidCounterCacheError, "You must define an association with name `foo' on model CommentWithInvalidAssociation")
+  end
+
+  it "should thown a sensible error if the field is not defined on the target object" do
+    lambda {
+      CounterCacheFixtureModels.module_eval do
+        class CommentWithBadRefenceField
+          include MongoMapper::Document
+
+          belongs_to :post,
+                     :class_name => "CounterCacheFixtureModels::Post"
+
+          counter_cache :post, :field => :invalid_field
+        end
+      end
+    }.should raise_error(MongoMapper::Plugins::CounterCache::InvalidCounterCacheError, "Missing `key :invalid_field, Integer, :default => 0' on model CounterCacheFixtureModels::Post")
+  end
+end