Merge pull request #24844 from arthurnn/arthurnn/conn

Refactor connection handler

11 files changed, 109 insertions, 92 deletions

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 e389d818fd..4ba8ee2706 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
@@ -778,8 +778,7 @@ module ActiveRecord
     end
 
     # ConnectionHandler is a collection of ConnectionPool objects. It is used
-    # for keeping separate connection pools for Active Record models that connect
-    # to different databases.
+    # for keeping separate connection pools that connect to different databases.
     #
     # For example, suppose that you have 5 models, with the following hierarchy:
     #
@@ -821,6 +820,10 @@ module ActiveRecord
     # ConnectionHandler accessible via ActiveRecord::Base.connection_handler.
     # All Active Record models use this handler to determine the connection pool that they
     # should use.
+    #
+    # The ConnectionHandler class is not coupled with the Active models, as it has no knowlodge
+    # about the model. The model, needs to pass a specification name to the handler,
+    # in order to lookup the correct connection pool.
     class ConnectionHandler
       def initialize
         # These caches are keyed by klass.name, NOT klass. Keying them by klass
@@ -829,9 +832,6 @@ module ActiveRecord
         @owner_to_pool = Concurrent::Map.new(:initial_capacity => 2) do |h,k|
           h[k] = Concurrent::Map.new(:initial_capacity => 2)
         end
-        @class_to_pool = Concurrent::Map.new(:initial_capacity => 2) do |h,k|
-          h[k] = Concurrent::Map.new
-        end
       end
 
       def connection_pool_list
@@ -839,10 +839,8 @@ module ActiveRecord
       end
       alias :connection_pools :connection_pool_list
 
-      def establish_connection(owner, spec)
-        @class_to_pool.clear
-        raise RuntimeError, "Anonymous class is not allowed." unless owner.name
-        owner_to_pool[owner.name] = ConnectionAdapters::ConnectionPool.new(spec)
+      def establish_connection(spec)
+        owner_to_pool[spec.name] = ConnectionAdapters::ConnectionPool.new(spec)
       end
 
       # Returns true if there are any active connections among the connection
@@ -873,18 +871,18 @@ module ActiveRecord
       # active or defined connection: if it is the latter, it will be
       # opened and set as the active connection for the class it was defined
       # for (not necessarily the current class).
-      def retrieve_connection(klass) #:nodoc:
-        pool = retrieve_connection_pool(klass)
-        raise ConnectionNotEstablished, "No connection pool for #{klass}" unless pool
+      def retrieve_connection(spec_name) #:nodoc:
+        pool = retrieve_connection_pool(spec_name)
+        raise ConnectionNotEstablished, "No connection pool with id #{spec_name} found." unless pool
         conn = pool.connection
-        raise ConnectionNotEstablished, "No connection for #{klass} in connection pool" unless conn
+        raise ConnectionNotEstablished, "No connection for #{spec_name} in connection pool" unless conn
         conn
       end
 
       # Returns true if a connection that's accessible to this class has
       # already been opened.
-      def connected?(klass)
-        conn = retrieve_connection_pool(klass)
+      def connected?(spec_name)
+        conn = retrieve_connection_pool(spec_name)
         conn && conn.connected?
       end
 
@@ -892,9 +890,8 @@ module ActiveRecord
       # connection and the defined connection (if they exist). The result
       # can be used as an argument for establish_connection, for easily
       # re-establishing the connection.
-      def remove_connection(owner)
-        if pool = owner_to_pool.delete(owner.name)
-          @class_to_pool.clear
+      def remove_connection(spec_name)
+        if pool = owner_to_pool.delete(spec_name)
           pool.automatic_reconnect = false
           pool.disconnect!
           pool.spec.config
@@ -910,14 +907,18 @@ module ActiveRecord
       # #fetch is significantly slower than #[]. So in the nil case, no caching will
       # take place, but that's ok since the nil case is not the common one that we wish
       # to optimise for.
