From 84c8e7cae241ec6050609347c533f09f304e39c1 Mon Sep 17 00:00:00 2001
From: Cyril Mougel <cyril.mougel@gmail.com>
Date: Fri, 3 Feb 2012 18:02:14 +0100
Subject: [PATCH] add log_duration system to add time spend in database

---
 lib/mongo/collection.rb   | 11 ++++++-----
 lib/mongo/cursor.rb       |  3 ++-
 lib/mongo/util/logging.rb | 11 ++++++++---
 test/connection_test.rb   |  9 +++++++++
 4 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/lib/mongo/collection.rb b/lib/mongo/collection.rb
index bbc5b72..9bdc14a 100644
--- a/lib/mongo/collection.rb
+++ b/lib/mongo/collection.rb
@@ -87,6 +87,7 @@ module Mongo
       @db, @name  = db, name
       @connection = @db.connection
       @logger     = @connection.logger
+      @log_duration = @connection.log_duration
       @cache_time = @db.cache_time
       @cache = Hash.new(0)
       unless pk_factory
@@ -264,8 +265,8 @@ module Mongo
     # @return [OrderedHash, Nil]
     #   a single document or nil if no result is found.
     #
-    # @param [Hash, ObjectId, Nil] spec_or_object_id a hash specifying elements 
-    #   which must be present for a document to be included in the result set or an 
+    # @param [Hash, ObjectId, Nil] spec_or_object_id a hash specifying elements
+    #   which must be present for a document to be included in the result set or an
     #   instance of ObjectId to be used as the value for an _id query.
     #   If nil, an empty selector, {}, will be used.
     #
@@ -413,7 +414,7 @@ module Mongo
     # @option opts [Boolean] :upsert (+false+) if true, performs an upsert (update or insert)
     # @option opts [Boolean] :multi (+false+) update all documents matching the selector, as opposed to
     #   just the first matching document. Note: only works in MongoDB 1.1.3 or later.
-    # @option opts [Boolean] :safe (+false+) 
+    # @option opts [Boolean] :safe (+false+)
     #   If true, check that the save succeeded. OperationFailure
     #   will be raised on an error. Note that a safe check requires an extra
     #   round-trip to the database. Safe options provided here will override any safe
@@ -906,12 +907,12 @@ module Mongo
           if [Mongo::ASCENDING, Mongo::DESCENDING, Mongo::GEO2D, Mongo::GEOHAYSTACK].include?(f[1])
             field_spec[f[0].to_s] = f[1]
           else
-            raise MongoArgumentError, "Invalid index field #{f[1].inspect}; " + 
+            raise MongoArgumentError, "Invalid index field #{f[1].inspect}; " +
               "should be one of Mongo::ASCENDING (1), Mongo::DESCENDING (-1) or Mongo::GEO2D ('2d')."
           end
         end
       else
-        raise MongoArgumentError, "Invalid index specification #{spec.inspect}; " + 
+        raise MongoArgumentError, "Invalid index specification #{spec.inspect}; " +
           "should be either a string, symbol, or an array of arrays."
       end
       field_spec
diff --git a/lib/mongo/cursor.rb b/lib/mongo/cursor.rb
index 452913d..9579562 100644
--- a/lib/mongo/cursor.rb
+++ b/lib/mongo/cursor.rb
@@ -43,6 +43,7 @@ module Mongo
       @collection = collection
       @connection = @db.connection
       @logger     = @connection.logger
+      @log_duration = @connection.log_duration
 
       # Query selector
       @selector   = opts[:selector] || {}
@@ -199,7 +200,7 @@ module Mongo
     # This method overrides any sort order specified in the Collection#find
     # method, and only the last sort applied has an effect.
     #
-    # @param [Symbol, Array] key_or_list either 1) a key to sort by or 2) 
+    # @param [Symbol, Array] key_or_list either 1) a key to sort by or 2)
     #   an array of [key, direction] pairs to sort by. Direction should
     #   be specified as Mongo::ASCENDING (or :ascending / :asc) or Mongo::DESCENDING (or :descending / :desc)
     #
diff --git a/lib/mongo/util/logging.rb b/lib/mongo/util/logging.rb
index 48a22a3..b6f303d 100644
--- a/lib/mongo/util/logging.rb
+++ b/lib/mongo/util/logging.rb
@@ -25,14 +25,15 @@ module Mongo
 
     # Execute the block and log the operation described by name and payload.
     def instrument(name, payload = {}, &blk)
+      time_elapse = @log_duration ? Time.now : nil
       res = yield
-      log_operation(name, payload)
+      log_operation(name, payload, time_elapse)
       res
     end
 
     protected
 
-    def log_operation(name, payload)
+    def log_operation(name, payload, time_elapse=nil)
       @logger ||= nil
       return unless @logger
       msg = "#{payload[:database]}['#{payload[:collection]}'].#{name}("
@@ -40,7 +41,11 @@ module Mongo
       msg += ".skip(#{payload[:skip]})"  if payload[:skip]
       msg += ".limit(#{payload[:limit]})"  if payload[:limit]
       msg += ".sort(#{payload[:order]})"  if payload[:order]
-      @logger.debug "MONGODB #{msg}"
+      if time_elapse
+        @logger.debug "MONGODB (#{Time.now - time_elapse}s) #{msg}"
+      else
+        @logger.debug "MONGODB #{msg}"
+      end
     end
 
   end
diff --git a/test/connection_test.rb b/test/connection_test.rb
index f467673..8a64357 100644
--- a/test/connection_test.rb
+++ b/test/connection_test.rb
@@ -137,6 +137,15 @@ class TestConnection < Test::Unit::TestCase
     assert output.string.include?("admin['$cmd'].find")
   end
 
+  def test_logging_duration
+    output = StringIO.new
+    logger = Logger.new(output)
+    logger.level = Logger::DEBUG
+    connection = standard_connection(:logger => logger, :log_duration => true).db(MONGO_TEST_DB)
+    assert output.string.index(/\(0.\d+s\)/)
+    assert output.string.include?("admin['$cmd'].find")
+  end
+
   def test_connection_logger
     output = StringIO.new
     logger = Logger.new(output)