Fix schema_up_to_date connection

This method was failing to return the connection to the correct state
which would break tests when eager loading the application. What's
happening is that when we check `schema_up_to_date` when the application
is eager loaded all the models are eager loaded as well so it sees that
ApplicationRecord has a connection already but that connection is set to
the wrong one. It doesn't get replaced with the right one so we see
errors about looking for tables in the wrong database. This doesn't
happen when eager loading is off because ApplicationRecord isn't set and
gets initialized with the correct connection (from it's
`connects_to/establish_connection` call).

I also refactored this to pull `needs_update` into a method and make the
loop easier to read.

Author: eileencodes

activerecord/lib/active_record/migration.rb

@@ -642,11 +642,7 @@ def check_pending!(connection = Base.connection)
       end
 
       def load_schema_if_pending!
-        needs_update = !db_configs_in_current_env.all? do |db_config|
-          Tasks::DatabaseTasks.schema_up_to_date?(db_config, ActiveRecord.schema_format)
-        end
-
-        if needs_update
+        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
@@ -690,6 +686,12 @@ def check_pending_migrations # :nodoc:
       end
 
       private
+        def any_schema_needs_update?
+          !db_configs_in_current_env.all? do |db_config|
+            Tasks::DatabaseTasks.schema_up_to_date?(db_config, ActiveRecord.schema_format)
+          end
+        end
+
         def pending_migrations
           prev_db_config = Base.connection_db_config
           pending_migrations = []

activerecord/lib/active_record/tasks/database_tasks.rb

@@ -378,6 +378,7 @@ def load_schema(db_config, format = ActiveRecord.schema_format, file = nil) # :n
       end
 
       def schema_up_to_date?(configuration, format = ActiveRecord.schema_format, file = nil)
+        original_db_config = ActiveRecord::Base.connection_db_config
         db_config = resolve_configuration(configuration)
 
         file ||= schema_dump_path(db_config)
@@ -391,6 +392,8 @@ def schema_up_to_date?(configuration, format = ActiveRecord.schema_format, file
         return false unless connection.internal_metadata.table_exists?
 
         connection.internal_metadata[:schema_sha1] == schema_sha1(file)
+      ensure
+        ActiveRecord::Base.establish_connection(original_db_config)
       end
 
       def reconstruct_from_schema(db_config, format = ActiveRecord.schema_format, file = nil) # :nodoc:

railties/test/application/rake/multi_dbs_test.rb

@@ -1259,6 +1259,49 @@ class TwoMigration < ActiveRecord::Migration::Current
           end
         end
       end
+
+      test "after schema is loaded test run on the correct connections" do
+        require "#{app_path}/config/environment"
+        app_file "config/database.yml", <<-YAML
+          development:
+            primary:
+              database: db/default.sqlite3
+              adapter: sqlite3
+            animals:
+              database: db/development_animals.sqlite3
+              adapter: sqlite3
+              migrations_paths: db/animals_migrate
+          test:
+            primary:
+              database: db/default_test.sqlite3
+              adapter: sqlite3
+            animals:
+              database: db/test_animals.sqlite3
+              adapter: sqlite3
+              migrations_paths: db/animals_migrate
+        YAML
+
+        Dir.chdir(app_path) do
+          generate_models_for_animals
+
+          File.open("test/models/book_test.rb", "w") do |file|
+            file.write(<<~EOS)
+              require "test_helper"
+
+              class BookTest < ActiveSupport::TestCase
+                test "a book" do
+                  assert Book.first
+                end
+              end
+            EOS
+          end
+
+          rails "db:migrate"
+          rails "db:schema:dump"
+          output = rails "test"
+          assert_match(/1 runs, 1 assertions, 0 failures, 0 errors, 0 skips/, output)
+        end
+      end
     end
   end
 end