Ensure ActiveRecord session store's connections are checked in after each request [#1927 state:resolved]

3 files changed, 31 insertions, 8 deletions

diff --git a/actionpack/lib/action_controller/dispatcher.rb b/actionpack/lib/action_controller/dispatcher.rb
index ec40b5c4e6..07931e4a4a 100644
--- a/actionpack/lib/action_controller/dispatcher.rb
+++ b/actionpack/lib/action_controller/dispatcher.rb
@@ -13,7 +13,6 @@ module ActionController
         end
 
         if defined?(ActiveRecord)
-          after_dispatch :checkin_connections
           to_prepare(:activerecord_instantiate_observers) { ActiveRecord::Base.instantiate_observers }
         end
 
@@ -115,11 +114,5 @@ module ActionController
     def flush_logger
       Base.logger.flush
     end
-
-    def checkin_connections
-      # Don't return connection (and peform implicit rollback) if this request is a part of integration test
-      return if @env.key?("rack.test")
-      ActiveRecord::Base.clear_active_connections!
-    end
   end
 end
diff --git a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
index 901b17124c..aac84cc5f4 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
@@ -351,5 +351,21 @@ module ActiveRecord
         retrieve_connection_pool klass.superclass
       end
     end
+
+    class ConnectionManagement
+      def initialize(app)
+        @app = app
+      end
+
+      def call(env)
+        @app.call(env)
+      ensure
+        # Don't return connection (and peform implicit rollback) if
+        # this request is a part of integration test
+        unless env.key?("rack.test")
+          ActiveRecord::Base.clear_active_connections!
+        end
+      end
+    end
   end
 end
diff --git a/railties/lib/initializer.rb b/railties/lib/initializer.rb
index a31ae9422e..09affe9e36 100644
--- a/railties/lib/initializer.rb
+++ b/railties/lib/initializer.rb
@@ -176,6 +176,9 @@ module Rails
       # the framework is now fully initialized
       after_initialize
 
+      # Setup database middleware after initializers have run
+      initialize_database_middleware
+
       # Prepare dispatcher callbacks and run 'prepare' callbacks
       prepare_dispatcher
 
@@ -410,7 +413,18 @@ Run `rake gems:install` to install the missing gems.
       if configuration.frameworks.include?(:active_record)
         ActiveRecord::Base.configurations = configuration.database_configuration
         ActiveRecord::Base.establish_connection
-        configuration.middleware.use ActiveRecord::QueryCache
+      end
+    end
+
+    def initialize_database_middleware
+      if configuration.frameworks.include?(:active_record)
+        if ActionController::Base.session_store == ActiveRecord::SessionStore
+          configuration.middleware.insert_before :"ActiveRecord::SessionStore", ActiveRecord::ConnectionAdapters::ConnectionManagement
+          configuration.middleware.insert_before :"ActiveRecord::SessionStore", ActiveRecord::QueryCache
+        else
+          configuration.middleware.use ActiveRecord::ConnectionAdapters::ConnectionManagement
+          configuration.middleware.use ActiveRecord::QueryCache
+        end
       end
     end