diff --git a/lib/msf/core/db_manager/migration.rb b/lib/msf/core/db_manager/migration.rb index f38b827409d3..32f5e0579eb8 100644 --- a/lib/msf/core/db_manager/migration.rb +++ b/lib/msf/core/db_manager/migration.rb @@ -23,33 +23,22 @@ def add_rails_engine_migration_paths # @see ActiveRecord::MigrationContext.migrate def migrate(config=nil, verbose=false) ran = [] - # Rails 5 changes ActiveRecord parents means to migrate outside - # the `rake` task framework has to dig a little lower into ActiveRecord - # to set up the DB connection capable of interacting with migration. - previouslyConnected = ActiveRecord::Base.connected? - unless previouslyConnected - ApplicationRecord.remove_connection - ActiveRecord::Base.establish_connection(config) - end + ActiveRecord::Migration.verbose = verbose ActiveRecord::Base.connection_pool.with_connection do begin - context = default_migration_context - if needs_migration?(context) - ran = context.migrate + with_migration_context do |context| + if context.needs_migration? + ran = context.migrate + end end - # ActiveRecord::Migrator#migrate rescues all errors and re-raises them - # as StandardError + # ActiveRecord::Migrator#migrate rescues all errors and re-raises them as StandardError rescue StandardError => error self.error = error elog('DB.migrate threw an exception', error: error) end end - unless previouslyConnected - ActiveRecord::Base.remove_connection - ApplicationRecord.establish_connection(config) - end # Since the connections that existed before the migrations ran could # have outdated column information, reset column information for all # ApplicationRecord descendents to prevent missing method errors for @@ -57,15 +46,14 @@ def migrate(config=nil, verbose=false) # information was cached. reset_column_information - return ran + ran end # Determine if the currently established database connection needs migration # - # @param [ActiveRecord::MigrationContext,snil] context The migration context to check. Will default if not supplied # @return [Boolean] True if migration is required, false otherwise - def needs_migration?(context = default_migration_context) - ActiveRecord::Base.connection_pool.with_connection do + def needs_migration? + with_migration_context do |context| return context.needs_migration? end end @@ -77,6 +65,10 @@ def needs_migration?(context = default_migration_context) private + def with_migration_context + yield ActiveRecord::MigrationContext.new(gather_engine_migration_paths) + end + # @return [ActiveRecord::MigrationContext] def default_migration_context ActiveRecord::MigrationContext.new(gather_engine_migration_paths, ActiveRecord::SchemaMigration)