refactor migration process for Rails 5

As noted in https://github.com/rails/rails/issues/36544 using
ActiveRecord migrations internally is not truly supported. This
workaround is valid for Rails 5 and might be easily adjusted
in Rails 6 although that is still TBD.

lib/msf/core/db_manager/connection.rb

@@ -10,12 +10,12 @@ module Msf::DBManager::Connection
   # and setting {#workspace}.
   #
   # @return [void]
-  def after_establish_connection
+  def after_establish_connection(opts={})
     self.migrated = false
 
     begin
       # Migrate the database, if needed
-      migrate
+      migrate(opts)
     rescue ::Exception => exception
       self.error = exception
       elog('DB.connect threw an exception', error: exception)
@@ -61,7 +61,7 @@ module Msf::DBManager::Connection
       elog('DB.connect threw an exception', error: e)
       return false
     ensure
-      after_establish_connection
+      after_establish_connection(nopts)
     end
 
     true

lib/msf/core/db_manager/migration.rb

@@ -28,18 +28,30 @@ module Msf::DBManager::Migration
 
   # Migrate database to latest schema version.
   #
+  # @param config [Hash] see ActiveRecord::Base.establish_connection
   # @param verbose [Boolean] see ActiveRecord::Migration.verbose
   # @return [Array<ActiveRecord::MigrationProxy] List of migrations that
   #   ran.
   #
-  # @see ActiveRecord::Migrator.migrate
-  def migrate(verbose=false)
+  # @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
-
-    ApplicationRecord.connection_pool.with_connection do
+    ActiveRecord::Base.connection_pool.with_connection do
       begin
-        ran = ActiveRecord::Migration.migrate(:up)
+        # When framework reached Rails 6 the path set here may be better suited a simple Array[]
+        context = ActiveRecord::MigrationContext.new(ActiveRecord::Migrator.migrations_paths)
+        if context.needs_migration?
+          ran = context.migrate
+        end
           # ActiveRecord::Migrator#migrate rescues all errors and re-raises them
           # as StandardError
       rescue StandardError => error
@@ -48,6 +60,10 @@ module Msf::DBManager::Migration
       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

spec/support/shared/examples/msf/db_manager/migration.rb

@@ -30,26 +30,36 @@ RSpec.shared_examples_for 'Msf::DBManager::Migration' do
       db_manager.migrate
     end
 
-    it 'should call ActiveRecord::Migration.migrate' do
-      expect(ActiveRecord::Migration).to receive(:migrate).with(:up)
+    it 'should create an ActiveRecord::MigrationContext' do
+      expect(ActiveRecord::MigrationContext).to receive(:new)
 
       migrate
     end
 
-    it 'should return migrations that were ran from ActiveRecord::Migrator.migrate' do
-      migrations = [double('Migration 1')]
-      expect(ActiveRecord::Migration).to receive(:migrate).and_return(migrations)
 
-      expect(migrate).to eq migrations
+    it 'should return an ActiveRecord::MigrationContext with known migrations' do
+      migrations_paths = [File.expand_path("../../../../../file_fixtures/migrate", __dir__)]
+      expect(ActiveRecord::Migrator).to receive(:migrations_paths).and_return(migrations_paths).exactly(3).times
+      result = migrate
+      expect(result.size).to eq 1
+      expect(result[0].name).to eq "TestDbMigration"
     end
 
+
+    # it 'should return migrations that were ran from ActiveRecord::Migrator.migrate' do
+    #   migrations = [double('Migration 1')]
+    #   expect(ActiveRecord::Migration).to receive(:migrate).and_return(migrations)
+    #
+    #   expect(migrate).to eq migrations
+    # end
+
     it 'should reset the column information' do
       expect(db_manager).to receive(:reset_column_information)
 
       migrate
     end
 
-    context 'with StandardError from ActiveRecord::Migration.migrate' do
+    context 'with StandardError from ActiveRecord::MigrationContext.migrate' do
       let(:standard_error) do
         StandardError.new(message)
       end
@@ -59,7 +69,10 @@ RSpec.shared_examples_for 'Msf::DBManager::Migration' do
       end
 
       before(:example) do
-        expect(ActiveRecord::Migration).to receive(:migrate).and_raise(standard_error)
+        mockContext = ActiveRecord::MigrationContext.new(nil)
+        expect(ActiveRecord::MigrationContext).to receive(:new).and_return(mockContext)
+        expect(mockContext).to receive(:needs_migration?).and_return(true)
+        expect(mockContext).to receive(:migrate).and_raise(standard_error)
       end
 
       it 'should set Msf::DBManager#error' do
@@ -80,7 +93,7 @@ RSpec.shared_examples_for 'Msf::DBManager::Migration' do
 
     context 'with verbose' do
       def migrate
-        db_manager.migrate(verbose)
+        db_manager.migrate(nil, verbose)
       end
 
       context 'false' do