Fix pending migration checks for multi-db

Co-authored-by: Guo Xiang Tan <[email protected]>

Fixes rails#37524

Author: hkdsun

activerecord/lib/active_record/migration.rb

@@ -598,7 +598,7 @@ def call(env)
         @mutex.synchronize do
           @watcher ||= build_watcher do
             @needs_check = true
-            ActiveRecord::Migration.check_pending!(connection)
+            ActiveRecord::Migration.check_pending_migrations
             @needs_check = false
           end
 
@@ -614,13 +614,11 @@ def call(env)
 
       private
         def build_watcher(&block)
-          paths = Array(connection.migration_context.migrations_paths)
+          current_environment = ActiveRecord::ConnectionHandling::DEFAULT_ENV.call
+          all_configs = ActiveRecord::Base.configurations.configs_for(env_name: current_environment)
+          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::Base.connection
-        end
     end
 
     class << self
@@ -638,7 +636,7 @@ def check_pending!(connection = Base.connection)
 
       def load_schema_if_pending!
         current_db_config = Base.connection_db_config
-        all_configs = ActiveRecord::Base.configurations.configs_for(env_name: Rails.env)
+        all_configs = db_configs_in_current_env
 
         needs_update = !all_configs.all? do |db_config|
           Tasks::DatabaseTasks.schema_up_to_date?(db_config, ActiveRecord.schema_format)
@@ -656,7 +654,8 @@ def load_schema_if_pending!
         # Establish a new connection, the old database may be gone (db:test:prepare uses purge)
         Base.establish_connection(current_db_config)
 
-        check_pending!
+        # Ensure all migrations have succeeded
+        check_pending_migrations(db_configs: all_configs)
       end
 
       def maintain_test_schema! # :nodoc:
@@ -681,6 +680,22 @@ def migrate(direction)
       def disable_ddl_transaction!
         @disable_ddl_transaction = true
       end
+
+      def check_pending_migrations(db_configs: db_configs_in_current_env) # :nodoc:
+        prev_db_config = Base.connection_db_config
+        db_configs.each do |db_config|
+          Base.establish_connection(db_config)
+          check_pending!
+        end
+      ensure
+        Base.establish_connection(prev_db_config)
+      end
+
+      private
+        def db_configs_in_current_env
+          current_environment = ActiveRecord::ConnectionHandling::DEFAULT_ENV.call
+          ActiveRecord::Base.configurations.configs_for(env_name: current_environment)
+        end
     end
 
     def disable_ddl_transaction # :nodoc:

activerecord/test/cases/migration/pending_migrations_test.rb

@@ -6,11 +6,16 @@ module ActiveRecord
   class Migration
     if current_adapter?(:SQLite3Adapter) && !in_memory_db?
       class PendingMigrationsTest < ActiveRecord::TestCase
+        self.use_transactional_tests = false
+
         setup do
           @migration_dir = Dir.mktmpdir("activerecord-migrations-")
+          @secondary_migration_dir = Dir.mktmpdir("activerecord-migrations-secondary-")
+          @original_configurations = ActiveRecord::Base.configurations
+          ActiveRecord::Base.configurations = base_config
+          ActiveRecord::Base.establish_connection(:primary)
 
           file = ActiveRecord::Base.connection.raw_connection.filename
-          @conn = ActiveRecord::Base.establish_connection adapter: "sqlite3", database: ":memory:", migrations_paths: @migration_dir
           source_db = SQLite3::Database.new file
           dest_db = ActiveRecord::Base.connection.raw_connection
           backup = SQLite3::Backup.new(dest_db, "main", source_db, "main")
@@ -23,19 +28,20 @@ class PendingMigrationsTest < ActiveRecord::TestCase
         end
 
         teardown do
-          @conn.release_connection if @conn
-          ActiveRecord::Base.establish_connection :arunit
+          ActiveRecord::Base.configurations = @original_configurations
+          ActiveRecord::Base.establish_connection(:arunit)
           FileUtils.rm_rf(@migration_dir)
+          FileUtils.rm_rf(@secondary_migration_dir)
         end
 
         def run_migrations
           migrator = Base.connection.migration_context
           capture(:stdout) { migrator.migrate }
         end
 
-        def create_migration(number, name)
+        def create_migration(number, name, migration_dir: nil)
           filename = "#{number}_#{name.underscore}.rb"
-          File.write(File.join(@migration_dir, filename), <<~RUBY)
+          File.write(File.join(migration_dir || @migration_dir, filename), <<~RUBY)
             class #{name.classify} < ActiveRecord::Migration::Current
             end
           RUBY
@@ -86,7 +92,7 @@ def test_okay_with_no_migrations
         def test_understands_migrations_created_out_of_order
           # With a prior file before even initialization
           create_migration "05", "create_bar"
-          run_migrations
+          quietly { run_migrations }
 
           check_pending = CheckPending.new(@app)
 
@@ -100,6 +106,65 @@ def test_understands_migrations_created_out_of_order
             check_pending.call({})
           end
         end
+
+        def test_with_multiple_database
+          create_migration "01", "create_bar", migration_dir: @secondary_migration_dir
+
+          assert_raises ActiveRecord::PendingMigrationError do
+            CheckPending.new(@app).call({})
+          end
+
+          begin
+            ActiveRecord::Base.establish_connection(:secondary)
+            quietly { run_migrations }
+          end
+          ActiveRecord::Base.establish_connection(:primary)
+
+          @app.expect :call, nil, [{}]
+          CheckPending.new(@app).call({})
+          @app.verify
+
+          # Now check exclusion if database_tasks is set to false for the db_config
+          create_migration "02", "create_foo", migration_dir: @secondary_migration_dir
+          assert_raises ActiveRecord::PendingMigrationError do
+            CheckPending.new(@app).call({})
+          end
+
+          new_config = base_config
+          new_config[ActiveRecord::ConnectionHandling::DEFAULT_ENV.call][:secondary][:database_tasks] = false
+          ActiveRecord::Base.configurations = new_config
+
+          @app.expect :call, nil, [{}]
+          CheckPending.new(@app).call({})
+          @app.verify
+        end
+
+        def test_with_stdlib_logger
+          old, ActiveRecord::Base.logger = ActiveRecord::Base.logger, ::Logger.new($stdout)
+          quietly do
+            assert_nothing_raised { ActiveRecord::Migration::CheckPending.new(Proc.new { }).call({}) }
+          end
+        ensure
+          ActiveRecord::Base.logger = old
+        end
+
+        private
+          def base_config
+            {
+              ActiveRecord::ConnectionHandling::DEFAULT_ENV.call => {
+                primary: {
+                  adapter: "sqlite3",
+                  database: "test/fixtures/fixture_database.sqlite3",
+                  migrations_paths: @migration_dir
+                },
+                secondary: {
+                  adapter: "sqlite3",
+                  database: "test/fixtures/fixture_database.sqlite3",
+                  migrations_paths: @secondary_migration_dir
+                }
+              }
+            }
+          end
       end
     end
   end