Move calls on Base connection to methods for rake tasks

This PR aims to contain calls to `Base.connection` and
`Base.establish_connection` in active record rake tasks and the
migration code. In a follow up PR I will swap out the `Base` class for a
temporary class which will allow us to stop clobbering `Base` in the
active record rails tasks.

This work is an important step to achieve moving away from Active
Record's dependence on `Base.connection` and
`Base.establish_connection`. The reliance on `Base.connection` is
problematic for sharding support and calling `Base.establish_connection`
in the Rake tasks (without any warning message) indicates it's ok for
applications to do the same (outside these tasks, it's not).

I've vetted the approach of swapping out `Base` for a temporary class in
another branch but decided that it would be easier to demonstrate and
contain the changes if I first contained these calls. The major changes
in this PR are:

* Contain `Base.connection` in `Tasks.migration_connection` and replace
calls to each
* Contain `Base.establish_connection` in `Tasks.establish_connection`
and replace calls to each
* Add a `with_temporary_connection_for_each` method for cases where we
need to loop over each config and set the connection back afterwards
* Add a `with_temporary_connection(db_config)` method for cases where we
have one config but need to establish a new connection and set the old
one back.
* Update every place we were looping through configs to establish a
connection with the new temporary connection methods.

There are a lot of changes here but I've pulled out everything that
didn't need to be in this commit into other PRs.

Once this is merged, I'll create the next PR that replaces `Base` the
new methods in `Tasks` with a temporary connection and then we will
officially be no longer clobbering `Base` in these tasks. This also
reduces complexity because we won't need to ensure we set
`Base.connection` back at the end. Once that is working and all internal
methods are using the new temp class I'll deprecate calls on
`Base.connection` in these methods. Most applications just override the
task but not the actual methods in `Tasks` so my hope is this will be
smooth-ish. However, nothing will stop applications from still using
`Base.connection` for a very long time if they still want to clobber.

Author: eileencodes

activerecord/lib/active_record/migration.rb

@@ -137,15 +137,15 @@ class PendingMigrationError < MigrationError # :nodoc:
       ActiveRecord::Tasks::DatabaseTasks.migrate
 
       if ActiveRecord.dump_schema_after_migration
-        ActiveRecord::Tasks::DatabaseTasks.dump_schema(
-          ActiveRecord::Base.connection_db_config
-        )
+        connection = ActiveRecord::Tasks::DatabaseTasks.migration_connection
+        ActiveRecord::Tasks::DatabaseTasks.dump_schema(connection.pool.db_config)
       end
     end
 
     def initialize(message = nil, pending_migrations: nil)
       if pending_migrations.nil?
-        pending_migrations = ActiveRecord::Base.connection.migration_context.open.pending_migrations
+        connection = ActiveRecord::Tasks::DatabaseTasks.migration_connection
+        pending_migrations = connection.migration_context.open.pending_migrations
       end
 
       super(message || detailed_migration_message(pending_migrations))
@@ -165,6 +165,10 @@ def detailed_migration_message(pending_migrations)
 
         message
       end
+
+      def connection
+        ActiveRecord::Tasks::DatabaseTasks.migration_connection
+      end
   end
 
   class ConcurrentMigrationError < MigrationError # :nodoc:
@@ -624,6 +628,10 @@ def build_watcher(&block)
           paths = all_configs.flat_map { |config| config.migrations_paths || Migrator.migrations_paths }.uniq
           @file_watcher.new([], paths.index_with(["rb"]), &block)
         end
+
+        def connection
+          ActiveRecord::Tasks::DatabaseTasks.migration_connection
+        end
     end
 
     class << self
@@ -635,7 +643,7 @@ def nearest_delegate # :nodoc:
       end
 
       # Raises <tt>ActiveRecord::PendingMigrationError</tt> error if any migrations are pending.
-      def check_pending!(connection = Base.connection)
+      def check_pending!(connection = ActiveRecord::Tasks::DatabaseTasks.migration_connection)
         if pending_migrations = connection.migration_context.pending_migrations
           raise ActiveRecord::PendingMigrationError.new(pending_migrations: pending_migrations)
         end
@@ -645,6 +653,7 @@ def load_schema_if_pending!
         if any_schema_needs_update?
           # Roundtrip to Rake to allow plugins to hook into database initialization.
           root = defined?(ENGINE_ROOT) ? ENGINE_ROOT : Rails.root
+
           FileUtils.cd(root) do
             Base.connection_handler.clear_all_connections!(:all)
             system("bin/rails db:test:prepare")
@@ -692,25 +701,24 @@ def any_schema_needs_update?
           end
         end
 
+        def db_configs_in_current_env
+          ActiveRecord::Base.configurations.configs_for(env_name: env)
+        end
+
         def pending_migrations
-          prev_db_config = Base.connection_db_config
           pending_migrations = []
 
-          db_configs_in_current_env.each do |db_config|
-            conn = Base.establish_connection(db_config).connection
-            if pending = conn.migration_context.open.pending_migrations
+          ActiveRecord::Tasks::DatabaseTasks.with_temporary_connection_for_each(env: env) do |connection|
+            if pending = connection.migration_context.open.pending_migrations
               pending_migrations << pending
             end
           end
 
           pending_migrations.flatten
-        ensure
-          Base.establish_connection(prev_db_config)
         end
 
-        def db_configs_in_current_env
-          current_environment = ActiveRecord::ConnectionHandling::DEFAULT_ENV.call
-          ActiveRecord::Base.configurations.configs_for(env_name: current_environment)
+        def env
+          ActiveRecord::ConnectionHandling::DEFAULT_ENV.call
         end
     end
 
@@ -893,7 +901,7 @@ def migrate(direction)
       end
 
       time = nil
-      ActiveRecord::Base.connection_pool.with_connection do |conn|
+      ActiveRecord::Tasks::DatabaseTasks.migration_connection.pool.with_connection do |conn|
         time = Benchmark.measure do
           exec_migration(conn, direction)
         end
@@ -957,7 +965,7 @@ def suppress_messages
     end
 
     def connection
-      @connection || ActiveRecord::Base.connection
+      @connection || ActiveRecord::Tasks::DatabaseTasks.migration_connection
     end
 
     def method_missing(method, *arguments, &block)
@@ -1138,8 +1146,8 @@ def initialize(migrations_paths, schema_migration = nil, internal_metadata = nil
       end
 
       @migrations_paths = migrations_paths
-      @schema_migration = schema_migration || SchemaMigration.new(ActiveRecord::Base.connection)
-      @internal_metadata = internal_metadata || InternalMetadata.new(ActiveRecord::Base.connection)
+      @schema_migration = schema_migration || SchemaMigration.new(connection)
+      @internal_metadata = internal_metadata || InternalMetadata.new(connection)
     end
 
     # Runs the migrations in the +migrations_path+.
@@ -1265,7 +1273,6 @@ def protected_environment? # :nodoc:
     end
 
     def last_stored_environment # :nodoc:
-      connection = ActiveRecord::Base.connection
       return nil unless connection.internal_metadata.enabled?
       return nil if current_version == 0
       raise NoEnvironmentInSchemaError unless connection.internal_metadata.table_exists?
@@ -1276,6 +1283,10 @@ def last_stored_environment # :nodoc:
     end
 
     private
+      def connection
+        ActiveRecord::Tasks::DatabaseTasks.migration_connection
+      end
+
       def migration_files
         paths = Array(migrations_paths)
         Dir[*paths.flat_map { |path| "#{path}/**/[0-9]*_*.rb" }]
@@ -1311,8 +1322,9 @@ class << self
 
       # For cases where a table doesn't exist like loading from schema cache
       def current_version
-        schema_migration = SchemaMigration.new(ActiveRecord::Base.connection)
-        internal_metadata = InternalMetadata.new(ActiveRecord::Base.connection)
+        connection = ActiveRecord::Tasks::DatabaseTasks.migration_connection
+        schema_migration = SchemaMigration.new(connection)
+        internal_metadata = InternalMetadata.new(connection)
 
         MigrationContext.new(migrations_paths, schema_migration, internal_metadata).current_version
       end
@@ -1388,6 +1400,10 @@ def load_migrated
     end
 
     private
+      def connection
+        ActiveRecord::Tasks::DatabaseTasks.migration_connection
+      end
+
       # Used for running a specific migration.
       def run_without_lock
         migration = migrations.detect { |m| m.version == @target_version }
@@ -1411,7 +1427,7 @@ def migrate_without_lock
       def record_environment
         return if down?
 
-        @internal_metadata[:environment] = ActiveRecord::Base.connection.migration_context.current_environment
+        @internal_metadata[:environment] = connection.migration_context.current_environment
       end
 
       def ran?(migration)
@@ -1481,23 +1497,22 @@ def down?
       # Wrap the migration in a transaction only if supported by the adapter.
       def ddl_transaction(migration, &block)
         if use_transaction?(migration)
-          Base.transaction(&block)
+          connection.transaction(&block)
         else
           yield
         end
       end
 
       def use_transaction?(migration)
-        !migration.disable_ddl_transaction && Base.connection.supports_ddl_transactions?
+        !migration.disable_ddl_transaction && connection.supports_ddl_transactions?
       end
 
       def use_advisory_lock?
-        Base.connection.advisory_locks_enabled?
+        connection.advisory_locks_enabled?
       end
 
       def with_advisory_lock
         lock_id = generate_migrator_advisory_lock_id
-        connection = ActiveRecord::Base.connection
 
         got_lock = connection.get_advisory_lock(lock_id)
         raise ConcurrentMigrationError unless got_lock
@@ -1513,7 +1528,7 @@ def with_advisory_lock
 
       MIGRATOR_SALT = 2053462845
       def generate_migrator_advisory_lock_id
-        db_name_hash = Zlib.crc32(Base.connection.current_database)
+        db_name_hash = Zlib.crc32(connection.current_database)
         MIGRATOR_SALT * db_name_hash
       end
   end