diff --git a/lib/clickhouse-activerecord/tasks.rb b/lib/clickhouse-activerecord/tasks.rb index b36d9e62..6b4a5a85 100644 --- a/lib/clickhouse-activerecord/tasks.rb +++ b/lib/clickhouse-activerecord/tasks.rb @@ -88,8 +88,34 @@ def migrate ActiveRecord::Migration.verbose = verbose_was end + def check_current_protected_environment!(db_config, migration_class = ActiveRecord::Migration) + with_temporary_pool(db_config, migration_class) 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 + private + def with_temporary_pool(db_config, migration_class, 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 establish_master_connection establish_connection @configuration end