activerecord 6.1.7 → 7.2.0

data/lib/active_record/tasks/database_tasks.rb

@@ -6,14 +6,16 @@ module ActiveRecord
   module Tasks # :nodoc:
     class DatabaseNotSupported < StandardError; end # :nodoc:
 
+    # = Active Record \DatabaseTasks
+    #
     # ActiveRecord::Tasks::DatabaseTasks is a utility class, which encapsulates
     # logic behind common tasks used to manage database and migrations.
     #
-    # The tasks defined here are used with Rails commands provided by Active Record.
+    # The tasks defined here are used with \Rails commands provided by Active Record.
     #
     # In order to use DatabaseTasks, a few config values need to be set. All the needed
-    # config values are set by Rails already, so it's necessary to do it only if you
-    # want to change the defaults or when you want to use Active Record outside of Rails
+    # config values are set by \Rails already, so it's necessary to do it only if you
+    # want to change the defaults or when you want to use Active Record outside of \Rails
     # (in such case after configuring the database tasks, you can also use the rake tasks
     # defined in Active Record).
     #
@@ -27,7 +29,7 @@ module ActiveRecord
     # * +seed_loader+: an object which will load seeds, it needs to respond to the +load_seed+ method.
     # * +root+: a path to the root of the application.
     #
-    # Example usage of DatabaseTasks outside Rails could look as such:
+    # Example usage of DatabaseTasks outside \Rails could look as such:
     #
     #   include ActiveRecord::Tasks
     #   DatabaseTasks.database_configuration = YAML.load_file('my_database_config.yml')
@@ -39,35 +41,33 @@ module ActiveRecord
       ##
       # :singleton-method:
       # Extra flags passed to database CLI tool (mysqldump/pg_dump) when calling db:schema:dump
+      # It can be used as a string/array (the typical case) or a hash (when you use multiple adapters)
+      # Example:
+      #   ActiveRecord::Tasks::DatabaseTasks.structure_dump_flags = {
+      #     mysql2: ['--no-defaults', '--skip-add-drop-table'],
+      #     postgres: '--no-tablespaces'
+      #   }
       mattr_accessor :structure_dump_flags, instance_accessor: false
 
       ##
       # :singleton-method:
       # Extra flags passed to database CLI tool when calling db:schema:load
+      # It can be used as a string/array (the typical case) or a hash (when you use multiple adapters)
       mattr_accessor :structure_load_flags, instance_accessor: false
 
       extend self
 
-      attr_writer :current_config, :db_dir, :migrations_paths, :fixtures_path, :root, :env, :seed_loader
-      deprecate :current_config=
+      attr_writer :db_dir, :migrations_paths, :fixtures_path, :root, :env, :seed_loader
       attr_accessor :database_configuration
 
       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
-
-          if ActiveRecord::Base.connection.migration_context.protected_environment?
-            raise ActiveRecord::ProtectedEnvironmentError.new(stored)
-          end
+      def check_protected_environments!(environment = env)
+        return if ENV["DISABLE_DATABASE_ENVIRONMENT_CHECK"]
 
-          if stored && stored != current
-            raise ActiveRecord::EnvironmentMismatchError.new(current: current, stored: stored)
-          end
+        configs_for(env_name: environment).each do |db_config|
+          check_current_protected_environment!(db_config)
         end
-      rescue ActiveRecord::NoDatabaseError
       end
 
       def register_task(pattern, task)
@@ -76,6 +76,7 @@ module ActiveRecord
       end
 
       register_task(/mysql/,        "ActiveRecord::Tasks::MySQLDatabaseTasks")
+      register_task(/trilogy/,      "ActiveRecord::Tasks::MySQLDatabaseTasks")
       register_task(/postgresql/,   "ActiveRecord::Tasks::PostgreSQLDatabaseTasks")
       register_task(/sqlite/,       "ActiveRecord::Tasks::SQLiteDatabaseTasks")
 
@@ -103,11 +104,6 @@ module ActiveRecord
         @env ||= Rails.env
       end
 
-      def spec
-        @spec ||= "primary"
-      end
-      deprecate spec: "please use name instead"
-
       def name
         @name ||= "primary"
       end
@@ -116,18 +112,6 @@ module ActiveRecord
         @seed_loader ||= Rails.application
       end
 
