From 1b22071b276a2a2b3e65e6db40c6a0aa92d9ebc3 Mon Sep 17 00:00:00 2001
From: Joshua Peek <josh@joshpeek.com>
Date: Tue, 24 Feb 2009 11:34:32 -0600
Subject: [PATCH] Ensure ActiveRecord session store's connections are checked
 in after each request [#1927 state:resolved]

---
 actionpack/lib/action_controller/dispatcher.rb   |  7 -------
 .../abstract/connection_pool.rb                  | 16 ++++++++++++++++
 railties/lib/initializer.rb                      | 16 +++++++++++++++-
 3 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/actionpack/lib/action_controller/dispatcher.rb b/actionpack/lib/action_controller/dispatcher.rb
index ec40b5c4e66d6..07931e4a4a50a 100644
--- a/actionpack/lib/action_controller/dispatcher.rb
+++ b/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
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 901b17124c271..aac84cc5f4996 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 @@ 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
diff --git a/railties/lib/initializer.rb b/railties/lib/initializer.rb
index a31ae9422e1d6..09affe9e3615e 100644
--- a/railties/lib/initializer.rb
+++ b/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