require 'active_record/connection_adapters/postgresql_adapter'

class ActiveRecord::ConnectionAdapters::PostgreSQLAdapter
  def query_with_stats(sql, name = nil)
    rows = nil
    elapsed = Benchmark.measure do
      rows = query_without_stats(sql, name)
    end
    msec = (elapsed.real * 1000).to_i
    record_query_stats(sql, name, msec, rows)
    rows
  end
  
  def execute_with_stats(sql, name = nil)
    result = nil
    elapsed = Benchmark.measure do
      result = execute_without_stats(sql, name)
    end
    msec = (elapsed.real * 1000).to_i
    if result.respond_to?(:rows)
      record_query_stats(sql, name, msec, result.rows)
    else
      record_query_stats(sql, name, msec, result)
    end
    result
  end
  
  alias_method_chain :query, :stats
  alias_method_chain :execute, :stats
end