Ensure ActiveRecord::Base.connection_pool.with_connection creates a new connection only when needed [#1752 state:resolved]

Signed-off-by: Pratik Naik <pratiknaik@gmail.com>

2 files changed, 35 insertions, 5 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 aac84cc5f4..e8e736bf38 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
@@ -107,13 +107,14 @@ module ActiveRecord
         checkin conn if conn
       end
 
-      # Reserve a connection, and yield it to a block. Ensure the connection is
-      # checked back in when finished.
+      # If a connection already exists yield it to the block.  If no connection
+      # exists checkout a connection, yield it to the block, and checkin the 
+      # connection when finished.
       def with_connection
-        conn = checkout
-        yield conn
+        fresh_connection = true unless @reserved_connections[current_connection_id]
+        yield connection
       ensure
-        checkin conn
+        release_connection if fresh_connection
       end
 
       # Returns true if a connection has already been opened.
diff --git a/activerecord/test/cases/pooled_connections_test.rb b/activerecord/test/cases/pooled_connections_test.rb
index 2649a9358a..bbd943fc6e 100644
--- a/activerecord/test/cases/pooled_connections_test.rb
+++ b/activerecord/test/cases/pooled_connections_test.rb
@@ -1,4 +1,6 @@
 require "cases/helper"
+require "models/project"
+require "timeout"
 
 class PooledConnectionsTest < ActiveRecord::TestCase
   def setup
@@ -89,6 +91,33 @@ class PooledConnectionsTest < ActiveRecord::TestCase
   ensure
     ActiveRecord::Base.connection_handler = old_handler
   end
+
+  def test_with_connection_nesting_safety
+    ActiveRecord::Base.establish_connection(@connection.merge({:pool => 1, :wait_timeout => 0.1}))
+
+    before_count = Project.count
+
+    add_record('one')
+
+    ActiveRecord::Base.connection.transaction do
+      add_record('two')
+      # Have another thread try to screw up the transaction
+      Thread.new do
+        raise ActiveRecord::Rollback
+        ActiveRecord::Base.connection_pool.release_connection
+      end.join rescue nil
+      add_record('three')
+    end
+
+    after_count = Project.count
+    assert_equal 3, after_count - before_count
+  end
+
+  private
+
+  def add_record(name)
+    ActiveRecord::Base.connection_pool.with_connection { Project.create! :name => name }
+  end
 end unless %w(FrontBase).include? ActiveRecord::Base.connection.adapter_name
 
 class AllowConcurrencyDeprecatedTest < ActiveRecord::TestCase