-      def retrieve_connection_pool(klass)
-        class_to_pool[klass.name] ||= begin
-          until pool = pool_for(klass)
-            klass = klass.superclass
-            break unless klass <= Base
+      def retrieve_connection_pool(spec_name)
+        owner_to_pool.fetch(spec_name) do
+          if ancestor_pool = pool_from_any_process_for(spec_name)
+            # A connection was established in an ancestor process that must have
+            # subsequently forked. We can't reuse the connection, but we can copy
+            # the specification and establish a new connection with it.
+            establish_connection(ancestor_pool.spec).tap do |pool|
+              pool.schema_cache = ancestor_pool.schema_cache if ancestor_pool.schema_cache
+            end
+          else
+            owner_to_pool[spec_name] = nil
           end
-
-          class_to_pool[klass.name] = pool
         end
       end
 
@@ -927,28 +928,9 @@ module ActiveRecord
         @owner_to_pool[Process.pid]
       end
 
-      def class_to_pool
-        @class_to_pool[Process.pid]
-      end
-
-      def pool_for(owner)
-        owner_to_pool.fetch(owner.name) {
-          if ancestor_pool = pool_from_any_process_for(owner)
-            # A connection was established in an ancestor process that must have
-            # subsequently forked. We can't reuse the connection, but we can copy
-            # the specification and establish a new connection with it.
-            establish_connection(owner, ancestor_pool.spec).tap do |pool|
-              pool.schema_cache = ancestor_pool.schema_cache if ancestor_pool.schema_cache
-            end
-          else
-            owner_to_pool[owner.name] = nil
-          end
-        }
-      end
-
-      def pool_from_any_process_for(owner)
-        owner_to_pool = @owner_to_pool.values.find { |v| v[owner.name] }
-        owner_to_pool && owner_to_pool[owner.name]
+      def pool_from_any_process_for(spec_name)
+        owner_to_pool = @owner_to_pool.values.find { |v| v[spec_name] }
+        owner_to_pool && owner_to_pool[spec_name]
       end
     end
   end
diff --git a/activerecord/lib/active_record/connection_adapters/connection_specification.rb b/activerecord/lib/active_record/connection_adapters/connection_specification.rb
index 4bc6447368..901c98b22b 100644
--- a/activerecord/lib/active_record/connection_adapters/connection_specification.rb
+++ b/activerecord/lib/active_record/connection_adapters/connection_specification.rb
@@ -3,10 +3,10 @@ require 'uri'
 module ActiveRecord
   module ConnectionAdapters
     class ConnectionSpecification #:nodoc:
-      attr_reader :config, :adapter_method
+      attr_reader :name, :config, :adapter_method
 
-      def initialize(config, adapter_method)
-        @config, @adapter_method = config, adapter_method
+      def initialize(name, config, adapter_method)
+        @name, @config, @adapter_method = name, config, adapter_method
       end
 
       def initialize_dup(original)
@@ -164,7 +164,7 @@ module ActiveRecord
         #   spec.config
         #   # => { "host" => "localhost", "database" => "foo", "adapter" => "sqlite3" }
         #
-        def spec(config)
+        def spec(config, name = nil)
           spec = resolve(config).symbolize_keys
 
           raise(AdapterNotSpecified, "database configuration does not specify adapter") unless spec.key?(:adapter)
@@ -179,7 +179,14 @@ module ActiveRecord
           end
 
           adapter_method = "#{spec[:adapter]}_connection"
-          ConnectionSpecification.new(spec, adapter_method)
+
+          name ||=
+            if config.is_a?(Symbol)
+              config.to_s
+            else
+              "primary"
+            end
+          ConnectionSpecification.new(name, spec, adapter_method)
         end
 
         private
diff --git a/activerecord/lib/active_record/connection_handling.rb b/activerecord/lib/active_record/connection_handling.rb
index a8b3d03ba5..ba763149cc 100644
--- a/activerecord/lib/active_record/connection_handling.rb
+++ b/activerecord/lib/active_record/connection_handling.rb
@@ -45,16 +45,20 @@ module ActiveRecord
     # The exceptions AdapterNotSpecified, AdapterNotFound and +ArgumentError+
     # may be returned on an error.
     def establish_connection(spec = nil)
