introduce a body proxy to ensure that query cache is enabled during streaming

3 files changed, 58 insertions, 3 deletions

diff --git a/activerecord/lib/active_record/connection_adapters/abstract/query_cache.rb b/activerecord/lib/active_record/connection_adapters/abstract/query_cache.rb
index 1db397f584..093c30aa42 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/query_cache.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/query_cache.rb
@@ -29,6 +29,14 @@ module ActiveRecord
         @query_cache_enabled = old
       end
 
+      def enable_query_cache!
+        @query_cache_enabled = true
+      end
+
+      def disable_query_cache!
+        @query_cache_enabled = false
+      end
+
       # Disable the query cache within the block.
       def uncached
         old, @query_cache_enabled = @query_cache_enabled, false
diff --git a/activerecord/lib/active_record/query_cache.rb b/activerecord/lib/active_record/query_cache.rb
index d9f85a4e5e..929998eb85 100644
--- a/activerecord/lib/active_record/query_cache.rb
+++ b/activerecord/lib/active_record/query_cache.rb
@@ -27,10 +27,31 @@ module ActiveRecord
       @app = app
     end
 
-    def call(env)
-      ActiveRecord::Base.cache do
-        @app.call(env)
+    class BodyProxy # :nodoc:
+      def initialize(original_cache_value, target)
+        @original_cache_value = original_cache_value
+        @target               = target
+      end
+
+      def each(&block)
+        @target.each(&block)
+      end
+
+      def close
+        @target.close if @target.respond_to?(:close)
+      ensure
+        unless @original_cache_value
+          ActiveRecord::Base.connection.disable_query_cache!
+        end
       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)]
+    end
   end
 end
diff --git a/activerecord/test/cases/query_cache_test.rb b/activerecord/test/cases/query_cache_test.rb
index 66ec592cca..b2e40c6b22 100644
--- a/activerecord/test/cases/query_cache_test.rb
+++ b/activerecord/test/cases/query_cache_test.rb
@@ -10,6 +10,7 @@ class QueryCacheTest < ActiveRecord::TestCase
 
   def setup
     Task.connection.clear_query_cache
+    ActiveRecord::Base.connection.disable_query_cache!
   end
 
   def test_middleware_delegates
@@ -39,6 +40,31 @@ class QueryCacheTest < ActiveRecord::TestCase
     mw.call({})
   end
 
+  def test_cache_on_during_body_write
+    streaming = Class.new do
+      def each
+        yield ActiveRecord::Base.connection.query_cache_enabled
+      end
+    end
+
+    mw = ActiveRecord::QueryCache.new lambda { |env|
+      [200, {}, streaming.new]
+    }
+    body = mw.call({}).last
+    body.each { |x| assert x, 'cache should be on' }
+    body.close
+    assert !ActiveRecord::Base.connection.query_cache_enabled, 'cache disabled'
+  end
+
+  def test_cache_off_after_close
+    mw = ActiveRecord::QueryCache.new lambda { |env| }
+    body = mw.call({}).last
+
+    assert ActiveRecord::Base.connection.query_cache_enabled, 'cache enabled'
+    body.close
+    assert !ActiveRecord::Base.connection.query_cache_enabled, 'cache disabled'
+  end
+
   def test_find_queries
     assert_queries(ActiveRecord::IdentityMap.enabled? ? 1 : 2) { Task.find(1); Task.find(1) }
   end