-      def current_config(options = {})
-        if options.has_key?(:config)
-          @current_config = options[:config]
-        else
-          env_name = options[:env] || env
-          name = options[:spec] || "primary"
-
-          @current_config ||= ActiveRecord::Base.configurations.configs_for(env_name: env_name, name: name)&.configuration_hash
-        end
-      end
-      deprecate :current_config
-
       def create(configuration, *arguments)
         db_config = resolve_configuration(configuration)
         database_adapter_for(db_config, *arguments).create
@@ -141,28 +125,20 @@ module ActiveRecord
       end
 
       def create_all
-        old_pool = ActiveRecord::Base.connection_handler.retrieve_connection_pool(ActiveRecord::Base.connection_specification_name)
+        db_config = migration_connection.pool.db_config
+
         each_local_configuration { |db_config| create(db_config) }
-        if old_pool
-          ActiveRecord::Base.connection_handler.establish_connection(old_pool.db_config)
-        end
+
+        migration_class.establish_connection(db_config)
       end
 
-      def setup_initial_database_yaml
+      def setup_initial_database_yaml # :nodoc:
         return {} unless defined?(Rails)
 
-        begin
-          Rails.application.config.load_database_yaml
-        rescue
-          unless ActiveRecord::Base.suppress_multiple_database_warning
-            $stderr.puts "Rails couldn't infer whether you are using multiple databases from your database.yml and can't generate the tasks for the non-primary databases. If you'd like to use this feature, please simplify your ERB."
-          end
-
-          {}
-        end
+        Rails.application.config.load_database_yaml
       end
 
-      def for_each(databases)
+      def for_each(databases) # :nodoc:
         return {} unless defined?(Rails)
 
         database_configs = ActiveRecord::DatabaseConfigurations.new(databases).configs_for(env_name: Rails.env)
@@ -171,12 +147,14 @@ module ActiveRecord
         return if database_configs.count == 1
 
         database_configs.each do |db_config|
+          next unless db_config.database_tasks?
+
           yield db_config.name
         end
       end
 
-      def raise_for_multi_db(environment = env, command:)
-        db_configs = ActiveRecord::Base.configurations.configs_for(env_name: environment)
+      def raise_for_multi_db(environment = env, command:) # :nodoc:
+        db_configs = configs_for(env_name: environment)
 
         if db_configs.count > 1
           dbs_list = []
@@ -191,7 +169,44 @@ module ActiveRecord
 
       def create_current(environment = env, name = nil)
         each_current_configuration(environment, name) { |db_config| create(db_config) }
-        ActiveRecord::Base.establish_connection(environment.to_sym)
+
+        migration_class.establish_connection(environment.to_sym)
+      end
+
+      def prepare_all
+        seed = false
+
+        each_current_configuration(env) do |db_config|
+          with_temporary_pool(db_config) do
+            begin
+              database_initialized = migration_connection_pool.schema_migration.table_exists?
+            rescue ActiveRecord::NoDatabaseError
+              create(db_config)
+              retry
+            end
+
+            unless database_initialized
+              if File.exist?(schema_dump_path(db_config))
+                load_schema(db_config, ActiveRecord.schema_format, nil)
+              end
+
+              seed = true
+            end
+          end
+        end
+
+        each_current_environment(env) do |environment|
+          db_configs_with_versions(environment).sort.each do |version, db_configs|
+            db_configs.each do |db_config|
+              with_temporary_pool(db_config) do
+                migrate(version)
+                dump_schema(db_config) if ActiveRecord.dump_schema_after_migration
+              end
+            end
+          end
+        end
+
+        load_seed if seed
       end
 
       def drop(configuration, *arguments)
@@ -215,51 +230,73 @@ module ActiveRecord
       end
 
       def truncate_tables(db_config)
-        ActiveRecord::Base.establish_connection(db_config)
-
-        connection = ActiveRecord::Base.connection
-        connection.truncate_tables(*connection.tables)
+        with_temporary_connection(db_config) do |conn|
+          conn.truncate_tables(*conn.tables)
+        end
       end
       private :truncate_tables
 
       def truncate_all(environment = env)
-        ActiveRecord::Base.configurations.configs_for(env_name: environment).each do |db_config|
+        configs_for(env_name: environment).each do |db_config|
           truncate_tables(db_config)
         end
       end
 
-      def migrate
-        check_target_version
-
+      def migrate(version = nil)
         scope = ENV["SCOPE"]
         verbose_was, Migration.verbose = Migration.verbose, verbose?
 
