Trace SQL queries in Rails 2

mysql_adapter_extensions.rb

module OpsApi
  # Helper methods for `ActiveRecord::ConnectionAdapters::MysqlAdapter` to enable SQL queries
  # tracing functionality. There are two goals:
  #
  # * Add source line into the query itself so it would be possible to understand where
  #   the query was issues from:
  #
  #       SELECT * FROM `admins` WHERE 
  #         ((`admins`.`user_id` = 1)) LIMIT 1
  #         /*OOPS:app/models/user.rb:72:in `admin?'*/
  #
  # * Add SQL tracing ability: record all SQL queries along with full stack traces:
  #
  #       with_sql_trace do
  #         Document.first.published?
  #       end
  #       ActiveRecord::ConnectionAdapters::MysqlAdapter.tap do |a|
  #         pp a.sql_trace
  #         pp a.sql_trace_analyzed
  #       end
  #
  module MysqlAdapterExtensions
    # A list of code points, where most specific ones are at the top, and generic are in the bottom.
    # Do not merge them into a single regular expression, so we will get more precise information
    # on where query was invoked from.
    INCLUDE_CODEPOINTS = [
      %r{^#{Rails.root}/(app/presenters/.*)},
      %r{^#{Rails.root}/(app/views/.*)},
      %r{^#{Rails.root}/(app/controllers/.*)},
      %r{^#{Rails.root}/(app/models/.*)},
      %r{^#{Rails.root}/(lib/.*)},
      %r{^#{Rails.root}/(spec/.*)},
      %r{^#{Rails.root}/(app/.*)},
      %r{^#{Rails.root}/(vendor/(?:gems|plugins)/.*)},
      %r{^#{Rails.root}/(.*)},
      %r{in `(irb)_binding'}
    ]
    # A regular expression used to skip certain code points (gems that do nothing
    # but add noice to the result).
    EXCLUDE_CODEPOINT = %r{^#{Rails.root}/(?:vendor/(?:rails|gems/(?:composite_primary_keys|db-charmer)|plugins/(?:paginating_find|acts_as_sluggable))|config/initializers/mysql_adapter_extensions\.rb|tmp/gems)}
    # A regular expression to exclude certain SQL queries from processing (who cares
    # where SHOW TABLES was issues from).
    EXCLUDE_SQL = %r{^(SET|SHOW)}

    def self.included(base)
      base.alias_method_chain :execute, :code_line_reporting
      base.extend ClassMethods
      # Cache current code revision on start
      base.current_code_revision
    end

    # Append codepoint info to every SQL query.
    def execute_with_code_line_reporting(sql, *args)
      # Skip noisy queries
      unless sql =~ EXCLUDE_SQL
        # Skip noisy codepoints
        stack = caller.reject { |line| line =~ EXCLUDE_CODEPOINT }
        # Save SQL and stack trace if SQL tracing enabled
        if self.class.sql_trace_enabled?
          self.class.sql_trace << { :sql => sql, :stack => stack }
        end
        # Find the best codepoint in stack trace and append it with current
        # code revision to SQL query as a comment
        INCLUDE_CODEPOINTS.each do |codepoint|
          if line = stack.grep(codepoint).first
            rev = self.class.current_code_revision ? "#{self.class.current_code_revision}:" : ''
            con = connection_name ? "/*DB:#{connection_name}*/" : ''
            sql = "#{sql} #{con}/*OOPS:#{rev}#{line[codepoint, 1]}*/"
            break
          end
        end
      end
      # We're done
      execute_without_code_line_reporting(sql, *args)
    end

    module ClassMethods
      # Detect the current code revision and memoize it for the future.
      def current_code_revision
        # Do we have the code revision memoized?
        return @@current_code_revision if defined?(@@current_code_revision)
        @@current_code_revision = if Rails.env.test? || Rails.env.development?
          # Do not add noise in development or test environment
          nil
        elsif File.exists?("#{Rails.root}/REVISION")
          # Capistrano-deployed application, we know where to get current revision
          File.read("#{Rails.root}/REVISION").chomp.strip
        else
          # Try to use git
          rev = `git rev-parse HEAD 2>/dev/null`.chomp.strip
          rev.empty? ? nil : rev
        end
      end

      # Enable SQL tracing.
      def sql_trace_enable!
        # Do nothing, if it was enabled before.
        unless sql_trace_enabled?
          # Clear old trace.
          @sql_trace = []
          @sql_trace_enabled = true
        end
      end

      # Disable SQL tracing.
      def sql_trace_disable!
        @sql_trace_enabled = false
      end

      # Check if SQL tracing enabled.
      def sql_trace_enabled?
        !!@sql_trace_enabled
      end

      # Get the latest SQL trace information. `Array` of `Hash`es, each contains
      # to keys:
      # 
      # * `:sql` - an actual SQL query;
      # * `:trace` - a Ruby stack trace.
      def sql_trace
        @sql_trace ||= []
      end

      # Get the latest SQL trace information in the same format as `sql_trace`,
      # but filter out all non-`SELECT` queries, and remove `Rails.root` from
      # all file paths in stack traces.
      def sql_trace_analyzed
        trace = sql_trace.reject { |row| !(row[:sql] =~ /^\s*SELECT/) }
        trace.each do |row|
          row[:stack].each { |line| line.sub!("#{Rails.root}/", '') }
          row[:stack].reject! { |row| row =~ %r{^/} }
        end
        trace
      end

      # Execute a block of code with SQL queries tracing enabled.
      def with_sql_trace
        sql_trace_enable!
        yield
      ensure
        sql_trace_disable!
      end
    end
  end

  # Helper method for `Object` to simplify SQL tracing.
  module ObjectExtensions
    # Execute a block of code with SQL queries tracing enabled.
    def with_sql_trace(&block)
      ActiveRecord::ConnectionAdapters::MysqlAdapter.with_sql_trace(&block)
    end
  end
end

ActiveRecord::ConnectionAdapters::MysqlAdapter.send(:include, OpsApi::MysqlAdapterExtensions)
Object.send(:include, OpsApi::ObjectExtensions)

Wow, that's cool!