fix protected env check for multi-db apps

Similarly to rails#46336, this change brings the `db:check_protected_environments` task up to parity with the rest of tasks that support multi-database applications.

Today, misconfiguration of non-primary databases in multi-db applications is not checked against the `ar_internal_metadata` table's `environment` row, posing a risk to production environments.

Similarly to my other PRs (rails#46097, rails#46336), we iterate through each of the relevant database configurations and re-establish `ActiveRecord::Base` using each configuration. This is safe to do since the `check_protected_environments!` method is only ever called from the `databases.rake` file/tasks.

**A hidden change, but critical (for the purposes that I'm working on), that is happening here** is that database configurations with `database_tasks: false` will be excluded from the logic in this task – due to the usage of the `ActiveRecord::Base.configurations.configs_for` method, which by default has `include_hidden: false`!

Author: hkdsun

activerecord/lib/active_record/tasks/database_tasks.rb

@@ -60,20 +60,16 @@ module DatabaseTasks
 
       LOCAL_HOSTS = ["127.0.0.1", "localhost"]
 
-      def check_protected_environments!
-        unless ENV["DISABLE_DATABASE_ENVIRONMENT_CHECK"]
-          current = ActiveRecord::Base.connection.migration_context.current_environment
-          stored  = ActiveRecord::Base.connection.migration_context.last_stored_environment
+      def check_protected_environments!(environment = env)
+        return if ENV["DISABLE_DATABASE_ENVIRONMENT_CHECK"]
 
-          if ActiveRecord::Base.connection.migration_context.protected_environment?
-            raise ActiveRecord::ProtectedEnvironmentError.new(stored)
-          end
-
-          if stored && stored != current
-            raise ActiveRecord::EnvironmentMismatchError.new(current: current, stored: stored)
-          end
+        original_db_config = ActiveRecord::Base.connection_db_config
+        configs_for(env_name: environment).each do |db_config|
+          ActiveRecord::Base.establish_connection(db_config)
+          check_current_protected_environment!
         end
-      rescue ActiveRecord::NoDatabaseError
+      ensure
+        ActiveRecord::Base.establish_connection(original_db_config)
       end
 
       def register_task(pattern, task)
@@ -589,6 +585,20 @@ def structure_load_flags_for(adapter)
             structure_load_flags
           end
         end
+
+        def check_current_protected_environment!
+          current = ActiveRecord::Base.connection.migration_context.current_environment
+          stored  = ActiveRecord::Base.connection.migration_context.last_stored_environment
+
+          if ActiveRecord::Base.connection.migration_context.protected_environment?
+            raise ActiveRecord::ProtectedEnvironmentError.new(stored)
+          end
+
+          if stored && stored != current
+            raise ActiveRecord::EnvironmentMismatchError.new(current: current, stored: stored)
+          end
+        rescue ActiveRecord::NoDatabaseError
+        end
     end
   end
 end

activerecord/lib/active_record/tasks/sqlite_database_tasks.rb

@@ -30,8 +30,8 @@ def drop
       end
 
       def purge
-        drop
         connection.disconnect!
+        drop
       rescue NoDatabaseError
       ensure
         create

activerecord/test/cases/tasks/database_tasks_test.rb

@@ -55,7 +55,17 @@ def assert_called_for_configs(method_name, configs, &block)
     sqlite3:    :sqlite_tasks
   }
 
-  class DatabaseTasksUtilsTask < ActiveRecord::TestCase
+  class DatabaseTasksCheckProtectedEnvironmentsTest < ActiveRecord::TestCase
+    self.use_transactional_tests = false
+
+    def setup
+      recreate_metadata_tables
+    end
+
+    def teardown
+      recreate_metadata_tables
+    end
+
     def test_raises_an_error_when_called_with_protected_environment
       protected_environments = ActiveRecord::Base.protected_environments
       current_env            = ActiveRecord::Base.connection.migration_context.current_environment
@@ -70,12 +80,12 @@ def test_raises_an_error_when_called_with_protected_environment
       ) do
         assert_not_includes protected_environments, current_env
         # Assert no error
-        ActiveRecord::Tasks::DatabaseTasks.check_protected_environments!
+        ActiveRecord::Tasks::DatabaseTasks.check_protected_environments!("arunit")
 
         ActiveRecord::Base.protected_environments = [current_env]
 
         assert_raise(ActiveRecord::ProtectedEnvironmentError) do
-          ActiveRecord::Tasks::DatabaseTasks.check_protected_environments!
+          ActiveRecord::Tasks::DatabaseTasks.check_protected_environments!("arunit")
         end
       end
     ensure
