proxy body responses so we close database connections after body is flushed

3 files changed, 49 insertions, 9 deletions

diff --git a/activerecord/CHANGELOG b/activerecord/CHANGELOG
index 3b6a7d7208..e536d2b408 100644
--- a/activerecord/CHANGELOG
+++ b/activerecord/CHANGELOG
@@ -1,5 +1,8 @@
 *Rails 3.1.0 (unreleased)*
 
+* ConnectionManagement middleware is changed to clean up the connection pool
+  after the rack body has been flushed.
+
 * Added an update_column method on ActiveRecord. This new method updates a given attribute on an object, skipping validations and callbacks.
   It is recommended to use #update_attribute unless you are sure you do not want to execute any callback, including the modification of
   the updated_at column. It should not be called on new records.
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 7a900055a9..5e12f80263 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
@@ -417,18 +417,35 @@ module ActiveRecord
     end
 
     class ConnectionManagement
+      class Proxy # :nodoc:
+        attr_reader :body, :testing
+
+        def initialize(body, testing = false)
+          @body    = body
+          @testing = testing
+        end
+
+        def each(&block)
+          body.each(&block)
+        end
+
+        def close
+          body.close if body.respond_to?(:close)
+
+          # Don't return connection (and perform implicit rollback) if
+          # this request is a part of integration test
+          ActiveRecord::Base.clear_active_connections! unless testing
+        end
+      end
+
       def initialize(app)
         @app = app
       end
 
       def call(env)
-        @app.call(env)
-      ensure
-        # Don't return connection (and perform implicit rollback) if
-        # this request is a part of integration test
-        unless env.key?("rack.test")
-          ActiveRecord::Base.clear_active_connections!
-        end
+        status, headers, body = @app.call(env)
+
+        [status, headers, Proxy.new(body, env.key?('rack.test'))]
       end
     end
   end
diff --git a/activerecord/test/cases/connection_management_test.rb b/activerecord/test/cases/connection_management_test.rb
index 1313e28bb1..3734f8e5f0 100644
--- a/activerecord/test/cases/connection_management_test.rb
+++ b/activerecord/test/cases/connection_management_test.rb
@@ -11,7 +11,7 @@ module ActiveRecord
 
         def call(env)
           @calls << env
-          [200, {}, [['hi mom']]]
+          [200, {}, ['hi mom']]
         end
       end
 
@@ -32,11 +32,31 @@ module ActiveRecord
         assert_equal [@env], @app.calls
       end
 
-      test "clears active connections after each call" do
+      def test_connections_are_active_after_call
         @management.call(@env)
+        assert ActiveRecord::Base.connection_handler.active_connections?
+      end
+
+      def test_body_responds_to_each
+        _, _, body = @management.call(@env)
+        bits = []
+        body.each { |bit| bits << bit }
+        assert_equal ['hi mom'], bits
+      end
+
+      def test_connections_are_cleared_after_body_close
+        _, _, body = @management.call(@env)
+        body.close
         assert !ActiveRecord::Base.connection_handler.active_connections?
       end
 
+      def test_active_connections_are_not_cleared_on_body_close_during_test
+        @env['rack.test'] = true
+        _, _, body = @management.call(@env)
+        body.close
+        assert ActiveRecord::Base.connection_handler.active_connections?
+      end
+
       test "doesn't clear active connections when running in a test case" do
         @env['rack.test'] = true
         @management.call(@env)