fix Rails 4 multisite so connection handlers don't keep spinning up new pools

vendor/gems/rails_multisite/Gemfile

@@ -1,11 +1,14 @@
 source 'https://rubygems.org'
 
 group :test do
-  gem 'rails'  
-  gem 'rspec'  
+  gem 'rails'
+  gem 'rspec'
   gem 'activerecord'
   gem 'sqlite3'
 end
 
+gem 'guard'
+gem 'guard-rspec'
+
 # Specify your gem's dependencies in rails_multisite.gemspec
 gemspec

vendor/gems/rails_multisite/Guardfile

@@ -0,0 +1,9 @@
+# A sample Guardfile
+# More info at https://github.com/guard/guard#readme
+
+guard :rspec do
+  watch(%r{^spec/.+_spec\.rb$})
+  watch(%r{^lib/rails_multisite/(.+)\.rb$})     { |m| "spec/#{m[1]}_spec.rb" }
+  watch('spec/spec_helper.rb')  { "spec" }
+end
+

vendor/gems/rails_multisite/lib/rails_multisite/connection_management.rb

@@ -16,16 +16,23 @@ module RailsMultisite
           handler = @@connection_handlers[spec]
           unless handler
             handler = ActiveRecord::ConnectionAdapters::ConnectionHandler.new
+            if rails4?
+              handler.establish_connection(ActiveRecord::Base, spec)
+            end
             @@connection_handlers[spec] = handler
           end
         else
           handler = @@default_connection_handler
+          if rails4? && !@@established_default
+            handler.establish_connection(ActiveRecord::Base, spec)
+            @@established_default = true
+          end
         end
+
         ActiveRecord::Base.connection_handler = handler
-        if rails4?
-          ActiveRecord::Base.connection_handler.establish_connection(ActiveRecord::Base, spec)
-        else
-          ActiveRecord::Base.connection_handler.establish_connection("ActiveRecord::Base", spec)
+
+        unless rails4?
+          handler.establish_connection("ActiveRecord::Base", spec)
         end
       end
     end
@@ -100,17 +107,19 @@ module RailsMultisite
         @@host_spec_cache[host] = @@default_spec
       end
 
+      @@default_connection_handler = ActiveRecord::Base.connection_handler
+
       # inject our connection_handler pool
       # WARNING MONKEY PATCH
       #
       # see: https://github.com/rails/rails/issues/8344#issuecomment-10800848
-      #
-      @@default_connection_handler = ActiveRecord::Base.connection_handler
-      ActiveRecord::Base.send :include, NewConnectionHandler if ActiveRecord::VERSION::MAJOR == 3
-
-      ActiveRecord::Base.connection_handler = @@default_connection_handler
+      if ActiveRecord::VERSION::MAJOR == 3
+        ActiveRecord::Base.send :include, NewConnectionHandler
+        ActiveRecord::Base.connection_handler = @@default_connection_handler
+      end
 
       @@connection_handlers = {}
+      @@established_default = false
     end
 
     module NewConnectionHandler

vendor/gems/rails_multisite/spec/connection_management_spec.rb

@@ -1,16 +1,26 @@
 require 'spec_helper'
 require 'rails_multisite'
 
+class Person < ActiveRecord::Base; end
+
 describe RailsMultisite::ConnectionManagement do
 
   subject { RailsMultisite::ConnectionManagement }
 
+  def with_connection(db)
+    subject.establish_connection(db: db)
+    yield ActiveRecord::Base.connection.raw_connection
+  ensure
+    ActiveRecord::Base.connection_handler.clear_active_connections!
+  end
+
   context 'default' do
     its(:all_dbs) { should == ['default']}
 
     context 'current' do
       before do
         subject.establish_connection(db: 'default')
+        ActiveRecord::Base.establish_connection
       end
 
       its(:current_db) { should == 'default' }
@@ -25,6 +35,7 @@ describe RailsMultisite::ConnectionManagement do
       subject.config_filename = "spec/fixtures/two_dbs.yml"
       subject.load_settings!
     end
+
     its(:all_dbs) { should == ['default', 'second']}
 
     context 'second db' do
@@ -36,6 +47,53 @@ describe RailsMultisite::ConnectionManagement do
       its(:current_hostname) { should == "second.localhost" }
     end
 
+    context 'data partitioning' do
+      after do
+        ['default','second'].each do |db|
+          with_connection(db) do |cnn|
+            cnn.execute("drop table people") rescue nil
+          end
+        end
+      end
+
+      it 'partitions data correctly' do
+        col1 = []
+        col2 = []
+
+        ['default','second'].map do |db|
+
+          with_connection(db) do |cnn|
+            cnn.execute("create table if not exists people(id INTEGER PRIMARY KEY AUTOINCREMENT, db)")
+          end
+        end
+
+        SQLite3::Database.query_log.clear
+
+        5.times do
+          ['default','second'].map do |db|
+            Thread.new do
+              with_connection(db) do |cnn|
+                Person.create!(db: db)
+              end
+            end
+          end.map(&:join)
+        end
+
+        lists = []
+        ['default', 'second'].each do |db|
+          with_connection(db) do |cnn|
+            lists << Person.order(:id).to_a.map{|p| [p.id, p.db]}
+          end
+        end
+
+        lists[1].should == (1..5).map{|id| [id, "second"]}
+        lists[0].should == (1..5).map{|id| [id, "default"]}
+
+        # puts SQLite3::Database.query_log.map{|args, caller, oid| "#{oid} #{args.join.inspect}"}.join("\n")
+
+      end
+    end
+
   end
 
 end

vendor/gems/rails_multisite/spec/fixtures/database.yml

@@ -1,6 +1,6 @@
 test:
   adapter: sqlite3
-  database: rails_multisite
+  database: "db.test"
   timeout: 5000
   host_names:
     - default.localhost

vendor/gems/rails_multisite/spec/fixtures/two_dbs.yml

@@ -1,9 +1,6 @@
 second:
   adapter: sqlite3
-  database: second_db
-  username: username
-  password: password
-  db_id: 1
+  database: "db1.test"
   host_names:
     - second.localhost
     - 2nd.localhost

vendor/gems/rails_multisite/spec/spec_helper.rb

@@ -6,6 +6,27 @@ require 'active_record'
 ENV["RAILS_ENV"] ||= 'test'
 RSpec.configure do |config|
 
+  require 'sqlite3'
+  class SQLite3::Database
+    def self.query_log
+      @@query_log ||= []
+    end
+
+    alias_method :old_execute, :execute
+    alias_method :old_prepare, :prepare
+
+    def execute(*args,&blk)
+      self.class.query_log << [args, caller, Thread.current.object_id]
+      old_execute(*args,&blk)
+    end
+
+    def prepare(*args,&blk)
+      self.class.query_log << [args, caller, Thread.current.object_id]
+      old_prepare(*args,&blk)
+    end
+
+  end
+
   config.color_enabled = true
 
   config.before(:suite) do