From c397758d1d752df6bf1b03408cc068d5a657c5a7 Mon Sep 17 00:00:00 2001
From: Kyle Banker <kylebanker@gmail.com>
Date: Sun, 6 Sep 2009 02:25:49 +0800
Subject: [PATCH] Added #limit and #offset methods for cursor objects

Signed-off-by: Mike Dirolf <mike@dirolf.com>
---
 lib/mongo/cursor.rb | 30 ++++++++++++++++++
 test/test_cursor.rb | 76 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 106 insertions(+)

diff --git a/lib/mongo/cursor.rb b/lib/mongo/cursor.rb
index a1e327d..355fbbb 100644
--- a/lib/mongo/cursor.rb
+++ b/lib/mongo/cursor.rb
@@ -74,6 +74,30 @@ module Mongo
       raise OperationFailure, "Count failed: #{response['errmsg']}"
     end
 
+    # Sets a limit on the number of results returned by the query.
+    # Returns a cursor object.
+    #
+    # Note: this method overrides any limit specified in the #find method.
+    def limit(number_to_return)
+      check_modifiable
+      raise ArgumentError, "limit requires an integer" unless number_to_return.is_a? Integer
+
+      @number_to_return       = number_to_return
+      @query.number_to_return = number_to_return
+      return self
+    end
+    
+    # Sets an offset on the query results. Returns a cursor object.
+    #
+    # Note: this method overrides any offset specified in the #find method.
+    def offset(number_to_skip)
+      check_modifiable
+      raise ArgumentError, "limit requires an integer" unless number_to_skip.is_a? Integer
+
+      @query.number_to_skip = number_to_skip
+      return self
+    end
+
     # Iterate over each document in this cursor, yielding it to the given
     # block.
     #
@@ -226,5 +250,11 @@ module Mongo
     def to_s
       "DBResponse(flags=#@result_flags, cursor_id=#@cursor_id, start=#@starting_from, n_returned=#@n_returned)"
     end
+
+    def check_modifiable
+      if @query_run || @closed
+        raise InvalidOperation, "Cannot modify the query once it has been run or closed."
+      end
+    end
   end
 end
diff --git a/test/test_cursor.rb b/test/test_cursor.rb
index f879013..dece653 100644
--- a/test/test_cursor.rb
+++ b/test/test_cursor.rb
@@ -58,6 +58,82 @@ class CursorTest < Test::Unit::TestCase
     assert_equal 0, @@db['acollectionthatdoesn'].count()
   end
 
+  def test_limit
+    @@coll.clear
+
+    10.times do |i|
+      @@coll.save("x" => i)
+    end
+    assert_equal 10, @@coll.find().count()
+
+    results = @@coll.find().limit(5).to_a
+    assert_equal 5, results.length
+  end
+
+  def test_limit_exceptions
+    assert_raise ArgumentError do 
+      cursor = @@coll.find().limit('not-an-integer')
+    end
+
+    cursor      = @@coll.find()
+    firstResult = cursor.next_object()
+    assert_raise InvalidOperation, "Cannot modify the query once it has been run or closed." do 
+      cursor.limit(1)
+    end
+
+    cursor = @@coll.find()
+    cursor.close
+    assert_raise InvalidOperation, "Cannot modify the query once it has been run or closed." do 
+      cursor.limit(1)
+    end
+  end
+
+  def test_offset
+    @@coll.clear
+
+    10.times do |i|
+      @@coll.save("x" => i)
+    end
+    assert_equal 10, @@coll.find().count()
+
+    all_results    = @@coll.find().to_a
+    offset_results = @@coll.find().offset(2).to_a
+    assert_equal 10, all_results.length
+    assert_equal 8,  offset_results.length
+
+    assert_equal all_results.slice(2...10), offset_results
+  end
+
+  def test_offset_exceptions
+    assert_raise ArgumentError do 
+      cursor = @@coll.find().offset('not-an-integer')
+    end
+
+    cursor      = @@coll.find()
+    firstResult = cursor.next_object()
+    assert_raise InvalidOperation, "Cannot modify the query once it has been run or closed." do 
+      cursor.offset(1)
+    end
+
+    cursor = @@coll.find()
+    cursor.close
+    assert_raise InvalidOperation, "Cannot modify the query once it has been run or closed." do 
+      cursor.offset(1)
+    end
+  end
+
+  def test_limit_offset_chaining
+    @@coll.clear
+    10.times do |i|
+      @@coll.save("x" => i)
+    end
+    
+    all_results = @@coll.find().to_a
+    limited_offset_results = @@coll.find().limit(5).offset(3).to_a
+    
+    assert_equal all_results.slice(3...8), limited_offset_results
+  end
+
   def test_close_no_query_sent
     begin
       cursor = @@coll.find('a' => 1)