1 files changed, 17 insertions, 25 deletions
diff --git a/activerecord/lib/active_record/query_cache.rb b/activerecord/lib/active_record/query_cache.rb
index dcb2bd3d84..f451ed1764 100644
--- a/activerecord/lib/active_record/query_cache.rb
+++ b/activerecord/lib/active_record/query_cache.rb
@@ -23,34 +23,26 @@ module ActiveRecord
       end
     end
 
-    def initialize(app)
-      @app = app
-    end
-
-    def call(env)
-      connection    = ActiveRecord::Base.connection
-      enabled       = connection.query_cache_enabled
-      connection_id = ActiveRecord::Base.connection_id
-      connection.enable_query_cache!
-
-      response = @app.call(env)
-      response[2] = Rack::BodyProxy.new(response[2]) do
-        restore_query_cache_settings(connection_id, enabled)
+    def self.install_executor_hooks(executor = ActiveSupport::Executor)
+      executor.to_run do
+        connection    = ActiveRecord::Base.connection
+        enabled       = connection.query_cache_enabled
+        connection_id = ActiveRecord::Base.connection_id
+        connection.enable_query_cache!
+
+        @restore_query_cache_settings = lambda do
+          ActiveRecord::Base.connection_id = connection_id
+          ActiveRecord::Base.connection.clear_query_cache
+          ActiveRecord::Base.connection.disable_query_cache! unless enabled
+        end
       end
 
-      response
-    rescue Exception => e
-      restore_query_cache_settings(connection_id, enabled)
-      raise e
-    end
-
-    private
+      executor.to_complete do
+        @restore_query_cache_settings.call if defined?(@restore_query_cache_settings)
 
-    def restore_query_cache_settings(connection_id, enabled)
-      ActiveRecord::Base.connection_id = connection_id
-      ActiveRecord::Base.connection.clear_query_cache
-      ActiveRecord::Base.connection.disable_query_cache! unless enabled
+        # FIXME: This should be skipped when env['rack.test']
+        ActiveRecord::Base.clear_active_connections!
+      end
     end
-
   end
 end