Add check_current_protected_environment! to Tasks. (#222)

Author: marcw

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