+      raise RuntimeError, "Anonymous class is not allowed." unless name
+
       spec     ||= DEFAULT_ENV.call.to_sym
       resolver =   ConnectionAdapters::ConnectionSpecification::Resolver.new configurations
-      spec     =   resolver.spec(spec)
+      # TODO: uses name on establish_connection, for backwards compatibility
+      spec     =   resolver.spec(spec, self == Base ? "primary" : name)
+      self.connection_specification_name = spec.name
 
       unless respond_to?(spec.adapter_method)
         raise AdapterNotFound, "database configuration specifies nonexistent #{spec.config[:adapter]} adapter"
       end
 
       remove_connection
-      connection_handler.establish_connection self, spec
+      connection_handler.establish_connection spec
     end
 
     class MergeAndResolveDefaultUrlConfig # :nodoc:
@@ -87,6 +91,17 @@ module ActiveRecord
       retrieve_connection
     end
 
+    attr_writer :connection_specification_name
+
+    # Return the specification id from this class otherwise look it up
+    # in the parent.
+    def connection_specification_name
+      unless defined?(@connection_specification_name)
+        @connection_specification_name = self == Base ? "primary" : superclass.connection_specification_name
+      end
+      @connection_specification_name
+    end
+
     def connection_id
       ActiveRecord::RuntimeRegistry.connection_id ||= Thread.current.object_id
     end
@@ -106,20 +121,20 @@ module ActiveRecord
     end
 
     def connection_pool
-      connection_handler.retrieve_connection_pool(self) or raise ConnectionNotEstablished
+      connection_handler.retrieve_connection_pool(connection_specification_name) or raise ConnectionNotEstablished
     end
 
     def retrieve_connection
-      connection_handler.retrieve_connection(self)
+      connection_handler.retrieve_connection(connection_specification_name)
     end
 
     # Returns +true+ if Active Record is connected.
     def connected?
-      connection_handler.connected?(self)
+      connection_handler.connected?(connection_specification_name)
     end
 
-    def remove_connection(klass = self)
-      connection_handler.remove_connection(klass)
+    def remove_connection(name = connection_specification_name)
+      connection_handler.remove_connection(name)
     end
 
     def clear_cache! # :nodoc:
diff --git a/activerecord/lib/active_record/core.rb b/activerecord/lib/active_record/core.rb
index 5d74631e32..f936e865e4 100644
--- a/activerecord/lib/active_record/core.rb
+++ b/activerecord/lib/active_record/core.rb
@@ -257,7 +257,7 @@ module ActiveRecord
       # Returns the Arel engine.
       def arel_engine # :nodoc:
         @arel_engine ||=
-          if Base == self || connection_handler.retrieve_connection_pool(self)
+          if Base == self || connection_handler.retrieve_connection_pool(connection_specification_name)
             self
           else
             superclass.arel_engine
diff --git a/activerecord/lib/active_record/tasks/database_tasks.rb b/activerecord/lib/active_record/tasks/database_tasks.rb
index 9aea5b360b..0df46d54df 100644
--- a/activerecord/lib/active_record/tasks/database_tasks.rb
+++ b/activerecord/lib/active_record/tasks/database_tasks.rb
@@ -117,10 +117,10 @@ module ActiveRecord
       end
 
       def create_all
-        old_pool = ActiveRecord::Base.connection_handler.retrieve_connection_pool(ActiveRecord::Base)
+        old_pool = ActiveRecord::Base.connection_handler.retrieve_connection_pool(ActiveRecord::Base.connection_specification_name)
         each_local_configuration { |configuration| create configuration }
         if old_pool
-          ActiveRecord::Base.connection_handler.establish_connection(ActiveRecord::Base, old_pool.spec)
+          ActiveRecord::Base.connection_handler.establish_connection(old_pool.spec)
         end
       end
 