-        Base.connection.migration_context.migrate(target_version) do |migration|
-          scope.blank? || scope == migration.scope
+        check_target_version
+
+        migration_connection_pool.migration_context.migrate(target_version) do |migration|
+          if version.blank?
+            scope.blank? || scope == migration.scope
+          else
+            migration.version == version
+          end
+        end.tap do |migrations_ran|
+          Migration.write("No migrations ran. (using #{scope} scope)") if scope.present? && migrations_ran.empty?
         end
 
-        ActiveRecord::Base.clear_cache!
+        migration_connection_pool.schema_cache.clear!
       ensure
         Migration.verbose = verbose_was
       end
 
+      def db_configs_with_versions(environment = env) # :nodoc:
+        db_configs_with_versions = Hash.new { |h, k| h[k] = [] }
+
+        with_temporary_pool_for_each(env: environment) do |pool|
+          db_config = pool.db_config
+          versions_to_run = pool.migration_context.pending_migration_versions
+          target_version = ActiveRecord::Tasks::DatabaseTasks.target_version
+
+          versions_to_run.each do |version|
+            next if target_version && target_version != version
+            db_configs_with_versions[version] << db_config
+          end
+        end
+
+        db_configs_with_versions
+      end
+
       def migrate_status
-        unless ActiveRecord::Base.connection.schema_migration.table_exists?
+        unless migration_connection_pool.schema_migration.table_exists?
           Kernel.abort "Schema migrations table does not exist yet."
         end
 
         # output
-        puts "\ndatabase: #{ActiveRecord::Base.connection_db_config.database}\n\n"
+        puts "\ndatabase: #{migration_connection_pool.db_config.database}\n\n"
         puts "#{'Status'.center(8)}  #{'Migration ID'.ljust(14)}  Migration Name"
         puts "-" * 50
-        ActiveRecord::Base.connection.migration_context.migrations_status.each do |status, version, name|
+        migration_connection_pool.migration_context.migrations_status.each do |status, version, name|
           puts "#{status.center(8)}  #{version.ljust(14)}  #{name}"
         end
         puts
       end
 
       def check_target_version
-        if target_version && !(Migration::MigrationFilenameRegexp.match?(ENV["VERSION"]) || /\A\d+\z/.match?(ENV["VERSION"]))
+        if target_version && !Migration.valid_version_format?(ENV["VERSION"])
           raise "Invalid format of target version: `VERSION=#{ENV['VERSION']}`"
         end
       end
@@ -269,7 +306,7 @@ module ActiveRecord
       end
 
       def charset_current(env_name = env, db_name = name)
-        db_config = ActiveRecord::Base.configurations.configs_for(env_name: env_name, name: db_name)
+        db_config = configs_for(env_name: env_name, name: db_name)
         charset(db_config)
       end
 
@@ -279,7 +316,7 @@ module ActiveRecord
       end
 
       def collation_current(env_name = env, db_name = name)
-        db_config = ActiveRecord::Base.configurations.configs_for(env_name: env_name, name: db_name)
+        db_config = configs_for(env_name: env_name, name: db_name)
         collation(db_config)
       end
 
@@ -299,27 +336,30 @@ module ActiveRecord
 
       def purge_current(environment = env)
         each_current_configuration(environment) { |db_config| purge(db_config) }
-        ActiveRecord::Base.establish_connection(environment.to_sym)
+
+        migration_class.establish_connection(environment.to_sym)
       end
 
       def structure_dump(configuration, *arguments)
         db_config = resolve_configuration(configuration)
         filename = arguments.delete_at(0)
-        database_adapter_for(db_config, *arguments).structure_dump(filename, structure_dump_flags)
+        flags = structure_dump_flags_for(db_config.adapter)
+        database_adapter_for(db_config, *arguments).structure_dump(filename, flags)
       end
 
       def structure_load(configuration, *arguments)
         db_config = resolve_configuration(configuration)
         filename = arguments.delete_at(0)
-        database_adapter_for(db_config, *arguments).structure_load(filename, structure_load_flags)
+        flags = structure_load_flags_for(db_config.adapter)
+        database_adapter_for(db_config, *arguments).structure_load(filename, flags)
       end
 
