use Rack::BodyProxy in activerecord middlewares

3 files changed, 23 insertions, 86 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 b8f99adc22..d69f02d504 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
@@ -406,35 +406,6 @@ module ActiveRecord
     end
 
     class ConnectionManagement
-      class Proxy # :nodoc:
-        attr_reader :body, :testing
-
-        def initialize(body, testing = false)
-          @body    = body
-          @testing = testing
-        end
-
-        def method_missing(method_sym, *arguments, &block)
-          @body.send(method_sym, *arguments, &block)
-        end
-
-        def respond_to?(method_sym, include_private = false)
-          super || @body.respond_to?(method_sym)
-        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
@@ -442,9 +413,12 @@ module ActiveRecord
       def call(env)
         testing = env.key?('rack.test')
 
-        status, headers, body = @app.call(env)
+        response = @app.call(env)
+        response[2] = ::Rack::BodyProxy.new(response[2]) do
+          ActiveRecord::Base.clear_active_connections! unless testing
+        end
 
-        [status, headers, Proxy.new(body, testing)]
+        response
       rescue
         ActiveRecord::Base.clear_active_connections! unless testing
         raise
diff --git a/activerecord/lib/active_record/identity_map.rb b/activerecord/lib/active_record/identity_map.rb
index 680d9ffea0..d9777bb2f6 100644
--- a/activerecord/lib/active_record/identity_map.rb
+++ b/activerecord/lib/active_record/identity_map.rb
@@ -123,24 +123,6 @@ module ActiveRecord
     end
 
     class Middleware
-      class Body #:nodoc:
-        def initialize(target, original)
-          @target   = target
-          @original = original
-        end
-
-        def each(&block)
-          @target.each(&block)
-        end
-
-        def close
-          @target.close if @target.respond_to?(:close)
-        ensure
-          IdentityMap.enabled = @original
-          IdentityMap.clear
-        end
-      end
-
       def initialize(app)
         @app = app
       end
@@ -148,8 +130,14 @@ module ActiveRecord
       def call(env)
         enabled = IdentityMap.enabled
         IdentityMap.enabled = true
-        status, headers, body = @app.call(env)
-        [status, headers, Body.new(body, enabled)]
+
+        response = @app.call(env)
+        response[2] = Rack::BodyProxy.new(response[2]) do
+          IdentityMap.enabled = enabled
+          IdentityMap.clear
+        end
+
+        response
       end
     end
   end
diff --git a/activerecord/lib/active_record/query_cache.rb b/activerecord/lib/active_record/query_cache.rb
index 466d148901..9701898415 100644
--- a/activerecord/lib/active_record/query_cache.rb
+++ b/activerecord/lib/active_record/query_cache.rb
@@ -27,47 +27,22 @@ module ActiveRecord
       @app = app
     end
 
-    class BodyProxy # :nodoc:
-      def initialize(original_cache_value, target, connection_id)
-        @original_cache_value = original_cache_value
-        @target               = target
-        @connection_id        = connection_id
-      end
-
-      def method_missing(method_sym, *arguments, &block)
-        @target.send(method_sym, *arguments, &block)
-      end
-
-      def respond_to?(method_sym, include_private = false)
-        super || @target.respond_to?(method_sym)
-      end
-
-      def each(&block)
-        @target.each(&block)
-      end
+    def call(env)
+      enabled       = ActiveRecord::Base.connection.query_cache_enabled
+      connection_id = ActiveRecord::Base.connection_id
+      ActiveRecord::Base.connection.enable_query_cache!
 
-      def close
-        @target.close if @target.respond_to?(:close)
-      ensure
-        ActiveRecord::Base.connection_id = @connection_id
+      response = @app.call(env)
+      response[2] = Rack::BodyProxy.new(response[2]) do
+        ActiveRecord::Base.connection_id = connection_id
         ActiveRecord::Base.connection.clear_query_cache
-        unless @original_cache_value
-          ActiveRecord::Base.connection.disable_query_cache!
-        end
+        ActiveRecord::Base.connection.disable_query_cache! unless enabled
       end
-    end
 
-    def call(env)
-      old = ActiveRecord::Base.connection.query_cache_enabled
-      ActiveRecord::Base.connection.enable_query_cache!
-
-      status, headers, body = @app.call(env)
-      [status, headers, BodyProxy.new(old, body, ActiveRecord::Base.connection_id)]
+      response
     rescue Exception => e
       ActiveRecord::Base.connection.clear_query_cache
-      unless old
-        ActiveRecord::Base.connection.disable_query_cache!
-      end
+      ActiveRecord::Base.connection.disable_query_cache! unless enabled
       raise e
     end
   end