mongo_ext.rb

module MongoInstrumenter
  class LogSubscriber < ActiveSupport::LogSubscriber
    def self.runtime=(value)
      Thread.current["mongoid_query_runtime"] = value
    end

    def self.runtime
      Thread.current["mongoid_query_runtime"] ||= 0
    end

    def self.reset_runtime
      rt, self.runtime = runtime, 0
      rt
    end

    def initialize
      super
      @odd_or_even = false
    end

    def query(event)
      self.class.runtime += event.duration
      logger = event.payload[:logger]
      return unless logger.logger.debug?

      name = '%s (%.1fms)' % ["MONGODB", event.duration]
      query  = event.payload[:query].squeeze(' ')

      if odd?
        name = color(name, CYAN, true)
        query  = color(query, nil, true)
      else
        name = color(name, MAGENTA, true)
      end

      debug "  #{name}  #{query}"
    end

    def odd?
      @odd_or_even = !@odd_or_even
    end
  end
  
  module CollectionInstrumenter
    extend ActiveSupport::Concern
    
    included do
      private
      def instrumenter
        @instrumenter ||= ActiveSupport::Notifications.instrumenter
      end
    
      def remove_with_instrumenter(selector={}, opts={})
        instrument "#{@db.name}['#{@name}'].remove(#{selector.inspect})" do
          remove_without_instrumenter(*args)
        end
      end

      def update_with_instrumenter(*args)
        instrument "#{@db.name}['#{@name}'].update(#{selector.inspect}, #{document.inspect})" do
          update_without_instrumenter(*args)
        end
      end

      def insert_documents_with_instrumenter(*args)
        instrument "#{@db.name}['#{collection_name}'].insert(#{documents.inspect})" do
          insert_documents_without_instrumenter(*args)
        end
      end
      alias_method_chain :remove, :instrumenter
      alias_method_chain :update, :instrumenter
      alias_method_chain :insert_documents, :instrumenter

      def instrument(query)
        instrumenter.instrument("query.mongoid", :query => query, :logger => @logger) do
          yield
        end
      end
    end
  end

  module CursorInstrumenter
    extend ActiveSupport::Concern
    
    included do
      private
      def instrumenter
        @instrumenter ||= ActiveSupport::Notifications.instrumenter
      end
      
      def send_initial_query_with_instrumenter
        instrument query_log_message do
          send_initial_query_without_instrumenter
        end
      end
      alias_method_chain :send_initial_query, :instrumenter
      
      def instrument(query)
        instrumenter.instrument("query.mongoid", :query => query, :logger => @logger) do
          yield
        end
      end
    end
  end
end



MongoInstrumenter::LogSubscriber.attach_to :mongoid

Mongo::Collection.send :include, MongoInstrumenter::CollectionInstrumenter
Mongo::Cursor.send :include, MongoInstrumenter::CursorInstrumenter