Modify connection pool callbacks to be compatible w/ new style

Signed-off-by: Yehuda Katz <wycats@gmail.com>

2 files changed, 27 insertions, 8 deletions

diff --git a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
index 12253eac3f..377f2a44c5 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
@@ -211,9 +211,10 @@ module ActiveRecord
       # calling +checkout+ on this pool.
       def checkin(conn)
         @connection_mutex.synchronize do
-          conn.run_callbacks :checkin
-          @checked_out.delete conn
-          @queue.signal
+          conn.run_callbacks :checkin do
+            @checked_out.delete conn
+            @queue.signal
+          end
         end
       end
 
@@ -255,9 +256,10 @@ module ActiveRecord
       end
 
       def checkout_and_verify(c)
-        c.verify!
-        c.run_callbacks :checkout
-        @checked_out << c
+        c.run_callbacks :checkout do
+          c.verify!
+          @checked_out << c
+        end
         c
       end
     end
diff --git a/activerecord/test/cases/pooled_connections_test.rb b/activerecord/test/cases/pooled_connections_test.rb
index bb9013c2a1..f9eea3d118 100644
--- a/activerecord/test/cases/pooled_connections_test.rb
+++ b/activerecord/test/cases/pooled_connections_test.rb
@@ -4,14 +4,14 @@ require "timeout"
 
 class PooledConnectionsTest < ActiveRecord::TestCase
   def setup
-    super
+    @per_test_teardown = []
     @connection = ActiveRecord::Base.remove_connection
   end
 
   def teardown
     ActiveRecord::Base.clear_all_connections!
     ActiveRecord::Base.establish_connection(@connection)
-    super
+    @per_test_teardown.each {|td| td.call }
   end
 
   def checkout_connections
@@ -113,6 +113,23 @@ class PooledConnectionsTest < ActiveRecord::TestCase
     assert_equal 3, after_count - before_count
   end
 
+  def test_connection_pool_callbacks
+    checked_out, checked_in = false, false
+    ActiveRecord::ConnectionAdapters::AbstractAdapter.class_eval do
+      set_callback(:checkout, :after) { checked_out = true }
+      set_callback(:checkin, :before) { checked_in = true }
+    end
+    @per_test_teardown << proc do
+      ActiveRecord::ConnectionAdapters::AbstractAdapter.class_eval do
+        reset_callbacks :checkout
+        reset_callbacks :checkin
+      end
+    end
+    checkout_checkin_connections 1, 1
+    assert checked_out
+    assert checked_in
+  end
+
   private
 
   def add_record(name)