Merge pull request #37443 from eileencodes/use-database-config-objects-in-railties-dbconsole

Use DatabaseConfig objects in dbconsole

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

@@ -22,20 +22,22 @@ class ExclusiveConnectionTimeoutError < ConnectionTimeoutError
   module ConnectionAdapters
     module AbstractPool # :nodoc:
       def get_schema_cache(connection)
-        self.schema_cache ||= SchemaCache.new(connection)
-        schema_cache.connection = connection
-        schema_cache
+        @schema_cache ||= SchemaCache.new(connection)
+        @schema_cache.connection = connection
+        @schema_cache
       end
 
       def set_schema_cache(cache)
-        self.schema_cache = cache
+        @schema_cache = cache
       end
     end
 
     class NullPool # :nodoc:
       include ConnectionAdapters::AbstractPool
 
-      attr_accessor :schema_cache
+      def initialize
+        @schema_cache = nil
+      end
     end
 
     # Connection pool base class for managing Active Record database
@@ -369,11 +371,9 @@ def run
       include QueryCache::ConnectionPoolConfiguration
       include ConnectionAdapters::AbstractPool
 
-      attr_accessor :automatic_reconnect, :checkout_timeout
+      attr_accessor :automatic_reconnect, :checkout_timeout, :schema_cache
       attr_reader :db_config, :size, :reaper
 
-      delegate :schema_cache, :schema_cache=, to: :db_config
-
       # Creates a new ConnectionPool object. +spec+ is a ConnectionSpecification
       # object which describes database connection information (e.g. adapter,
       # host name, username, password, etc), as well as the maximum size for
@@ -998,15 +998,35 @@ def checkout_and_verify(c)
     # about the model. The model needs to pass a specification name to the handler,
     # in order to look up the correct connection pool.
     class ConnectionHandler
-      FINALIZER = lambda { |_| ActiveSupport::ForkTracker.check! }
-      private_constant :FINALIZER
+      def self.create_owner_to_pool # :nodoc:
+        Concurrent::Map.new(initial_capacity: 2) do |h, k|
+          # Discard the parent's connection pools immediately; we have no need
+          # of them
+          discard_unowned_pools(h)
+
+          h[k] = Concurrent::Map.new(initial_capacity: 2)
+        end
+      end
+
+      def self.unowned_pool_finalizer(pid_map) # :nodoc:
+        lambda do |_|
+          discard_unowned_pools(pid_map)
+        end
+      end
+
+      def self.discard_unowned_pools(pid_map) # :nodoc:
+        pid_map.each do |pid, pools|
+          pools.values.compact.each(&:discard!) unless pid == Process.pid
+        end
+      end
 
       def initialize
         # These caches are keyed by spec.name (ConnectionSpecification#name).
-        @owner_to_config = Concurrent::Map.new(initial_capacity: 2)
+        @owner_to_pool = ConnectionHandler.create_owner_to_pool
 
-        # Backup finalizer: if the forked child skipped Kernel#fork the early discard has not occurred
-        ObjectSpace.define_finalizer self, FINALIZER
+        # Backup finalizer: if the forked child never needed a pool, the above
+        # early discard has not occurred
+        ObjectSpace.define_finalizer self, ConnectionHandler.unowned_pool_finalizer(@owner_to_pool)
       end
 
       def prevent_writes # :nodoc:
@@ -1030,11 +1050,11 @@ def while_preventing_writes(enabled = true)
       end
 
       def connection_pool_names # :nodoc:
-        owner_to_config.keys
+        owner_to_pool.keys
       end
 
       def connection_pool_list
-        owner_to_config.values.compact.map(&:connection_pool)
+        owner_to_pool.values.compact
       end
       alias :connection_pools :connection_pool_list
 
@@ -1054,11 +1074,11 @@ def establish_connection(config)
           payload[:config] = db_config.configuration_hash
         end
 
-        owner_to_config[spec.name] = db_config
-
         message_bus.instrument("!connection.active_record", payload) do
-          db_config.connection_pool
+          owner_to_pool[spec.name] = ConnectionAdapters::ConnectionPool.new(db_config)
         end
