Making .cache on context more consistent

[Fix #2667]

lib/mongoid/contextual/mongo.rb

@@ -56,7 +56,7 @@ def count(document = false, &block)
           return collection.find(criteria.and(_id: document.id).selector).count
         end
         return query.count(document) if document
-        cached? ? @count ||= query.count : query.count
+        try_cache(:count) { query.count }
       end
 
       # Delete all documents in the database that match the selector.
@@ -143,7 +143,12 @@ def each(&block)
       #
       # @since 3.0.0
       def exists?
-        @exists ||= check_existence
+        return !documents.empty? if cached? && cache_loaded?
+        return @count > 0 if instance_variable_defined?(:@count)
+
+        try_cache(:exists) do
+          !!(query.dup.select(_id: 1).limit(1).first)
+        end
       end
 
       # Run an explain on the criteria.
@@ -189,9 +194,8 @@ def find_and_modify(update, options = {})
       #
       # @since 3.0.0
       def first
-        if cached? && cache_loaded?
-          documents.first
-        else
+        return documents.first if cached? && cache_loaded?
+        try_cache(:first) do
           with_sorting do
             with_eager_loading(query.first)
           end
@@ -274,8 +278,10 @@ def initialize(criteria)
       #
       # @since 3.0.0
       def last
-        with_inverse_sorting do
-          with_eager_loading(query.first)
+        try_cache(:last) do
+          with_inverse_sorting do
+            with_eager_loading(query.first)
+          end
         end
       end
 
@@ -406,21 +412,21 @@ def update_all(attributes = nil)
 
       private
 
-      # Checks if any documents exist in the database.
+      # yield the block given or return the cached value
       #
-      # @api private
-      #
-      # @example Check for document existsence.
-      #   context.check_existence
+      # @param [ String, Symbol ] key The instance variable name
       #
-      # @return [ true, false ] If documents exist.
+      # @return the result of the block
       #
-      # @since 3.1.0
-      def check_existence
-        if cached? && cache_loaded?
-          !documents.empty?
+      # @since 3.1.4
+      def try_cache(key, &block)
+        unless cached?
+          yield
         else
-          @count ? @count > 0 : !query.dup.select(_id: 1).limit(1).entries.first.nil?
+          unless ret = instance_variable_get("@#{key}")
+            instance_variable_set("@#{key}", ret = yield)
+          end
+          ret
         end
       end
 

spec/mongoid/contextual/mongo_spec.rb

@@ -480,6 +480,29 @@
       end
     end
 
+    context "when caching is not enabled" do
+
+      let(:criteria) do
+        Band.where(name: "Depeche Mode")
+      end
+
+      let(:context) do
+        described_class.new(criteria)
+      end
+
+      context "when exists? already called" do
+
+        before do
+          context.exists?
+        end
+
+        it "hits the database again" do
+          context.should_receive(:query).once.and_call_original
+          context.should be_exists
+        end
+      end
+    end
+
     context "when caching is enabled" do
 
       let(:criteria) do
@@ -788,6 +811,18 @@
             expect(context.send(method)).to eq(depeche_mode)
           end
         end
+
+        context "when first method was called before" do
+
+          before do
+            context.first
+          end
+
+          it "returns the first document without touching the database" do
+            context.should_receive(:query).never
+            context.send(method).should eq(depeche_mode)
+          end
+        end
       end
     end
   end