Make connection_handler overridable per thread

3 files changed, 75 insertions, 3 deletions

diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md
index 94d86f6550..3fe8d7ec2e 100644
--- a/activerecord/CHANGELOG.md
+++ b/activerecord/CHANGELOG.md
@@ -1,5 +1,11 @@
 ## Rails 4.0.0 (unreleased) ##
 
+*   Allow ActiveRecord::Base.connection_handler to have thread affinity and be
+    settable, this effectively allows ActiveRecord to be used in a multi threaded
+    setup with multiple connections to multiple dbs.
+
+    *Sam Saffron*
+
 *   `rename_column` preserves auto_increment in mysql migrations.
     Fixes #3493.
 
@@ -8,7 +14,7 @@
 *   PostgreSQL geometric type point is supported by ActiveRecord. Fixes #7324.
 
     *Martin Schuerrer*
- 
+
 *   Add suport for concurrent indexing in PostgreSQL adapter via the
     `algorithm: :concurrently` option
 
diff --git a/activerecord/lib/active_record/core.rb b/activerecord/lib/active_record/core.rb
index 013f9b92b9..1ed05c1bd4 100644
--- a/activerecord/lib/active_record/core.rb
+++ b/activerecord/lib/active_record/core.rb
@@ -77,8 +77,17 @@ module ActiveRecord
       mattr_accessor :disable_implicit_join_references, instance_writer: false
       self.disable_implicit_join_references = false
 
-      class_attribute :connection_handler, instance_writer: false
-      self.connection_handler = ConnectionAdapters::ConnectionHandler.new
+      class_attribute :default_connection_handler, instance_writer: false
+      
+      def self.connection_handler
+        Thread.current[:active_record_connection_handler] || self.default_connection_handler
+      end
+
+      def self.connection_handler=(handler)
+        Thread.current[:active_record_connection_handler] = handler
+      end
+
+      self.default_connection_handler = ConnectionAdapters::ConnectionHandler.new
     end
 
     module ClassMethods
diff --git a/activerecord/test/cases/base_test.rb b/activerecord/test/cases/base_test.rb
index f9524bc6bd..69eef500f5 100644
--- a/activerecord/test/cases/base_test.rb
+++ b/activerecord/test/cases/base_test.rb
@@ -1578,4 +1578,61 @@ class BasicsTest < ActiveRecord::TestCase
     klass = Class.new(ActiveRecord::Base)
     assert_equal ['foo'], klass.all.merge!(select: 'foo').select_values
   end
+
+  test "connection_handler can be overriden" do
+    klass = Class.new(ActiveRecord::Base)
+    orig_handler = klass.connection_handler
+    new_handler = ActiveRecord::ConnectionAdapters::ConnectionHandler.new
+    thread_connection_handler = nil
+
+    t = Thread.new do
+      klass.connection_handler = new_handler
+      thread_connection_handler = klass.connection_handler
+    end
+    t.join
+
+    assert_equal klass.connection_handler, orig_handler
+    assert_equal thread_connection_handler, new_handler
+  end
+
+  test "new threads get default the default connection handler" do
+    klass = Class.new(ActiveRecord::Base)
+    orig_handler = klass.connection_handler
+    handler = nil
+
+    t = Thread.new do
+      handler = klass.connection_handler
+    end
+    t.join
+
+    assert_equal handler, orig_handler
+    assert_equal klass.connection_handler, orig_handler
+    assert_equal klass.default_connection_handler, orig_handler
+  end
+
+  test "changing a connection handler in a main thread does not poison the other threads" do
+    klass = Class.new(ActiveRecord::Base)
+    orig_handler = klass.connection_handler
+    new_handler = ActiveRecord::ConnectionAdapters::ConnectionHandler.new
+    after_handler = nil
+    is_set = false
+
+    t = Thread.new do
+      klass.connection_handler = new_handler
+      is_set = true
+      Thread.stop
+      after_handler = klass.connection_handler
+    end
+
+    while(!is_set)
+      Thread.pass
+    end
+
+    klass.connection_handler = orig_handler
+    t.wakeup
+    t.join
+
+    assert_equal after_handler, new_handler
+    assert_equal orig_handler, klass.connection_handler
+  end
 end