Merge pull request rapid7#19645 from adeherdt-r7/MS-9862-rails-upgrade-preparation-migration-manager

MS-9862 Ruby on Rails Upgrade Preparation : Migration

Author: adeherdt-r7

lib/msf/core/db_manager/migration.rb

@@ -23,49 +23,37 @@ 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
     # column methods for columns created in migrations after the column
     # 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)