From 845ba3013e3a4726c1f9fe892695ca2cea7ea059 Mon Sep 17 00:00:00 2001
From: eileencodes <eileencodes@gmail.com>
Date: Thu, 27 Oct 2022 11:28:40 -0400
Subject: [PATCH] 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.
---
 activerecord/lib/active_record/migration.rb   | 12 +++---
 .../lib/active_record/tasks/database_tasks.rb |  3 ++
 .../test/application/rake/multi_dbs_test.rb   | 43 +++++++++++++++++++
 3 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/activerecord/lib/active_record/migration.rb b/activerecord/lib/active_record/migration.rb
index 1cd8093b9ff..ebd18b28fe5 100644
--- a/activerecord/lib/active_record/migration.rb
+++ b/activerecord/lib/active_record/migration.rb
@@ -642,11 +642,7 @@ module ActiveRecord
       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 @@ module ActiveRecord
       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 = []
diff --git a/activerecord/lib/active_record/tasks/database_tasks.rb b/activerecord/lib/active_record/tasks/database_tasks.rb
index 507ee2819bd..3acf2635533 100644
--- a/activerecord/lib/active_record/tasks/database_tasks.rb
+++ b/activerecord/lib/active_record/tasks/database_tasks.rb
@@ -378,6 +378,7 @@ module ActiveRecord
       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 @@ module ActiveRecord
         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:
diff --git a/railties/test/application/rake/multi_dbs_test.rb b/railties/test/application/rake/multi_dbs_test.rb
index 6a6916b7cf0..53e243dd595 100644
--- a/railties/test/application/rake/multi_dbs_test.rb
+++ b/railties/test/application/rake/multi_dbs_test.rb
@@ -1259,6 +1259,49 @@ module ApplicationTests
           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