From 2d854a0b518459c841d292ec6f3ca03d43e789ef Mon Sep 17 00:00:00 2001
From: Frank Macreery <frank@macreery.com>
Date: Wed, 27 Nov 2013 11:47:47 -0500
Subject: [PATCH] Add @cache to Moped::Cursor for out-of-block iteration

The Mongo Ruby driver (and thus Mongoid 2.x) supported iterating over a
cursor outside a block, using such operations as take(). For example:

    batch_size = 10
    cursor = session[:users].find({}).cursor
    while (users = cursor.take(batch_size)).present?
      # Do stuff...
    end
    # Do more stuff...

This was a convenient pattern for implementing lots of algorithms in a
more readable way than if they were written using each().
---
 lib/moped/cursor.rb       | 10 +++++-----
 spec/moped/cursor_spec.rb | 20 ++++++++++++++++++++
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/lib/moped/cursor.rb b/lib/moped/cursor.rb
index 83f2693..e4cda99 100644
--- a/lib/moped/cursor.rb
+++ b/lib/moped/cursor.rb
@@ -11,7 +11,7 @@ class Cursor
     # @attribute [r] kill_cursor_op The kill cursor message.
     # @attribute [r] query_op The query message.
     # @attribute [r] session The session.
-    attr_reader :get_more_op, :kill_cursor_op, :query_op, :session
+    attr_reader :get_more_op, :kill_cursor_op, :query_op, :session, :cache
 
     # Iterate over the results of the query.
     #
@@ -24,12 +24,12 @@ class Cursor
     #
     # @since 1.0.0
     def each
-      documents = load_docs
-      documents.each { |doc| yield doc }
+      @cache ||= load_docs
+      yield @cache.shift while @cache.any?
       while more?
         return kill if limited? && @limit <= 0
-        documents = get_more
-        documents.each { |doc| yield doc }
+        @cache = get_more
+        yield @cache.shift while @cache.any?
       end
     end
 
diff --git a/spec/moped/cursor_spec.rb b/spec/moped/cursor_spec.rb
index 835db1a..717cb89 100644
--- a/spec/moped/cursor_spec.rb
+++ b/spec/moped/cursor_spec.rb
@@ -67,5 +67,25 @@
         cursor.request_limit.should eq(100)
       end
     end
+
+    context "when the cursor is iterated upon out-of-block" do
+
+      before do
+        session[:users].insert({ "name" => "create" })
+      end
+
+      let(:query) do
+        session[:users].find.limit(1)
+      end
+
+      let(:cursor) do
+        described_class.new(session, query.operation)
+      end
+
+      it "advances the cursor_id" do
+        cursor.take(1)
+        cursor.take(1).should be_empty
+      end
+    end
   end
 end