+
+        owner_to_pool[spec.name]
       end
 
       # Returns true if there are any active connections among the connection
@@ -1123,21 +1143,42 @@ def connected?(spec_name)
       # can be used as an argument for #establish_connection, for easily
       # re-establishing the connection.
       def remove_connection(spec_name)
-        if db_config = owner_to_config.delete(spec_name)
-          db_config.disconnect!
-          db_config.configuration_hash
+        if pool = owner_to_pool.delete(spec_name)
+          pool.automatic_reconnect = false
+          pool.disconnect!
+          pool.db_config.configuration_hash
         end
       end
 
-      # Retrieving the connection pool happens a lot, so we cache it in @owner_to_config.
+      # Retrieving the connection pool happens a lot, so we cache it in @owner_to_pool.
       # This makes retrieving the connection pool O(1) once the process is warm.
       # When a connection is established or removed, we invalidate the cache.
       def retrieve_connection_pool(spec_name)
-        owner_to_config[spec_name]&.connection_pool
+        owner_to_pool.fetch(spec_name) do
+          # Check if a connection was previously established in an ancestor process,
+          # which may have been forked.
+          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.db_config.configuration_hash.merge(name: spec_name)).tap do |pool|
+              pool.schema_cache = ancestor_pool.schema_cache if ancestor_pool.schema_cache
+            end
+          else
+            owner_to_pool[spec_name] = nil
+          end
+        end
       end
 
       private
-        attr_reader :owner_to_config
+        def owner_to_pool
+          @owner_to_pool[Process.pid]
+        end
+
+        def pool_from_any_process_for(spec_name)
+          owner_to_pool = @owner_to_pool.values.reverse.find { |v| v[spec_name] }
+          owner_to_pool && owner_to_pool[spec_name]
+        end
     end
   end
 end

activerecord/lib/active_record/database_configurations/database_config.rb

@@ -6,60 +6,11 @@ class DatabaseConfigurations
     # UrlConfig respectively. It will never return a DatabaseConfig object,
     # as this is the parent class for the types of database configuration objects.
     class DatabaseConfig # :nodoc:
-      include Mutex_m
-
       attr_reader :env_name, :spec_name
 
-      attr_accessor :schema_cache
-
-      INSTANCES = ObjectSpace::WeakMap.new
-      private_constant :INSTANCES
-
-      class << self
-        def discard_pools!
-          INSTANCES.each_key(&:discard_pool!)
-        end
-      end
-
       def initialize(env_name, spec_name)
-        super()
         @env_name = env_name
         @spec_name = spec_name
-        @pool = nil
-
-        INSTANCES[self] = self
-      end
-
-      def disconnect!
-        ActiveSupport::ForkTracker.check!
-
-        return unless @pool
-
-        synchronize do
-          return unless @pool
-
-          @pool.automatic_reconnect = false
-          @pool.disconnect!
-        end
-
-        nil
-      end
-
-      def connection_pool
-        ActiveSupport::ForkTracker.check!
-
-        @pool || synchronize { @pool ||= ConnectionAdapters::ConnectionPool.new(self) }
-      end
-
-      def discard_pool!
-        return unless @pool
-
-        synchronize do
-          return unless @pool
-
-          @pool.discard!
-          @pool = nil
-        end
       end
 
       def config
@@ -108,5 +59,3 @@ def for_current_env?
     end
   end
 end
-
-ActiveSupport::ForkTracker.after_fork { ActiveRecord::DatabaseConfigurations::DatabaseConfig.discard_pools! }

activerecord/lib/active_record/test_fixtures.rb

@@ -192,7 +192,8 @@ def setup_shared_connection_pool
         ActiveRecord::Base.connection_handlers.values.each do |handler|
           if handler != writing_handler
             handler.connection_pool_names.each do |name|
-              handler.send(:owner_to_config)[name] = writing_handler.send(:owner_to_config)[name]
+              writing_connection = writing_handler.retrieve_connection_pool(name)
+              handler.send(:owner_to_pool)[name] = writing_connection
             end
           end
         end

activerecord/test/cases/fixtures_test.rb

