From fce980fc47830ae088f2528ec98c9ce303d92346 Mon Sep 17 00:00:00 2001
From: Jim Menard <jim@10gen.com>
Date: Wed, 14 Jan 2009 10:42:56 -0500
Subject: [PATCH] Fixed $hint and $explain support.

---
 lib/mongo/cursor.rb                | 29 ++++++++++--------------
 lib/mongo/message/query_message.rb | 36 +++++++++++++++++++-----------
 lib/mongo/query.rb                 |  8 +++++++
 3 files changed, 43 insertions(+), 30 deletions(-)

diff --git a/lib/mongo/cursor.rb b/lib/mongo/cursor.rb
index c23362b..35f5011 100644
--- a/lib/mongo/cursor.rb
+++ b/lib/mongo/cursor.rb
@@ -125,13 +125,15 @@ module XGen
 
         # Returns an explain plan record.
         def explain
-          sel = OrderedHash.new
-          sel['query'] = @query.selector
-          sel['$explain'] = true
-          c = Cursor.new(@db, @collection, Query.new(sel))
-          e = c.next_object
+          old_val = @query.explain
+          @query.explain = true
+
+          c = Cursor.new(@db, @collection, @query)
+          explanation = c.next_object
           c.close
-          e
+
+          @query.explain = old_val
+          explanation
         end
 
         # Close the cursor.
@@ -210,18 +212,11 @@ module XGen
           # Run query first time we request an object from the wire
           unless @query_run
             hints = @hint_fields || @collection.hint_fields
-            query = if hints
-                      h = {}
-                      hints.each { |field| h[field] = 1 }
-                      sel = OrderedHash.new
-                      sel['query'] = @query.selector
-                      sel['$hint'] = h
-                      Query.new(sel)
-                    else
-                      @query
-                    end
-            @db.send_query_message(QueryMessage.new(@db.name, @collection.name, query))
+            old_hints = @query.hint_fields
+            @query.hint_fields = hints
+            @db.send_query_message(QueryMessage.new(@db.name, @collection.name, @query))
             @query_run = true
+            @query.hint_fields = old_hints
             read_all
           end
         end
diff --git a/lib/mongo/message/query_message.rb b/lib/mongo/message/query_message.rb
index b7e39c7..554c69b 100644
--- a/lib/mongo/message/query_message.rb
+++ b/lib/mongo/message/query_message.rb
@@ -18,21 +18,31 @@ module XGen
           write_int(query.number_to_skip)
           write_int(query.number_to_return)
           sel = query.selector
-          if query.order_by && query.order_by.length > 0
+          if query.contains_special_fields
             sel = OrderedHash.new
             sel['query'] = query.selector
-            sel['orderby'] = case query.order_by
-                             when String
-                               {query.order_by => 1}
-                             when Array
-                               h = OrderedHash.new
-                               query.order_by.each { |ob| h[ob] = 1 }
-                               h
-                             when Hash # Should be an ordered hash, but this message doesn't care
-                               query.order_by
-                             else
-                               raise "illegal order_by: is a #{query.order_by.class.name}, must be String, Array, Hash, or OrderedHash"
-                             end
+            if query.order_by && query.order_by.length > 0
+              sel['orderby'] = case query.order_by
+                               when String
+                                 {query.order_by => 1}
+                               when Array
+                                 h = OrderedHash.new
+                                 query.order_by.each { |ob| h[ob] = 1 }
+                                 h
+                               when Hash # Should be an ordered hash, but this message doesn't care
+                                 query.order_by
+                               else
+                                 raise "illegal order_by: is a #{query.order_by.class.name}, must be String, Array, Hash, or OrderedHash"
+                               end
+            end
+            if query.hint_fields && query.hint_fields.length > 0
+              hints = OrderedHash.new
+              query.hint_fields.each { |hf| hints[hf] = 1 }
+              sel['$hint'] = hints
+            end
+            if query.explain
+              sel['$explain'] = true
+            end
                                
           end
           write_doc(sel)
diff --git a/lib/mongo/query.rb b/lib/mongo/query.rb
index 48a913f..5c63172 100644
--- a/lib/mongo/query.rb
+++ b/lib/mongo/query.rb
@@ -27,6 +27,10 @@ module XGen
       class Query
 
         attr_accessor :number_to_skip, :number_to_return, :order_by
+        # If true, $explain will be set in QueryMessage that uses this query.
+        attr_accessor :explain
+        # Either +nil+ or an array of hint field names.
+        attr_accessor :hint_fields
         attr_reader :selector   # writer defined below
 
         # sel :: A hash describing the query. See the Mongo docs for details.
@@ -97,6 +101,10 @@ module XGen
             nil
           end
         end
+
+        def contains_special_fields
+          (@order_by != nil && @order_by.length > 0) || @explain || @hint_fields
+        end
       end
     end
   end