@@ -96,11 +106,11 @@ def test_raises_an_error_when_called_with_protected_environment_which_name_is_a_
       ) do
         assert_not_includes protected_environments, current_env
         # Assert no error
-        ActiveRecord::Tasks::DatabaseTasks.check_protected_environments!
+        ActiveRecord::Tasks::DatabaseTasks.check_protected_environments!("arunit")
 
         ActiveRecord::Base.protected_environments = [current_env.to_sym]
         assert_raise(ActiveRecord::ProtectedEnvironmentError) do
-          ActiveRecord::Tasks::DatabaseTasks.check_protected_environments!
+          ActiveRecord::Tasks::DatabaseTasks.check_protected_environments!("arunit")
         end
       end
     ensure
@@ -119,12 +129,89 @@ def test_raises_an_error_if_no_migrations_have_been_made
       assert_not_predicate internal_metadata, :table_exists?
 
       assert_raises(ActiveRecord::NoEnvironmentInSchemaError) do
-        ActiveRecord::Tasks::DatabaseTasks.check_protected_environments!
+        ActiveRecord::Tasks::DatabaseTasks.check_protected_environments!("arunit")
       end
     ensure
       schema_migration.delete_version("1")
       internal_metadata.create_table
     end
+
+    private
+      def recreate_metadata_tables
+        schema_migration = ActiveRecord::Base.connection.schema_migration
+        schema_migration.drop_table
+        schema_migration.create_table
+
+        internal_metadata = ActiveRecord::Base.connection.internal_metadata
+        internal_metadata.drop_table
+        internal_metadata.create_table
+      end
+  end
+
+  if current_adapter?(:SQLite3Adapter) && !in_memory_db?
+    class DatabaseTasksCheckProtectedEnvironmentsMultiDatabaseTest < ActiveRecord::TestCase
+      self.use_transactional_tests = false
+
+      def test_with_multiple_databases
+        env = ActiveRecord::ConnectionHandling::DEFAULT_ENV.call
+        with_multi_db_configurations(env) do
+          protected_environments = ActiveRecord::Base.protected_environments
+          current_env = ActiveRecord::Base.connection.migration_context.current_environment
+          assert_equal current_env, env
+
+          ActiveRecord::Base.establish_connection(:primary)
+          ActiveRecord::Base.connection.internal_metadata.create_table_and_set_flags(current_env)
+
+          ActiveRecord::Base.establish_connection(:secondary)
+          ActiveRecord::Base.connection.internal_metadata.create_table_and_set_flags(current_env)
+
+          assert_not_includes protected_environments, current_env
+          # Assert not raises
+          ActiveRecord::Tasks::DatabaseTasks.check_protected_environments!("test")
+
+          ActiveRecord::Base.establish_connection(:secondary)
+          connection = ActiveRecord::Base.connection
+          schema_migration = connection.schema_migration
+          schema_migration.create_table
+          schema_migration.create_version("1")
+
+          ActiveRecord::Base.protected_environments = [current_env.to_sym]
+          assert_raise(ActiveRecord::ProtectedEnvironmentError) do
+            ActiveRecord::Tasks::DatabaseTasks.check_protected_environments!("test")
+          end
+        ensure
+          ActiveRecord::Base.protected_environments = protected_environments
+        end
+      end
+
+      private
+        def with_multi_db_configurations(env)
+          old_configurations = ActiveRecord::Base.configurations
+          ActiveRecord::Base.configurations = {
+            env => {
+              primary: {
+                adapter: "sqlite3",
+                database: "test/fixtures/fixture_database.sqlite3",
+              },
+              secondary: {
+                adapter: "sqlite3",
+                database: "test/fixtures/fixture_database_2.sqlite3",
+              }
+            }
+          }
+
+          ActiveRecord::Base.establish_connection(:primary)
+          yield
+        ensure
+          [:primary, :secondary].each do |db|
+            ActiveRecord::Base.establish_connection(db)
+            ActiveRecord::Base.connection.schema_migration.drop_table
+            ActiveRecord::Base.connection.internal_metadata.drop_table
+          end
+          ActiveRecord::Base.configurations = old_configurations
+          ActiveRecord::Base.establish_connection(:arunit)
+        end
+    end
   end
 
   class DatabaseTasksRegisterTask < ActiveRecord::TestCase

railties/test/application/rake/multi_dbs_test.rb

@@ -443,6 +443,9 @@ def generate_models_for_animals
         require "#{app_path}/config/environment"
         ActiveRecord::Base.configurations.configs_for(env_name: Rails.env).each do |db_config|
           db_create_and_drop_namespace db_config.name, db_config.database
+        ensure
+          # secondary databases might have been created by check_protected_environments task
+          rails("db:drop:all")
         end
       end