@@ -1390,9 +1390,6 @@ def self.file_fixture_path
 
 class MultipleDatabaseFixturesTest < ActiveRecord::TestCase
   test "enlist_fixture_connections ensures multiple databases share a connection pool" do
-    old_handlers = ActiveRecord::Base.connection_handlers
-    ActiveRecord::Base.connection_handlers = {}
-
     with_temporary_connection_pool do
       ActiveRecord::Base.connects_to database: { writing: :arunit, reading: :arunit2 }
 
@@ -1409,16 +1406,17 @@ class MultipleDatabaseFixturesTest < ActiveRecord::TestCase
       assert_equal rw_conn, ro_conn
     end
   ensure
-    ActiveRecord::Base.connection_handlers = old_handlers
+    ActiveRecord::Base.connection_handlers = { writing: ActiveRecord::Base.connection_handler }
   end
 
   private
     def with_temporary_connection_pool
-      db_config = ActiveRecord::Base.connection_handler.send(:owner_to_config).fetch("primary")
-      new_pool = ActiveRecord::ConnectionAdapters::ConnectionPool.new(db_config)
+      old_pool = ActiveRecord::Base.connection_handler.retrieve_connection_pool(ActiveRecord::Base.connection_specification_name)
+      new_pool = ActiveRecord::ConnectionAdapters::ConnectionPool.new(ActiveRecord::Base.connection_pool.db_config)
+      ActiveRecord::Base.connection_handler.send(:owner_to_pool)["primary"] = new_pool
 
-      db_config.stub(:connection_pool, new_pool) do
-        yield
-      end
+      yield
+    ensure
+      ActiveRecord::Base.connection_handler.send(:owner_to_pool)["primary"] = old_pool
     end
 end

activerecord/test/cases/query_cache_test.rb

@@ -555,12 +555,13 @@ def test_clear_query_cache_is_called_on_all_connections
 
   private
     def with_temporary_connection_pool
-      db_config = ActiveRecord::Base.connection_handler.send(:owner_to_config).fetch("primary")
-      new_pool = ActiveRecord::ConnectionAdapters::ConnectionPool.new(db_config)
+      old_pool = ActiveRecord::Base.connection_handler.retrieve_connection_pool(ActiveRecord::Base.connection_specification_name)
+      new_pool = ActiveRecord::ConnectionAdapters::ConnectionPool.new(ActiveRecord::Base.connection_pool.db_config)
+      ActiveRecord::Base.connection_handler.send(:owner_to_pool)["primary"] = new_pool
 
-      db_config.stub(:connection_pool, new_pool) do
-        yield
-      end
+      yield
+    ensure
+      ActiveRecord::Base.connection_handler.send(:owner_to_pool)["primary"] = old_pool
     end
 
     def middleware(&app)

activerecord/test/cases/unconnected_test.rb

@@ -13,7 +13,10 @@ def setup
     @specification = ActiveRecord::Base.remove_connection
 
     # Clear out connection info from other pids (like a fork parent) too
-    ActiveRecord::DatabaseConfigurations::DatabaseConfig.discard_pools!
+    pool_map = ActiveRecord::Base.connection_handler.instance_variable_get(:@owner_to_pool)
+    (pool_map.keys - [Process.pid]).each do |other_pid|
+      pool_map.delete(other_pid)
+    end
   end
 
   teardown do

railties/lib/rails/commands/dbconsole/dbconsole_command.rb

@@ -17,7 +17,7 @@ def initialize(options = {})
     def start
       ENV["RAILS_ENV"] ||= @options[:environment] || environment
 
