# NewRelic instrumentation for ActiveRecord
if defined?(ActiveRecord::Base) && !NewRelic::Control.instance['skip_ar_instrumentation']
  
  module NewRelic
    module Agent
      module Instrumentation
        module ActiveRecordInstrumentation

          def self.included(instrumented_class)
            instrumented_class.class_eval do
              alias_method :log_without_newrelic_instrumentation, :log
              alias_method :log, :log_with_newrelic_instrumentation
              protected :log
            end
          end
          
          def active_record_all_stats
            NewRelic::Agent.instance.stats_engine.get_stats_no_scope("ActiveRecord/all")
          end
          
          def log_with_newrelic_instrumentation(sql, name, &block)
            # Capture db config if we are going to try to get the explain plans
            if (defined?(ActiveRecord::ConnectionAdapters::MysqlAdapter) && self.is_a?(ActiveRecord::ConnectionAdapters::MysqlAdapter)) ||
                (defined?(ActiveRecord::ConnectionAdapters::PostgreSQLAdapter) && self.is_a?(ActiveRecord::ConnectionAdapters::PostgreSQLAdapter))
              supported_config = @config
            end
            if name && (parts = name.split " ") && parts.size == 2
              model = parts.first
              operation = parts.last.downcase
              metric_name = case operation
                            when 'load' then 'find'
                            when 'indexes', 'columns' then nil # fall back to DirectSQL
                            when 'destroy', 'find', 'save', 'create' then operation
                            when 'update' then 'save'
                            else
                              if model == 'Join'
                                operation
                              end
                            end
              metric = "ActiveRecord/#{model}/#{metric_name}" if metric_name
            end      
            if metric.nil? && sql =~ /^(select|update|insert|delete)/i
              # Could not determine the model/operation so let's find a better
              # metric.  If it doesn't match the regex, it's probably a show
              # command or some DDL which we'll ignore.
              metric = "Database/SQL/#{$1.downcase}"
            end
            
            if !metric
              log_without_newrelic_instrumentation(sql, name, &block)
            else
              self.class.trace_method_execution_with_scope metric, true, true do        
                t0 = Time.now.to_f
                result = log_without_newrelic_instrumentation(sql, name, &block)
                duration = Time.now.to_f - t0
                
                NewRelic::Agent.instance.transaction_sampler.notice_sql(sql, supported_config, duration)
                # Record in the overall summary metric
                active_record_all_stats.record_data_point(duration)
                # Record in the summary metric for this operation
                NewRelic::Agent.instance.stats_engine.get_stats_no_scope("ActiveRecord/#{metric_name}").record_data_point(duration) if metric_name
                result
              end
            end
          end
          
        end
        
        # instrumentation to catch logged SQL statements in sampled transactions
        ActiveRecord::ConnectionAdapters::AbstractAdapter.module_eval do
          include ::NewRelic::Agent::Instrumentation::ActiveRecordInstrumentation
        end

        # This instrumentation will add an extra scope to the transaction traces
        # which will show the code surrounding the query, inside the model find_by_sql
        # method.
        ActiveRecord::Base.class_eval do
          class << self
            add_method_tracer :find_by_sql, 'ActiveRecord/#{self.name}/find_by_sql', :metric => false
          end
        end
        
      end
    end
  end
end