module ActiveRecord
module ConnectionAdapters
class ConnectionPool
def cleanup_connection
return yield if Thread.current[:__AR__cleanup_connection]
begin
Thread.current[:__AR__cleanup_connection] = true
yield
ensure
release_connection
Thread.current[:__AR__cleanup_connection] = false
end
end
end
end
class Base
class << self
def cleanup_connection(&block)
connection_pool.cleanup_connection(&block)
end
# comment out this redefinition once you've wrapped all necessary methods
alias_method :connection_without_cleanup_connection_check, :connection
def connection(*a)
if ! Thread.current[:__AR__cleanup_connection]
puts "connection called outside of cleanup_connection block", caller, "\n"
end
connection_without_cleanup_connection_check(*a)
end
end
end
end
methods_to_wrap = {
(class<<ActiveRecord::Base;self;end) => [
:find, :find_every, :find_by_sql, :transaction, :count, :create, :delete, :count_by_sql,
:update, :destroy, :cache, :uncached, :quoted_table_name, :columns, :exists?, :update_all,
:increment_counter, :decrement_counter, :delete_all, :table_exists?, :update_counters,
],
ActiveRecord::Base => [:quoted_id, :valid?],
ActiveRecord::Associations::AssociationCollection => [:initialize, :find, :find_target, :load_target, :count],
ActiveRecord::Associations::HasAndBelongsToManyAssociation => [:create],
ActiveRecord::Associations::HasManyThroughAssociation => [:construct_conditions],
ActiveRecord::Associations::HasOneAssociation => [:construct_sql],
ActiveRecord::Associations::ClassMethods => [:collection_reader_method, :configure_dependency_for_has_many],
ActiveRecord::Calculations::ClassMethods => [:calculate],
}
methods_to_wrap[Test::Unit::TestSuite] = [:run] if defined?(Test::Unit::TestSuite)
methods_to_wrap.each do |klass, methods|
methods.each do |method|
klass.class_eval do
alias_method_chain method, :cleanup_connection do |target, punc|
eval %{
def #{target}_with_cleanup_connection#{punc}(*a, &b)
ActiveRecord::Base.connection_pool.cleanup_connection do
#{target}_without_cleanup_connection#{punc}(*a, &b)
end
end
}
end
end
end
end