diff --git a/activerecord/test/cases/connection_adapters/adapter_leasing_test.rb b/activerecord/test/cases/connection_adapters/adapter_leasing_test.rb
index 580568c8ac..c7ca428ab7 100644
--- a/activerecord/test/cases/connection_adapters/adapter_leasing_test.rb
+++ b/activerecord/test/cases/connection_adapters/adapter_leasing_test.rb
@@ -37,7 +37,7 @@ module ActiveRecord
       end
 
       def test_close
-        pool = Pool.new(ConnectionSpecification.new({}, nil))
+        pool = Pool.new(ConnectionSpecification.new("primary", {}, nil))
         pool.insert_connection_for_test! @adapter
         @adapter.pool = pool
 
diff --git a/activerecord/test/cases/connection_adapters/connection_handler_test.rb b/activerecord/test/cases/connection_adapters/connection_handler_test.rb
index 9b1865e8bb..fc5ca8865b 100644
--- a/activerecord/test/cases/connection_adapters/connection_handler_test.rb
+++ b/activerecord/test/cases/connection_adapters/connection_handler_test.rb
@@ -4,43 +4,41 @@ module ActiveRecord
   module ConnectionAdapters
     class ConnectionHandlerTest < ActiveRecord::TestCase
       def setup
-        @klass    = Class.new(Base)   { def self.name; 'klass';    end }
-        @subklass = Class.new(@klass) { def self.name; 'subklass'; end }
-
         @handler = ConnectionHandler.new
-        @pool    = @handler.establish_connection(@klass, Base.connection_pool.spec)
+        resolver = ConnectionAdapters::ConnectionSpecification::Resolver.new Base.configurations
+        @spec_name = "primary"
+        @pool    = @handler.establish_connection(resolver.spec(:arunit, @spec_name))
+      end
+
+      def test_establish_connection_uses_spec_name
+        config = {"readonly" => {"adapter" => 'sqlite3'}}
+        resolver = ConnectionAdapters::ConnectionSpecification::Resolver.new(config)
+        spec =   resolver.spec(:readonly)
+        @handler.establish_connection(spec)
+
+        assert_not_nil @handler.retrieve_connection_pool('readonly')
+      ensure
+        @handler.remove_connection('readonly')
       end
 
       def test_retrieve_connection
-        assert @handler.retrieve_connection(@klass)
+        assert @handler.retrieve_connection(@spec_name)
       end
 
       def test_active_connections?
         assert !@handler.active_connections?
-        assert @handler.retrieve_connection(@klass)
+        assert @handler.retrieve_connection(@spec_name)
         assert @handler.active_connections?
         @handler.clear_active_connections!
         assert !@handler.active_connections?
       end
 
-      def test_retrieve_connection_pool_with_ar_base
-        assert_nil @handler.retrieve_connection_pool(ActiveRecord::Base)
-      end
-
       def test_retrieve_connection_pool
-        assert_not_nil @handler.retrieve_connection_pool(@klass)
-      end
-
-      def test_retrieve_connection_pool_uses_superclass_when_no_subclass_connection
-        assert_not_nil @handler.retrieve_connection_pool(@subklass)
+        assert_not_nil @handler.retrieve_connection_pool(@spec_name)
       end
 
-      def test_retrieve_connection_pool_uses_superclass_pool_after_subclass_establish_and_remove
-        sub_pool = @handler.establish_connection(@subklass, Base.connection_pool.spec)
-        assert_same sub_pool, @handler.retrieve_connection_pool(@subklass)
-
-        @handler.remove_connection @subklass
-        assert_same @pool, @handler.retrieve_connection_pool(@subklass)
+      def test_retrieve_connection_pool_with_invalid_id
+        assert_nil @handler.retrieve_connection_pool("foo")
       end
 
       def test_connection_pools
