Merge pull request rails#46270 from eileencodes/extract-calls-on-base-for-rake-tasks

Move calls on Base connection to methods for rake tasks

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