use Rack::BodyProxy in activerecord middlewares

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

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

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 close
-        @target.close if @target.respond_to?(:close)
-      ensure
-        ActiveRecord::Base.connection_id = @connection_id
-        ActiveRecord::Base.connection.clear_query_cache
-        unless @original_cache_value
-          ActiveRecord::Base.connection.disable_query_cache!
-        end
-      end
-    end
-
     def call(env)
-      old = ActiveRecord::Base.connection.query_cache_enabled
+      enabled       = ActiveRecord::Base.connection.query_cache_enabled
+      connection_id = ActiveRecord::Base.connection_id
       ActiveRecord::Base.connection.enable_query_cache!
 
-      status, headers, body = @app.call(env)
-      [status, headers, BodyProxy.new(old, body, ActiveRecord::Base.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
+        ActiveRecord::Base.connection.disable_query_cache! unless enabled
+      end
+
+      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

activerecord/test/cases/connection_management_test.rb

@@ -1,4 +1,5 @@
 require "cases/helper"
+require "rack"
 
 module ActiveRecord
   module ConnectionAdapters
@@ -80,9 +81,10 @@ module ActiveRecord
 
       test "proxy is polite to it's body and responds to it" do
         body = Class.new(String) { def to_path; "/path"; end }.new
-        proxy = ConnectionManagement::Proxy.new(body)
-        assert proxy.respond_to?(:to_path)
-        assert_equal proxy.to_path, "/path"
+        app = lambda { |_| [200, {}, body] }
+        response_body = ConnectionManagement.new(app).call(@env)[2]
+        assert response_body.respond_to?(:to_path)
+        assert_equal response_body.to_path, "/path"
       end
     end
   end

activerecord/test/cases/identity_map/middleware_test.rb

@@ -1,4 +1,5 @@
 require "cases/helper"
+require "rack"
 
 module ActiveRecord
   module IdentityMap
@@ -19,6 +20,7 @@ module ActiveRecord
         called = false
         mw = Middleware.new lambda { |env|
           called = true
+          [200, {}, nil]
         }
         mw.call({})
         assert called, 'middleware delegated'
@@ -27,6 +29,7 @@ module ActiveRecord
       def test_im_enabled_during_delegation
         mw = Middleware.new lambda { |env|
           assert IdentityMap.enabled?, 'identity map should be enabled'
+          [200, {}, nil]
         }
         mw.call({})
       end

activerecord/test/cases/query_cache_test.rb

@@ -3,7 +3,7 @@ require 'models/topic'
 require 'models/task'
 require 'models/category'
 require 'models/post'
-
+require 'rack'
 
 class QueryCacheTest < ActiveRecord::TestCase
   fixtures :tasks, :topics, :categories, :posts, :categories_posts
@@ -43,6 +43,7 @@ class QueryCacheTest < ActiveRecord::TestCase
     called = false
     mw = ActiveRecord::QueryCache.new lambda { |env|
       called = true
+      [200, {}, nil]
     }
     mw.call({})
     assert called, 'middleware should delegate'
@@ -53,6 +54,7 @@ class QueryCacheTest < ActiveRecord::TestCase
       Task.find 1
       Task.find 1
       assert_equal 1, ActiveRecord::Base.connection.query_cache.length
+      [200, {}, nil]
     }
     mw.call({})
   end
@@ -62,6 +64,7 @@ class QueryCacheTest < ActiveRecord::TestCase
 
     mw = ActiveRecord::QueryCache.new lambda { |env|
       assert ActiveRecord::Base.connection.query_cache_enabled, 'cache on'
+      [200, {}, nil]
     }
     mw.call({})
   end
@@ -83,7 +86,7 @@ class QueryCacheTest < ActiveRecord::TestCase
   end
 
   def test_cache_off_after_close
-    mw = ActiveRecord::QueryCache.new lambda { |env| }
+    mw = ActiveRecord::QueryCache.new lambda { |env| [200, {}, nil] }
     body = mw.call({}).last
 
     assert ActiveRecord::Base.connection.query_cache_enabled, 'cache enabled'
@@ -94,6 +97,7 @@ class QueryCacheTest < ActiveRecord::TestCase
   def test_cache_clear_after_close
     mw = ActiveRecord::QueryCache.new lambda { |env|
       Post.find(:first)
+      [200, {}, nil]
     }
     body = mw.call({}).last
 
@@ -244,14 +248,3 @@ class QueryCacheExpiryTest < ActiveRecord::TestCase
     end
   end
 end
-
-class QueryCacheBodyProxyTest < ActiveRecord::TestCase
-
-  test "is polite to it's body and responds to it" do
-    body = Class.new(String) { def to_path; "/path"; end }.new
-    proxy = ActiveRecord::QueryCache::BodyProxy.new(nil, body, ActiveRecord::Base.connection_id)
-    assert proxy.respond_to?(:to_path)
-    assert_equal proxy.to_path, "/path"
-  end
-
-end