@@ -79,7 +77,7 @@ module ActiveRecord
 
           pid = fork {
             rd.close
-            pool = @handler.retrieve_connection_pool(@klass)
+            pool = @handler.retrieve_connection_pool(@spec_name)
             wr.write Marshal.dump pool.schema_cache.size
             wr.close
             exit!
diff --git a/activerecord/test/cases/connection_adapters/connection_specification_test.rb b/activerecord/test/cases/connection_adapters/connection_specification_test.rb
index ea2196cda2..d204fce59f 100644
--- a/activerecord/test/cases/connection_adapters/connection_specification_test.rb
+++ b/activerecord/test/cases/connection_adapters/connection_specification_test.rb
@@ -4,7 +4,7 @@ module ActiveRecord
   module ConnectionAdapters
     class ConnectionSpecificationTest < ActiveRecord::TestCase
       def test_dup_deep_copy_config
-        spec = ConnectionSpecification.new({ :a => :b }, "bar")
+        spec = ConnectionSpecification.new("primary", { :a => :b }, "bar")
         assert_not_equal(spec.config.object_id, spec.dup.config.object_id)
       end
     end
diff --git a/activerecord/test/cases/connection_pool_test.rb b/activerecord/test/cases/connection_pool_test.rb
index efa3e0455e..a45ee281c7 100644
--- a/activerecord/test/cases/connection_pool_test.rb
+++ b/activerecord/test/cases/connection_pool_test.rb
@@ -335,11 +335,10 @@ module ActiveRecord
       # is called with an anonymous class
       def test_anonymous_class_exception
         anonymous = Class.new(ActiveRecord::Base)
-        handler = ActiveRecord::Base.connection_handler
 
-        assert_raises(RuntimeError) {
-          handler.establish_connection anonymous, nil
-        }
+        assert_raises(RuntimeError) do
+          anonymous.establish_connection
+        end
       end
 
       def test_pool_sets_connection_schema_cache
diff --git a/activerecord/test/cases/connection_specification/resolver_test.rb b/activerecord/test/cases/connection_specification/resolver_test.rb
index 358b6ad537..3bddaf32ec 100644
--- a/activerecord/test/cases/connection_specification/resolver_test.rb
+++ b/activerecord/test/cases/connection_specification/resolver_test.rb
@@ -116,6 +116,15 @@ module ActiveRecord
             "encoding" => "utf8" }, spec)
         end
 
+        def test_spec_name_on_key_lookup
+          spec = spec(:readonly, 'readonly' => {'adapter' => 'sqlite3'})
+          assert_equal "readonly", spec.name
+        end
+
+        def test_spec_name_with_inline_config
+          spec = spec({'adapter' => 'sqlite3'})
+          assert_equal "primary", spec.name, "should default to primary id"
+        end
       end
     end
   end
diff --git a/activerecord/test/cases/multiple_db_test.rb b/activerecord/test/cases/multiple_db_test.rb
index af4183a601..a4fbf579a1 100644
--- a/activerecord/test/cases/multiple_db_test.rb
+++ b/activerecord/test/cases/multiple_db_test.rb
@@ -24,6 +24,13 @@ class MultipleDbTest < ActiveRecord::TestCase
     assert_equal(ActiveRecord::Base.connection, Entrant.connection)
   end
 
+  def test_swapping_the_connection
+    old_spec_name, Course.connection_specification_name = Course.connection_specification_name, "primary"
+    assert_equal(Entrant.connection, Course.connection)
+  ensure
+    Course.connection_specification_name = old_spec_name
+  end
+
   def test_find
     c1 = Course.find(1)
     assert_equal "Ruby Development", c1.name
@@ -89,8 +96,8 @@ class MultipleDbTest < ActiveRecord::TestCase
   end
 
   def test_connection
-    assert_equal Entrant.arel_engine.connection, Bird.arel_engine.connection
-    assert_not_equal Entrant.arel_engine.connection, Course.arel_engine.connection
+    assert_equal Entrant.arel_engine.connection.object_id, Bird.arel_engine.connection.object_id
+    assert_not_equal Entrant.arel_engine.connection.object_id, Course.arel_engine.connection.object_id
   end
 
   unless in_memory_db?