-      def load_schema(db_config, format = ActiveRecord::Base.schema_format, file = nil) # :nodoc:
-        file ||= dump_filename(db_config.name, format)
+      def load_schema(db_config, format = ActiveRecord.schema_format, file = nil) # :nodoc:
+        file ||= schema_dump_path(db_config, format)
+        return unless file
 
         verbose_was, Migration.verbose = Migration.verbose, verbose? && ENV["VERBOSE"]
         check_schema_file(file)
-        ActiveRecord::Base.establish_connection(db_config)
 
         case format
         when :ruby
@@ -329,112 +369,111 @@ module ActiveRecord
         else
           raise ArgumentError, "unknown format #{format.inspect}"
         end
-        ActiveRecord::InternalMetadata.create_table
-        ActiveRecord::InternalMetadata[:environment] = db_config.env_name
-        ActiveRecord::InternalMetadata[:schema_sha1] = schema_sha1(file)
+
+        migration_connection_pool.internal_metadata.create_table_and_set_flags(db_config.env_name, schema_sha1(file))
       ensure
         Migration.verbose = verbose_was
       end
 
-      def schema_up_to_date?(configuration, format = ActiveRecord::Base.schema_format, file = nil, environment = nil, name = nil)
+      def schema_up_to_date?(configuration, format = ActiveRecord.schema_format, file = nil)
         db_config = resolve_configuration(configuration)
 
-        if environment || name
-          ActiveSupport::Deprecation.warn("`environment` and `name` will be removed as parameters in 7.0.0, you may now pass an ActiveRecord::DatabaseConfigurations::DatabaseConfig as `configuration` instead.")
-        end
-
-        name ||= db_config.name
-
-        file ||= dump_filename(name, format)
+        file ||= schema_dump_path(db_config)
 
-        return true unless File.exist?(file)
+        return true unless file && File.exist?(file)
 
-        ActiveRecord::Base.establish_connection(db_config)
+        with_temporary_pool(db_config) do |pool|
+          internal_metadata = pool.internal_metadata
+          return false unless internal_metadata.enabled?
+          return false unless internal_metadata.table_exists?
 
-        return false unless ActiveRecord::InternalMetadata.enabled?
-        return false unless ActiveRecord::InternalMetadata.table_exists?
-
-        ActiveRecord::InternalMetadata[:schema_sha1] == schema_sha1(file)
+          internal_metadata[:schema_sha1] == schema_sha1(file)
+        end
       end
 
-      def reconstruct_from_schema(db_config, format = ActiveRecord::Base.schema_format, file = nil) # :nodoc:
-        file ||= dump_filename(db_config.name, format)
+      def reconstruct_from_schema(db_config, format = ActiveRecord.schema_format, file = nil) # :nodoc:
+        file ||= schema_dump_path(db_config, format)
 
-        check_schema_file(file)
-
-        ActiveRecord::Base.establish_connection(db_config)
+        check_schema_file(file) if file
 
-        if schema_up_to_date?(db_config, format, file)
-          truncate_tables(db_config)
-        else
-          purge(db_config)
+        with_temporary_pool(db_config, clobber: true) do
+          if schema_up_to_date?(db_config, format, file)
+            truncate_tables(db_config) unless ENV["SKIP_TEST_DATABASE_TRUNCATE"]
+          else
+            purge(db_config)
+            load_schema(db_config, format, file)
+          end
+        rescue ActiveRecord::NoDatabaseError
+          create(db_config)
           load_schema(db_config, format, file)
         end
-      rescue ActiveRecord::NoDatabaseError
-        create(db_config)
-        load_schema(db_config, format, file)
       end
 
-      def dump_schema(db_config, format = ActiveRecord::Base.schema_format) # :nodoc:
+      def dump_schema(db_config, format = ActiveRecord.schema_format) # :nodoc:
+        return unless db_config.schema_dump
+
         require "active_record/schema_dumper"
-        filename = dump_filename(db_config.name, format)
-        connection = ActiveRecord::Base.connection
+        filename = schema_dump_path(db_config, format)
+        return unless filename
 
         FileUtils.mkdir_p(db_dir)
         case format
         when :ruby
           File.open(filename, "w:utf-8") do |file|
-            ActiveRecord::SchemaDumper.dump(ActiveRecord::Base.connection, file)
+            ActiveRecord::SchemaDumper.dump(migration_connection_pool, file)
           end
         when :sql
           structure_dump(db_config, filename)