-      case config[:adapter]
+      case db_config.adapter
       when /^(jdbc)?mysql/
         args = {
           host: "--host",
@@ -38,7 +38,7 @@ def start
           args << "-p"
         end
 
-        args << config[:database]
+        args << db_config.database
 
         find_cmd_and_exec(["mysql", "mysql5"], *args)
 
@@ -47,14 +47,14 @@ def start
         ENV["PGHOST"]     = config[:host] if config[:host]
         ENV["PGPORT"]     = config[:port].to_s if config[:port]
         ENV["PGPASSWORD"] = config[:password].to_s if config[:password] && @options[:include_password]
-        find_cmd_and_exec("psql", config[:database])
+        find_cmd_and_exec("psql", db_config.database)
 
       when "sqlite3"
         args = []
 
         args << "-#{@options[:mode]}" if @options[:mode]
         args << "-header" if @options[:header]
-        args << File.expand_path(config[:database], Rails.respond_to?(:root) ? Rails.root : nil)
+        args << File.expand_path(db_config.database, Rails.respond_to?(:root) ? Rails.root : nil)
 
         find_cmd_and_exec("sqlite3", *args)
 
@@ -64,7 +64,7 @@ def start
         if config[:username]
           logon = config[:username].dup
           logon << "/#{config[:password]}" if config[:password] && @options[:include_password]
-          logon << "@#{config[:database]}" if config[:database]
+          logon << "@#{db_config.database}" if db_config.database
         end
 
         find_cmd_and_exec("sqlplus", logon)
@@ -72,7 +72,7 @@ def start
       when "sqlserver"
         args = []
 
-        args += ["-D", "#{config[:database]}"] if config[:database]
+        args += ["-D", "#{db_config.database}"] if db_config.database
         args += ["-U", "#{config[:username]}"] if config[:username]
         args += ["-P", "#{config[:password]}"] if config[:password]
 
@@ -85,23 +85,29 @@ def start
         find_cmd_and_exec("sqsh", *args)
 
       else
-        abort "Unknown command-line client for #{config[:database]}."
+        abort "Unknown command-line client for #{db_config.database}."
       end
     end
 
     def config
-      @config ||= begin
-        # We need to check whether the user passed the database the
-        # first time around to show a consistent error message to people
-        # relying on 2-level database configuration.
-        if @options[:database] && configurations[database].blank?
-          raise ActiveRecord::AdapterNotSpecified, "'#{database}' database is not configured. Available configuration: #{configurations.inspect}"
-        elsif configurations[environment].blank? && configurations[database].blank?
-          raise ActiveRecord::AdapterNotSpecified, "'#{environment}' database is not configured. Available configuration: #{configurations.inspect}"
-        else
-          (configurations[database] || configurations[environment].presence).symbolize_keys
-        end
+      db_config.configuration_hash
+    end
+
+    def db_config
+      return @db_config if @db_config
+
+      # We need to check whether the user passed the database the
+      # first time around to show a consistent error message to people
+      # relying on 2-level database configuration.
+
+      @db_config = configurations.configs_for(env_name: environment, spec_name: database)
+
+      unless @db_config
+        raise ActiveRecord::AdapterNotSpecified,
+          "'#{database}' database is not configured for '#{environment}'. Available configuration: #{configurations.inspect}"
       end
+
+      @db_config
     end
 
     def environment

railties/test/commands/dbconsole_test.rb

@@ -4,6 +4,7 @@
 require "minitest/mock"
 require "rails/command"
 require "rails/commands/dbconsole/dbconsole_command"
+require "active_record/database_configurations"
 
 class Rails::DBConsoleTest < ActiveSupport::TestCase
   def setup
@@ -231,7 +232,7 @@ def test_specifying_a_missing_database
         Rails::Command.invoke(:dbconsole, ["--db", "i_do_not_exist"])
       end
 
-      assert_includes e.message, "'i_do_not_exist' database is not configured."
+      assert_includes e.message, "'i_do_not_exist' database is not configured for 'test'."
     end
   end
 
@@ -241,7 +242,7 @@ def test_specifying_a_missing_environment
         Rails::Command.invoke(:dbconsole)
       end
 
-      assert_includes e.message, "'test' database is not configured."
+      assert_includes e.message, "'primary' database is not configured for 'test'."
     end
   end
 
@@ -294,8 +295,10 @@ def find_cmd_and_exec(*args)
     attr_reader :dbconsole
 
     def start(config = {}, argv = [])
+      hash_config = ActiveRecord::DatabaseConfigurations::HashConfig.new("test", "primary", config)
+
       @dbconsole = make_dbconsole.new(parse_arguments(argv))
-      @dbconsole.stub(:config, config) do
+      @dbconsole.stub(:db_config, hash_config) do
         capture_abort { @dbconsole.start }
       end
     end