From 265f074f10db968bc6d5d911bbfad857b343f88b Mon Sep 17 00:00:00 2001
From: Kyle Banker <kylebanker@gmail.com>
Date: Mon, 12 Jul 2010 12:11:01 -0400
Subject: [PATCH] RUBY-147 check for CursorNotFound response flag

---
 lib/mongo.rb            | 15 +++++++++++----
 lib/mongo/collection.rb |  3 ++-
 lib/mongo/connection.rb | 11 ++++++++++-
 lib/mongo/cursor.rb     |  4 ++--
 test/cursor_test.rb     | 19 ++++++++++++++++++-
 5 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/lib/mongo.rb b/lib/mongo.rb
index 6d75e78..7ce1807 100644
--- a/lib/mongo.rb
+++ b/lib/mongo.rb
@@ -21,11 +21,18 @@ module Mongo
     OP_DELETE       = 2006
     OP_KILL_CURSORS = 2007
 
-    OP_QUERY_TAILABLE          = 2
-    OP_QUERY_SLAVE_OK          = 4
-    OP_QUERY_NO_CURSOR_TIMEOUT = 16
-  end
+    OP_QUERY_TAILABLE          = 2 ** 1
+    OP_QUERY_SLAVE_OK          = 2 ** 2
+    OP_QUERY_OPLOG_REPLAY      = 2 ** 3
+    OP_QUERY_NO_CURSOR_TIMEOUT = 2 ** 4
+    OP_QUERY_AWAIT_DATA        = 2 ** 5
+    OP_QUERY_EXHAUST           = 2 ** 6
 
+    REPLY_CURSOR_NOT_FOUND     = 2 ** 0
+    REPLY_QUERY_FAILURE        = 2 ** 1
+    REPLY_SHARD_CONFIG_STALE   = 2 ** 2
+    REPLY_AWAIT_CAPABLE        = 2 ** 3
+  end
 end
 
 require 'bson'
diff --git a/lib/mongo/collection.rb b/lib/mongo/collection.rb
index 740d624..7cb5d57 100644
--- a/lib/mongo/collection.rb
+++ b/lib/mongo/collection.rb
@@ -157,7 +157,8 @@ module Mongo
       raise RuntimeError, "Unknown options [#{opts.inspect}]" unless opts.empty?
 
       cursor = Cursor.new(self, :selector => selector, :fields => fields, :skip => skip, :limit => limit,
-                          :order => sort, :hint => hint, :snapshot => snapshot, :timeout => timeout, :batch_size => batch_size)
+        :order => sort, :hint => hint, :snapshot => snapshot, :timeout => timeout, :batch_size => batch_size)
+
       if block_given?
         yield cursor
         cursor.close()
diff --git a/lib/mongo/connection.rb b/lib/mongo/connection.rb
index 013ed29..ea03905 100644
--- a/lib/mongo/connection.rb
+++ b/lib/mongo/connection.rb
@@ -632,13 +632,22 @@ module Mongo
           "expected #{RESPONSE_HEADER_SIZE} bytes, saw #{header_buf.length}"
       end
       header_buf.rewind
-      result_flags     = header_buf.get_int
+      check_response_flags(header_buf.get_int)
       cursor_id        = header_buf.get_long
       starting_from    = header_buf.get_int
       number_remaining = header_buf.get_int
       [number_remaining, cursor_id]
     end
 
+    def check_response_flags(flags)
+      if flags & Mongo::Constants::REPLY_CURSOR_NOT_FOUND != 0
+        raise Mongo::OperationFailure, "Query response returned CURSOR_NOT_FOUND. " +
+          "Either an invalid cursor was specified, or the cursor may have timed out on the server."
+      elsif flags & Mongo::Constants::REPLY_QUERY_FAILURE != 0
+        raise Mongo::OperationFailure, "Query response returned QUERY_FAILURE."
+      end
+    end
+
     def read_documents(number_received, cursor_id, sock)
       docs = []
       number_remaining = number_received
diff --git a/lib/mongo/cursor.rb b/lib/mongo/cursor.rb
index 13f9977..5e66fb5 100644
--- a/lib/mongo/cursor.rb
+++ b/lib/mongo/cursor.rb
@@ -23,7 +23,7 @@ module Mongo
 
     attr_reader :collection, :selector, :admin, :fields,
       :order, :hint, :snapshot, :timeout,
-      :full_collection_name
+      :full_collection_name, :batch_size
 
     # Create a new cursor.
     #
@@ -237,7 +237,7 @@ module Mongo
         message = BSON::ByteBuffer.new([0, 0, 0, 0])
         message.put_int(1)
         message.put_long(@cursor_id)
-        @connection.send_message(Mongo::Constants::OP_KILL_CURSORS, message, "cursor.close")
+        @connection.send_message(Mongo::Constants::OP_KILL_CURSORS, message, "cursor.close #{@cursor_id}")
       end
       @cursor_id = 0
       @closed    = true
diff --git a/test/cursor_test.rb b/test/cursor_test.rb
index be5e317..4aab1ad 100644
--- a/test/cursor_test.rb
+++ b/test/cursor_test.rb
@@ -137,7 +137,7 @@ class CursorTest < Test::Unit::TestCase
     @@coll.save({:t => 't2'})
     @@coll.save({:t => 't2'})
 
-    assert_equal 3, @@coll.find({'_id' => {'$gt' => t1_id}, '_id' => {'$lt' => t2_id}}).count
+    assert_equal 3, @@coll.find({'_id' => {'$gt' => t1_id, '$lt' => t2_id}}).count
     @@coll.find({'_id' => {'$gt' => t2_id}}).each do |doc|
       assert_equal 't2', doc['t']
     end
@@ -396,4 +396,21 @@ class CursorTest < Test::Unit::TestCase
 
     assert_equal false, cursor.has_next?
   end
+
+  def test_cursor_invalid
+    @@coll.remove
+    1000.times do |n|
+      @@coll.insert({:a => n})
+    end
+
+    cursor = @@coll.find({})
+    cursor.next_document
+    cursor.close
+
+    assert_raise_error(Mongo::OperationFailure, "CURSOR_NOT_FOUND") do
+      999.times do
+        cursor.next_document
+      end
+    end
+  end
 end