-          if connection.schema_migration.table_exists?
+          if migration_connection_pool.schema_migration.table_exists?
             File.open(filename, "a") do |f|
-              f.puts connection.dump_schema_information
+              f.puts migration_connection.dump_schema_information
               f.print "\n"
             end
           end
         end
       end
 
-      def schema_file(format = ActiveRecord::Base.schema_format)
-        File.join(db_dir, schema_file_type(format))
-      end
+      def schema_dump_path(db_config, format = ActiveRecord.schema_format)
+        return ENV["SCHEMA"] if ENV["SCHEMA"]
 
-      def schema_file_type(format = ActiveRecord::Base.schema_format)
-        case format
-        when :ruby
-          "schema.rb"
-        when :sql
-          "structure.sql"
-        end
-      end
+        filename = db_config.schema_dump(format)
+        return unless filename
 
-      def dump_filename(db_config_name, format = ActiveRecord::Base.schema_format)
-        filename = if ActiveRecord::Base.configurations.primary?(db_config_name)
-          schema_file_type(format)
+        if File.dirname(filename) == ActiveRecord::Tasks::DatabaseTasks.db_dir
+          filename
         else
-          "#{db_config_name}_#{schema_file_type(format)}"
+          File.join(ActiveRecord::Tasks::DatabaseTasks.db_dir, filename)
         end
-
-        ENV["SCHEMA"] || File.join(ActiveRecord::Tasks::DatabaseTasks.db_dir, filename)
       end
 
-      def cache_dump_filename(db_config_name, schema_cache_path: nil)
-        filename = if ActiveRecord::Base.configurations.primary?(db_config_name)
-          "schema_cache.yml"
+      def cache_dump_filename(db_config_or_name, schema_cache_path: nil)
+        if db_config_or_name.is_a?(DatabaseConfigurations::DatabaseConfig)
+          schema_cache_path ||
+            db_config_or_name.schema_cache_path ||
+            schema_cache_env ||
+            db_config_or_name.default_schema_cache_path(ActiveRecord::Tasks::DatabaseTasks.db_dir)
         else
-          "#{db_config_name}_schema_cache.yml"
-        end
+          ActiveRecord.deprecator.warn(<<~MSG.squish)
+            Passing a database name to `cache_dump_filename` is deprecated and will be removed in Rails 8.0. Pass a
+            `ActiveRecord::DatabaseConfigurations::DatabaseConfig` object instead.
+          MSG
+
+          filename = if ActiveRecord::Base.configurations.primary?(db_config_or_name)
+            "schema_cache.yml"
+          else
+            "#{db_config_or_name}_schema_cache.yml"
+          end
 
-        schema_cache_path || ENV["SCHEMA_CACHE"] || File.join(ActiveRecord::Tasks::DatabaseTasks.db_dir, filename)
+          schema_cache_path || schema_cache_env || File.join(ActiveRecord::Tasks::DatabaseTasks.db_dir, filename)
+        end
       end
 
-      def load_schema_current(format = ActiveRecord::Base.schema_format, file = nil, environment = env)
+      def load_schema_current(format = ActiveRecord.schema_format, file = nil, environment = env)
         each_current_configuration(environment) do |db_config|
-          load_schema(db_config, format, file)
+          with_temporary_connection(db_config) do
+            load_schema(db_config, format, file)
+          end
         end
-        ActiveRecord::Base.establish_connection(environment.to_sym)
       end
 
       def check_schema_file(filename)
@@ -457,17 +496,70 @@ module ActiveRecord
 
       # Dumps the schema cache in YAML format for the connection into the file
       #
-      # ==== Examples:
-      #   ActiveRecord::Tasks::DatabaseTasks.dump_schema_cache(ActiveRecord::Base.connection, "tmp/schema_dump.yaml")
-      def dump_schema_cache(conn, filename)
-        conn.schema_cache.dump_to(filename)
+      # ==== Examples
+      #   ActiveRecord::Tasks::DatabaseTasks.dump_schema_cache(ActiveRecord::Base.lease_connection, "tmp/schema_dump.yaml")
+      def dump_schema_cache(conn_or_pool, filename)
+        conn_or_pool.schema_cache.dump_to(filename)
       end
 
       def clear_schema_cache(filename)
         FileUtils.rm_f filename, verbose: false
       end
 
