From 3842422d732341bfdbe35b27d658de347480ee6e Mon Sep 17 00:00:00 2001 From: Nick Sieger Date: Wed, 18 Feb 2009 14:26:09 -0600 Subject: [PATCH] Allow AR::Base.connection to take a block meaning cleanup connection after block finishes - idea based on coderrr's original patch for 1752 --- .../abstract/connection_pool.rb | 22 ++++++-- .../abstract/connection_specification.rb | 23 ++++++-- activerecord/test/cases/pooled_connections_test.rb | 56 ++++++++++++++++++++ 3 files changed, 89 insertions(+), 12 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 901b171..a66b8d9 100644 --- a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb +++ b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb @@ -91,11 +91,21 @@ module ActiveRecord # # #connection can be called any number of times; the connection is # held in a hash keyed by the thread id. - def connection - if conn = @reserved_connections[current_connection_id] - conn + def connection(&block) + fresh_connection = false + unless conn = @reserved_connections[current_connection_id] + fresh_connection = true + conn = @reserved_connections[current_connection_id] = checkout + end + if block_given? + begin + yield conn + ensure + release_connection if fresh_connection + end + true # else we raise ConnectionNotEstablished below else - @reserved_connections[current_connection_id] = checkout + conn end end @@ -321,9 +331,9 @@ module ActiveRecord # active or defined connection: if it is the latter, it will be # opened and set as the active connection for the class it was defined # for (not necessarily the current class). - def retrieve_connection(klass) #:nodoc: + def retrieve_connection(klass, &block) #:nodoc: pool = retrieve_connection_pool(klass) - (pool && pool.connection) or raise ConnectionNotEstablished + (pool && pool.connection(&block)) or raise ConnectionNotEstablished end # Returns true if a connection that's accessible to this class has diff --git a/activerecord/lib/active_record/connection_adapters/abstract/connection_specification.rb b/activerecord/lib/active_record/connection_adapters/abstract/connection_specification.rb index bbc290f..9ad69df 100644 --- a/activerecord/lib/active_record/connection_adapters/abstract/connection_specification.rb +++ b/activerecord/lib/active_record/connection_adapters/abstract/connection_specification.rb @@ -16,8 +16,19 @@ module ActiveRecord # Returns the connection currently associated with the class. This can # also be used to "borrow" the connection to do database work that isn't # easily done without going straight to SQL. - def connection - self.class.connection + # + # If called with a block, the connection is automatically released + # after the block is yielded. This block version is useful outside + # of normal actionpack dispatcher operation where connection + # cleanup must be managed by the application. + # + # User.connection do + # User.all.each do |u| + # # do something with each user + # end + # end # connection is released when the block completes + def connection(&block) + self.class.connection(&block) end # Establishes the connection to the database. Accepts a hash as input where @@ -111,16 +122,16 @@ module ActiveRecord # Returns the connection currently associated with the class. This can # also be used to "borrow" the connection to do database work unrelated # to any of the specific Active Records. - def connection - retrieve_connection + def connection(&block) + retrieve_connection(&block) end def connection_pool connection_handler.retrieve_connection_pool(self) end - def retrieve_connection - connection_handler.retrieve_connection(self) + def retrieve_connection(&block) + connection_handler.retrieve_connection(self, &block) end # Returns true if +ActiveRecord+ is connected. diff --git a/activerecord/test/cases/pooled_connections_test.rb b/activerecord/test/cases/pooled_connections_test.rb index 2649a93..478119e 100644 --- a/activerecord/test/cases/pooled_connections_test.rb +++ b/activerecord/test/cases/pooled_connections_test.rb @@ -1,4 +1,5 @@ require "cases/helper" +require "timeout" class PooledConnectionsTest < ActiveRecord::TestCase def setup @@ -89,6 +90,61 @@ class PooledConnectionsTest < ActiveRecord::TestCase ensure ActiveRecord::Base.connection_handler = old_handler end + + def test_cleanup_connection + ActiveRecord::Base.establish_connection(@connection.merge({:pool => 2, :wait_timeout => 10})) + threads = [] + 10.times do + threads << Thread.new do + ActiveRecord::Base.connection do |conn| + conn.execute("select 1") + end + end + end + + assert_nothing_raised do + Timeout.timeout(2) do + threads.each {|t| t.join } + end + end + end + + def test_cleanup_connection_nesting_safety + ActiveRecord::Base.establish_connection(@connection.merge({:pool => 1, :wait_timeout => 0.1})) + + before_count = ActiveRecord::Base.connection.execute("select * from projects").num_rows + + add_record('one') + + transaction do + add_record('two') + # have another thread try to screw up the transaction + Thread.new do + ActiveRecord::Base.connection.execute("rollback") + ActiveRecord::Base.connection_pool.release_connection + end.join rescue nil + add_record('three') + end + + after_count = ActiveRecord::Base.connection.execute("select * from projects").num_rows + assert_equal 3, after_count - before_count + end + + private + + def add_record(name) + ActiveRecord::Base.connection do |conn| + conn.execute("insert into projects (name) values ('#{name}')") + end + end + + def transaction + ActiveRecord::Base.connection do |conn| + conn.execute("start transaction") + yield + conn.execute("commit") + end + end end unless %w(FrontBase).include? ActiveRecord::Base.connection.adapter_name class AllowConcurrencyDeprecatedTest < ActiveRecord::TestCase -- 1.5.6