Ensure ActiveRecord session store's connections are checked in after …

…each request [#1927 state:resolved]

actionpack/lib/action_controller/dispatcher.rb

@@ -13,7 +13,6 @@ def define_dispatcher_callbacks(cache_classes)
         end
 
         if defined?(ActiveRecord)
-          after_dispatch :checkin_connections
           to_prepare(:activerecord_instantiate_observers) { ActiveRecord::Base.instantiate_observers }
         end
 
@@ -115,11 +114,5 @@ def _call(env)
     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

activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb

@@ -351,5 +351,21 @@ def retrieve_connection_pool(klass)
         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

railties/lib/initializer.rb

@@ -176,6 +176,9 @@ def process
       # 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 @@ def initialize_database
       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