RUBY-171 allow Cursor#to_a even after iterating; added Cursor#rewind; consistent Enumberable behavior for Cursor

Author: banker

lib/mongo/cursor.rb

@@ -81,6 +81,18 @@ def next_document
       doc
     end
 
+    # Reset this cursor on the server. Cursor options, such as the
+    # query string and the values for skip and limit, are preserved.
+    def rewind!
+      close
+      @cache.clear
+      @cursor_id  = nil
+      @closed     = false
+      @query_run  = false
+      @n_received = nil
+      true
+    end
+
     # Determine whether this cursor has any remaining results.
     #
     # @return [Boolean]
@@ -187,22 +199,18 @@ def each
 
     # Receive all the documents from this cursor as an array of hashes.
     #
-    # Note: use of this method is discouraged - in most cases, it's much more
+    # Notes:
+    #
+    # If you've already started iterating over the cursor, the array returned
+    # by this method contains only the remaining documents. See Cursor#rewind! if you
+    # need to reset the cursor.
+    #
+    # Use of this method is discouraged - in most cases, it's much more
     # efficient to retrieve documents as you need them by iterating over the cursor.
     #
     # @return [Array] an array of documents.
-    #
-    # @raise [InvalidOperation] if this cursor has already been used or if
-    #   this method has already been called on the cursor.
     def to_a
-      raise InvalidOperation, "can't call Cursor#to_a on a used cursor" if @query_run
-      rows = []
-      num_returned = 0
-      while has_next? && (@limit <= 0 || num_returned < @limit)
-        rows << next_document
-        num_returned += 1
-      end
-      rows
+      super
     end
 
     # Get the explain plan for this cursor.

test/cursor_test.rb

@@ -396,4 +396,36 @@ def test_cursor_invalid
     end
   end
 
+  def test_enumberables
+    @@coll.remove
+    100.times do |n|
+      @@coll.insert({:a => n})
+    end
+
+    assert_equal 100, @@coll.find.to_a.length
+    assert_equal 100, @@coll.find.to_set.length
+
+    cursor = @@coll.find
+    50.times { |n| cursor.next_document }
+    assert_equal 50, cursor.to_a.length
+  end
+
+  def test_rewind
+    @@coll.remove
+    100.times do |n|
+      @@coll.insert({:a => n})
+    end
+
+    cursor = @@coll.find
+    cursor.to_a
+    assert_equal [], cursor.map {|doc| doc }
+
+    cursor.rewind!
+    assert_equal 100, cursor.map {|doc| doc }.length
+
+    cursor.rewind!
+    5.times { cursor.next_document }
+    cursor.rewind!
+    assert_equal 100, cursor.map {|doc| doc }.length
+  end
 end

test/db_api_test.rb

@@ -445,25 +445,6 @@ def test_strict_create_collection
     @@db.drop_collection('foobar')
   end
 
-  def test_to_a
-    cursor = @@coll.find()
-    rows = cursor.to_a
-
-    assert_raise InvalidOperation do
-      cursor.to_a
-    end
-
-    cursor.each { |doc| fail "should be no docs in each now" }
-  end
-
-  def test_to_a_after_each
-    cursor = @@coll.find
-    cursor.each { |row| row }
-    assert_raise InvalidOperation do
-      cursor.to_a
-    end
-  end
-
   def test_where
     @@coll.insert('a' => 2)
     @@coll.insert('a' => 3)