Restore connection_id on error.

2 files changed, 22 insertions, 4 deletions

diff --git a/activerecord/lib/active_record/query_cache.rb b/activerecord/lib/active_record/query_cache.rb
index 9701898415..d64dee10fe 100644
--- a/activerecord/lib/active_record/query_cache.rb
+++ b/activerecord/lib/active_record/query_cache.rb
@@ -34,16 +34,22 @@ module ActiveRecord
 
       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
+        restore_query_cache_settings(connection_id, enabled)
       end
 
       response
     rescue Exception => e
+      restore_query_cache_settings(connection_id, enabled)
+      raise e
+    end
+
+    private
+
+    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
-      raise e
     end
+
   end
 end
diff --git a/activerecord/test/cases/query_cache_test.rb b/activerecord/test/cases/query_cache_test.rb
index 83e207a260..2d778e9e90 100644
--- a/activerecord/test/cases/query_cache_test.rb
+++ b/activerecord/test/cases/query_cache_test.rb
@@ -39,6 +39,18 @@ class QueryCacheTest < ActiveRecord::TestCase
     assert ActiveRecord::Base.connection.query_cache_enabled, 'cache on'
   end
 
+  def test_exceptional_middleware_assigns_original_connection_id_on_error
+    connection_id = ActiveRecord::Base.connection_id
+
+    mw = ActiveRecord::QueryCache.new lambda { |env|
+      ActiveRecord::Base.connection_id = self.object_id
+      raise "lol borked"
+    }
+    assert_raises(RuntimeError) { mw.call({}) }
+
+    assert_equal connection_id, ActiveRecord::Base.connection_id 
+  end
+
   def test_middleware_delegates
     called = false
     mw = ActiveRecord::QueryCache.new lambda { |env|