+      def with_temporary_pool_for_each(env: ActiveRecord::Tasks::DatabaseTasks.env, name: nil, clobber: false, &block) # :nodoc:
+        if name
+          db_config = ActiveRecord::Base.configurations.configs_for(env_name: env, name: name)
+          with_temporary_pool(db_config, clobber: clobber, &block)
+        else
+          ActiveRecord::Base.configurations.configs_for(env_name: env, name: name).each do |db_config|
+            with_temporary_pool(db_config, clobber: clobber, &block)
+          end
+        end
+      end
+
+      def with_temporary_connection(db_config, clobber: false, &block) # :nodoc:
+        with_temporary_pool(db_config, clobber: clobber) do |pool|
+          pool.with_connection(&block)
+        end
+      end
+
+      def migration_class # :nodoc:
+        ActiveRecord::Base
+      end
+
+      def migration_connection # :nodoc:
+        migration_class.lease_connection
+      end
+
+      def migration_connection_pool # :nodoc:
+        migration_class.connection_pool
+      end
+
       private
+        def schema_cache_env
+          if ENV["SCHEMA_CACHE"]
+            ActiveRecord.deprecator.warn(<<~MSG.squish)
+              Setting `ENV["SCHEMA_CACHE"]` is deprecated and will be removed in Rails 8.0.
+              Configure the `:schema_cache_path` in the database configuration instead.
+            MSG
+
+            nil
+          end
+        end
+
+        def with_temporary_pool(db_config, clobber: false)
+          original_db_config = migration_class.connection_db_config
+          pool = migration_class.connection_handler.establish_connection(db_config, clobber: clobber)
+
+          yield pool
+        ensure
+          migration_class.connection_handler.establish_connection(original_db_config, clobber: clobber)
+        end
+
+        def configs_for(**options)
+          Base.configurations.configs_for(**options)
+        end
+
         def resolve_configuration(configuration)
           Base.configurations.resolve(configuration)
         end
@@ -488,7 +580,7 @@ module ActiveRecord
         end
 
         def class_for_adapter(adapter)
-          _key, task = @tasks.each_pair.detect { |pattern, _task| adapter[pattern] }
+          _key, task = @tasks.reverse_each.detect { |pattern, _task| adapter[pattern] }
           unless task
             raise DatabaseNotSupported, "Rake tasks not supported by '#{adapter}' adapter"
           end
@@ -496,11 +588,8 @@ module ActiveRecord
         end
 
         def each_current_configuration(environment, name = nil)
-          environments = [environment]
-          environments << "test" if environment == "development" && !ENV["SKIP_TEST_DATABASE"] && !ENV["DATABASE_URL"]
-
-          environments.each do |env|
-            ActiveRecord::Base.configurations.configs_for(env_name: env).each do |db_config|
+          each_current_environment(environment) do |env|
+            configs_for(env_name: env).each do |db_config|
               next if name && name != db_config.name
 
               yield db_config
@@ -508,8 +597,14 @@ module ActiveRecord
           end
         end
 
+        def each_current_environment(environment, &block)
+          environments = [environment]
+          environments << "test" if environment == "development" && !ENV["SKIP_TEST_DATABASE"] && !ENV["DATABASE_URL"]
+          environments.each(&block)
+        end
+
         def each_local_configuration
-          ActiveRecord::Base.configurations.configs_for.each do |db_config|
+          configs_for.each do |db_config|
             next unless db_config.database
 
             if local_database?(db_config)
@@ -526,7 +621,40 @@ module ActiveRecord
         end
 
         def schema_sha1(file)
-          Digest::SHA1.hexdigest(File.read(file))
+          OpenSSL::Digest::SHA1.hexdigest(File.read(file))
+        end
+
+        def structure_dump_flags_for(adapter)
+          if structure_dump_flags.is_a?(Hash)
+            structure_dump_flags[adapter.to_sym]
+          else
+            structure_dump_flags
+          end
+        end
+
+        def structure_load_flags_for(adapter)
+          if structure_load_flags.is_a?(Hash)
+            structure_load_flags[adapter.to_sym]
+          else
+            structure_load_flags
+          end
+        end
+
+        def check_current_protected_environment!(db_config)
+          with_temporary_pool(db_config) do |pool|
+            migration_context = pool.migration_context
+            current = migration_context.current_environment
+            stored  = migration_context.last_stored_environment
+
+            if 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