Clear the correct query cache

This executor currently relies on `ActiveRecord::Base.connection` not
changing between `prepare` and `complete`. If something else returns
the current ActiveRecord connection to the pool early then this
`complete` call will fail to clear the correct query cache and restore
the original `query_cache_enabled` status.

This has for example been happening in Sidekiq:

     https://github.com/mperham/sidekiq/pull/3166

We can just keep track of the connection as part of the exector state.

2 files changed, 27 insertions, 4 deletions

diff --git a/activerecord/lib/active_record/query_cache.rb b/activerecord/lib/active_record/query_cache.rb
index c45c8c1697..c42c22ab09 100644
--- a/activerecord/lib/active_record/query_cache.rb
+++ b/activerecord/lib/active_record/query_cache.rb
@@ -28,12 +28,12 @@ module ActiveRecord
       enabled       = connection.query_cache_enabled
       connection.enable_query_cache!
 
-      enabled
+      [connection, enabled]
     end
 
-    def self.complete(enabled)
-      ActiveRecord::Base.connection.clear_query_cache
-      ActiveRecord::Base.connection.disable_query_cache! unless enabled
+    def self.complete((connection, enabled))
+      connection.clear_query_cache
+      connection.disable_query_cache! unless enabled
 
       unless ActiveRecord::Base.connected? && ActiveRecord::Base.connection.transaction_open?
         ActiveRecord::Base.clear_active_connections!
diff --git a/activerecord/test/cases/query_cache_test.rb b/activerecord/test/cases/query_cache_test.rb
index 4cd258695d..29b2deea26 100644
--- a/activerecord/test/cases/query_cache_test.rb
+++ b/activerecord/test/cases/query_cache_test.rb
@@ -51,6 +51,29 @@ class QueryCacheTest < ActiveRecord::TestCase
     assert !ActiveRecord::Base.connection.query_cache_enabled, "cache off"
   end
 
+  def test_exceptional_middleware_cleans_up_correct_cache
+    connection = ActiveRecord::Base.connection
+    called = false
+
+    mw = middleware { |env|
+      Task.find 1
+      Task.find 1
+      assert_equal 1, connection.query_cache.length
+
+      # Checkin connection early
+      ActiveRecord::Base.clear_active_connections!
+      # Make sure ActiveRecord::Base.connection doesn't checkout the same connection
+      ActiveRecord::Base.connection_pool.remove(connection)
+
+      called = true
+    }
+    mw.call({})
+
+    assert called
+    assert_equal 0, connection.query_cache.length
+    assert !connection.query_cache_enabled, "cache off"
+  end
+
   def test_exceptional_middleware_leaves_enabled_cache_alone
     ActiveRecord::Base.connection.enable_query_cache!