Migration testing for down database migrations

Adds a gitlab:db:migration_testing:down rake task to instrument down
migrations on the current branch.

Renames the gitlab:db:migration_testing rake task to
gitlab:db:migration_testing:up for symmetry

Keeps the gitlab:db:migration_testing task for now, so that this change
is backwards compatible with the migration testing pipeline, which does
not yet know about these changes.

Relates to
https://gitlab.com/gitlab-org/database-team/gitlab-com-database-testing/-/issues/18

lib/gitlab/database/migrations/instrumentation.rb

@@ -4,21 +4,21 @@ module Gitlab
   module Database
     module Migrations
       class Instrumentation
-        RESULT_DIR = Rails.root.join('tmp', 'migration-testing').freeze
         STATS_FILENAME = 'migration-stats.json'
 
         attr_reader :observations
 
-        def initialize(observer_classes = ::Gitlab::Database::Migrations::Observers.all_observers)
+        def initialize(result_dir:, observer_classes: ::Gitlab::Database::Migrations::Observers.all_observers)
           @observer_classes = observer_classes
           @observations = []
+          @result_dir = result_dir
         end
 
         def observe(version:, name:, &block)
           observation = Observation.new(version, name)
           observation.success = true
 
-          observers = observer_classes.map { |c| c.new(observation) }
+          observers = observer_classes.map { |c| c.new(observation, @result_dir) }
 
           exception = nil
 

lib/gitlab/database/migrations/observers/migration_observer.rb

@@ -5,11 +5,12 @@ module Gitlab
     module Migrations
       module Observers
         class MigrationObserver
-          attr_reader :connection, :observation
+          attr_reader :connection, :observation, :output_dir
 
-          def initialize(observation)
+          def initialize(observation, output_dir)
             @connection = ActiveRecord::Base.connection
             @observation = observation
+            @output_dir = output_dir
           end
 
           def before

lib/gitlab/database/migrations/observers/query_details.rb

@@ -6,7 +6,7 @@ module Gitlab
       module Observers
         class QueryDetails < MigrationObserver
           def before
-            file_path = File.join(Instrumentation::RESULT_DIR, "#{observation.version}_#{observation.name}-query-details.json")
+            file_path = File.join(output_dir, "#{observation.version}_#{observation.name}-query-details.json")
             @file = File.open(file_path, 'wb')
             @writer = Oj::StreamWriter.new(@file, {})
             @writer.push_array

lib/gitlab/database/migrations/observers/query_log.rb

@@ -7,7 +7,7 @@ module Gitlab
         class QueryLog < MigrationObserver
           def before
             @logger_was = ActiveRecord::Base.logger
-            file_path = File.join(Instrumentation::RESULT_DIR, "#{observation.version}_#{observation.name}.log")
+            file_path = File.join(output_dir, "#{observation.version}_#{observation.name}.log")
             @logger = Logger.new(file_path)
             ActiveRecord::Base.logger = @logger
           end

lib/gitlab/database/migrations/runner.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Migrations
+      class Runner
+        BASE_RESULT_DIR = Rails.root.join('tmp', 'migration-testing').freeze
+
+        class << self
+          def up(legacy_pipeline: false)
+            result_dir = if legacy_pipeline
+                           BASE_RESULT_DIR
+                         else
+                           BASE_RESULT_DIR.join('up')
+                         end
+
+            Runner.new(direction: :up, migrations: migrations_for_up, result_dir: result_dir)
+          end
+
+          def down
+            Runner.new(direction: :down, migrations: migrations_for_down, result_dir: BASE_RESULT_DIR.join('down'))
+          end
+
+          def migration_context
+            @migration_context ||= ApplicationRecord.connection.migration_context
+          end
+
+          private
+
+          def migrations_for_up
+            existing_versions = migration_context.get_all_versions.to_set
+
+            migration_context.migrations.reject do |migration|
+              existing_versions.include?(migration.version)
+            end
+          end
+
+          def migration_file_names_this_branch
+            `git diff --name-only origin/HEAD...HEAD db/post_migrate db/migrate`.split("\n")
+          end
+
+          def migrations_for_down
+            versions_this_branch = migration_file_names_this_branch.map do |m_name|
+              m_name.match(%r{^db/(post_)?migrate/(\d+)}) { |m| m.captures[1]&.to_i }
+            end.to_set
+
+            existing_versions = migration_context.get_all_versions.to_set
+            migration_context.migrations.select do |migration|
+              existing_versions.include?(migration.version) && versions_this_branch.include?(migration.version)
+            end
+          end
+        end
+
+        attr_reader :direction, :result_dir, :migrations
+
+        delegate :migration_context, to: :class
+
+        def initialize(direction:, migrations:, result_dir:)
+          raise "Direction must be up or down" unless %i[up down].include?(direction)
+
+          @direction = direction
+          @migrations = migrations
+          @result_dir = result_dir
+        end
+
+        def run
+          FileUtils.mkdir_p(result_dir)
+
+          verbose_was = ActiveRecord::Migration.verbose
+          ActiveRecord::Migration.verbose = true
+
+          sorted_migrations = migrations.sort_by(&:version)
+          sorted_migrations.reverse! if direction == :down
+
+          instrumentation = Instrumentation.new(result_dir: result_dir)
+
+          sorted_migrations.each do |migration|
+            instrumentation.observe(version: migration.version, name: migration.name) do
+              ActiveRecord::Migrator.new(direction, migration_context.migrations, migration_context.schema_migration, migration.version).run
+            end
+          end
+        ensure
+          if instrumentation
+            File.open(File.join(result_dir, Gitlab::Database::Migrations::Instrumentation::STATS_FILENAME), 'wb+') do |io|
+              io << instrumentation.observations.to_json
+            end
+          end
+
+          # We clear the cache here to mirror the cache clearing that happens at the end of `db:migrate` tasks
+          # This clearing makes subsequent rake tasks in the same execution pick up database schema changes caused by
+          # the migrations that were just executed
+          ApplicationRecord.clear_cache!
+          ActiveRecord::Migration.verbose = verbose_was
+        end
+      end
+    end
+  end
+end

lib/tasks/gitlab/db.rake

@@ -211,37 +211,22 @@ namespace :gitlab do
       exit 0
     end
 
-    desc 'Run migrations with instrumentation'
-    task migration_testing: :environment do
-      result_dir = Gitlab::Database::Migrations::Instrumentation::RESULT_DIR
-      FileUtils.mkdir_p(result_dir)
-
-      verbose_was = ActiveRecord::Migration.verbose
-      ActiveRecord::Migration.verbose = true
-
-      ctx = ActiveRecord::Base.connection.migration_context
-      existing_versions = ctx.get_all_versions.to_set
-
-      pending_migrations = ctx.migrations.reject do |migration|
-        existing_versions.include?(migration.version)
+    namespace :migration_testing do
+      desc 'Run migrations with instrumentation'
+      task up: :environment do
+        Gitlab::Database::Migrations::Runner.up.run
       end
 
-      instrumentation = Gitlab::Database::Migrations::Instrumentation.new
-
-      pending_migrations.each do |migration|
-        instrumentation.observe(version: migration.version, name: migration.name) do
-          ActiveRecord::Migrator.new(:up, ctx.migrations, ctx.schema_migration, migration.version).run
-        end
-      end
-    ensure
-      if instrumentation
-        File.open(File.join(result_dir, Gitlab::Database::Migrations::Instrumentation::STATS_FILENAME), 'wb+') do |io|
-          io << instrumentation.observations.to_json
-        end
+      desc 'Run down migrations in current branch with instrumentation'
+      task down: :environment do
+        Gitlab::Database::Migrations::Runner.down.run
       end
+    end
 
-      ActiveRecord::Base.clear_cache!
-      ActiveRecord::Migration.verbose = verbose_was
+    # TODO: Remove this rake task after migrating the database testing runner to :up / :down versions of it
+    desc 'Run migrations with instrumentation'
+    task migration_testing: :environment do
+      Gitlab::Database::Migrations::Runner.up(legacy_pipeline: true).run
     end
 
     desc 'Run all pending batched migrations'

spec/lib/gitlab/database/migrations/instrumentation_spec.rb

@@ -2,8 +2,13 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::Instrumentation do
+  let(:result_dir) { Dir.mktmpdir }
+
+  after do
+    FileUtils.rm_rf(result_dir)
+  end
   describe '#observe' do
-    subject { described_class.new }
+    subject { described_class.new(result_dir: result_dir) }
 
     let(:migration_name) { 'test' }
     let(:migration_version) { '12345' }
@@ -13,7 +18,7 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
     end
 
     context 'behavior with observers' do
-      subject { described_class.new([Gitlab::Database::Migrations::Observers::MigrationObserver]).observe(version: migration_version, name: migration_name) {} }
+      subject { described_class.new(observer_classes: [Gitlab::Database::Migrations::Observers::MigrationObserver], result_dir: result_dir).observe(version: migration_version, name: migration_name) {} }
 
       let(:observer) { instance_double('Gitlab::Database::Migrations::Observers::MigrationObserver', before: nil, after: nil, record: nil) }
 
@@ -24,7 +29,7 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
       it 'instantiates observer with observation' do
         expect(Gitlab::Database::Migrations::Observers::MigrationObserver)
           .to receive(:new)
-          .with(instance_of(Gitlab::Database::Migrations::Observation)) { |observation| expect(observation.version).to eq(migration_version) }
+          .with(instance_of(Gitlab::Database::Migrations::Observation), anything) { |observation| expect(observation.version).to eq(migration_version) }
           .and_return(observer)
 
         subject
@@ -58,7 +63,7 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
     end
 
     context 'on successful execution' do
-      subject { described_class.new.observe(version: migration_version, name: migration_name) {} }
+      subject { described_class.new(result_dir: result_dir).observe(version: migration_version, name: migration_name) {} }
 
       it 'records walltime' do
         expect(subject.walltime).not_to be_nil
@@ -78,7 +83,7 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
     end
 
     context 'upon failure' do
-      subject { described_class.new.observe(version: migration_version, name: migration_name) { raise 'something went wrong' } }
+      subject { described_class.new(result_dir: result_dir).observe(version: migration_version, name: migration_name) { raise 'something went wrong' } }
 
       it 'raises the exception' do
         expect { subject }.to raise_error(/something went wrong/)
@@ -93,7 +98,7 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
           # ignore
         end
 
-        let(:instance) { described_class.new }
+        let(:instance) { described_class.new(result_dir: result_dir) }
 
         it 'records walltime' do
           expect(subject.walltime).not_to be_nil
@@ -114,7 +119,7 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
     end
 
     context 'sequence of migrations with failures' do
-      subject { described_class.new }
+      subject { described_class.new(result_dir: result_dir) }
 
       let(:migration1) { double('migration1', call: nil) }
       let(:migration2) { double('migration2', call: nil) }

spec/lib/gitlab/database/migrations/observers/query_details_spec.rb

@@ -2,7 +2,7 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::Observers::QueryDetails do
-  subject { described_class.new(observation) }
+  subject { described_class.new(observation, directory_path) }
 
   let(:observation) { Gitlab::Database::Migrations::Observation.new(migration_version, migration_name) }
   let(:connection) { ActiveRecord::Base.connection }
@@ -14,10 +14,6 @@ RSpec.describe Gitlab::Database::Migrations::Observers::QueryDetails do
   let(:migration_version) { 20210422152437 }
   let(:migration_name) { 'test' }
 
-  before do
-    stub_const('Gitlab::Database::Migrations::Instrumentation::RESULT_DIR', directory_path)
-  end
-
   after do
     FileUtils.remove_entry(directory_path)
   end

spec/lib/gitlab/database/migrations/observers/query_log_spec.rb

@@ -2,7 +2,7 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::Observers::QueryLog do
-  subject { described_class.new(observation) }
+  subject { described_class.new(observation, directory_path) }
 
   let(:observation) { Gitlab::Database::Migrations::Observation.new(migration_version, migration_name) }
   let(:connection) { ActiveRecord::Base.connection }
@@ -11,10 +11,6 @@ RSpec.describe Gitlab::Database::Migrations::Observers::QueryLog do
   let(:migration_version) { 20210422152437 }
   let(:migration_name) { 'test' }
 
-  before do
-    stub_const('Gitlab::Database::Migrations::Instrumentation::RESULT_DIR', directory_path)
-  end
-
   after do
     FileUtils.remove_entry(directory_path)
   end

spec/lib/gitlab/database/migrations/observers/query_statistics_spec.rb

@@ -2,7 +2,7 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::Observers::QueryStatistics do
-  subject { described_class.new(observation) }
+  subject { described_class.new(observation, double("unused path")) }
 
   let(:observation) { Gitlab::Database::Migrations::Observation.new }
   let(:connection) { ActiveRecord::Base.connection }

spec/lib/gitlab/database/migrations/observers/total_database_size_change_spec.rb

@@ -2,7 +2,7 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::Observers::TotalDatabaseSizeChange do
-  subject { described_class.new(observation) }
+  subject { described_class.new(observation, double('unused path')) }
 
   let(:observation) { Gitlab::Database::Migrations::Observation.new }
   let(:connection) { ActiveRecord::Base.connection }

spec/lib/gitlab/database/migrations/runner_spec.rb

+# frozen_string_literal: true
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Migrations::Runner do
+  let(:result_dir) { Pathname.new(Dir.mktmpdir) }
+
+  let(:migration_runs) { [] } # This list gets populated as the runner tries to run migrations
+
+  # Tests depend on all of these lists being sorted in the order migrations would be applied
+  let(:applied_migrations_other_branches) { [double(ActiveRecord::Migration, version: 1, name: 'migration_complete_other_branch')] }
+
+  let(:applied_migrations_this_branch) do
+    [
+      double(ActiveRecord::Migration, version: 2, name: 'older_migration_complete_this_branch'),
+      double(ActiveRecord::Migration, version: 3, name: 'newer_migration_complete_this_branch')
+    ].sort_by(&:version)
+  end
+
+  let(:pending_migrations) do
+    [
+      double(ActiveRecord::Migration, version: 4, name: 'older_migration_pending'),
+      double(ActiveRecord::Migration, version: 5, name: 'newer_migration_pending')
+    ].sort_by(&:version)
+  end
+
+  before do
+    stub_const('Gitlab::Database::Migrations::Runner::BASE_RESULT_DIR', result_dir)
+    allow(ActiveRecord::Migrator).to receive(:new) do |dir, _all_migrations, _schema_migration_class, version_to_migrate|
+      migrator = double(ActiveRecord::Migrator)
+      expect(migrator).to receive(:run) do
+        migration_runs << OpenStruct.new(dir: dir, version_to_migrate: version_to_migrate)
+      end
+      migrator
+    end
+
+    all_versions = (applied_migrations_other_branches + applied_migrations_this_branch).map(&:version)
+    migrations = applied_migrations_other_branches + applied_migrations_this_branch + pending_migrations
+    ctx = double(ActiveRecord::MigrationContext, get_all_versions: all_versions, migrations: migrations, schema_migration: ActiveRecord::SchemaMigration)
+
+    allow(described_class).to receive(:migration_context).and_return(ctx)
+
+    names_this_branch = (applied_migrations_this_branch + pending_migrations).map { |m| "db/migrate/#{m.version}_#{m.name}.rb"}
+    allow(described_class).to receive(:migration_file_names_this_branch).and_return(names_this_branch)
+  end
+
+  after do
+    FileUtils.rm_rf(result_dir)
+  end
+
+  it 'creates the results dir when one does not exist' do
+    FileUtils.rm_rf(result_dir)
+
+    expect do
+      described_class.new(direction: :up, migrations: [], result_dir: result_dir).run
+    end.to change { Dir.exist?(result_dir) }.from(false).to(true)
+  end
+
+  describe '.up' do
+    context 'result directory' do
+      context 'legacy mode' do
+        it 'uses the root result directory' do
+          expect(described_class.up(legacy_pipeline: true).result_dir).to eq(result_dir)
+        end
+      end
+
+      context 'not legacy mode' do
+        it 'uses the /up subdirectory' do
+          expect(described_class.up.result_dir).to eq(result_dir.join('up'))
+        end
+      end
+    end
+
+    context 'migrations to run' do
+      subject(:up) { described_class.up }
+
+      it 'is the list of pending migrations' do
+        expect(up.migrations).to eq(pending_migrations)
+      end
+    end
+
+    context 'running migrations' do
+      subject(:up) { described_class.up }
+
+      it 'runs the unapplied migrations in version order', :aggregate_failures do
+        up.run
+
+        expect(migration_runs.map(&:dir)).to eq([:up, :up])
+        expect(migration_runs.map(&:version_to_migrate)).to eq(pending_migrations.map(&:version))
+      end
+    end
+  end
+
+  describe '.down' do
+    subject(:down) { described_class.down }
+
+    context 'result directory' do
+      it 'is the /down subdirectory' do
+        expect(down.result_dir).to eq(result_dir.join('down'))
+      end
+    end
+
+    context 'migrations to run' do
+      it 'is the list of migrations that are up and on this branch' do
+        expect(down.migrations).to eq(applied_migrations_this_branch)
+      end
+    end
+
+    context 'running migrations' do
+      it 'runs the applied migrations for the current branch in reverse order', :aggregate_failures do
+        down.run
+
+        expect(migration_runs.map(&:dir)).to eq([:down, :down])
+        expect(migration_runs.map(&:version_to_migrate)).to eq(applied_migrations_this_branch.reverse.map(&:version))
+      end
+    end
+  end
+end

spec/tasks/gitlab/db_rake_spec.rb

@@ -293,53 +293,37 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
   end
 
   describe '#migrate_with_instrumentation' do
-    subject { run_rake_task('gitlab:db:migration_testing') }
+    describe '#up' do
+      subject { run_rake_task('gitlab:db:migration_testing:up') }
 
-    let(:ctx) { double('ctx', migrations: all_migrations, schema_migration: double, get_all_versions: existing_versions) }
-    let(:instrumentation) { instance_double(Gitlab::Database::Migrations::Instrumentation, observations: observations) }
-    let(:existing_versions) { [1] }
-    let(:all_migrations) { [double('migration1', version: 1, name: 'test'), pending_migration] }
-    let(:pending_migration) { double('migration2', version: 2, name: 'test') }
-    let(:filename) { Gitlab::Database::Migrations::Instrumentation::STATS_FILENAME }
-    let(:result_dir) { Dir.mktmpdir }
-    let(:observations) { %w[some data] }
+      it 'delegates to the migration runner' do
+        expect(::Gitlab::Database::Migrations::Runner).to receive_message_chain(:up, :run)
 
-    before do
-      allow(ActiveRecord::Base.connection).to receive(:migration_context).and_return(ctx)
-      allow(Gitlab::Database::Migrations::Instrumentation).to receive(:new).and_return(instrumentation)
-      allow(ActiveRecord::Migrator).to receive_message_chain('new.run').with(any_args).with(no_args)
-
-      allow(instrumentation).to receive(:observe).and_yield
-
-      stub_const('Gitlab::Database::Migrations::Instrumentation::RESULT_DIR', result_dir)
-    end
-
-    after do
-      FileUtils.rm_rf(result_dir)
+        subject
+      end
     end
 
-    it 'creates result directory when one does not exist' do
-      FileUtils.rm_rf(result_dir)
+    describe '#down' do
+      subject { run_rake_task('gitlab:db:migration_testing:down') }
 
-      expect { subject }.to change { Dir.exist?(result_dir) }.from(false).to(true)
-    end
+      it 'delegates to the migration runner' do
+        expect(::Gitlab::Database::Migrations::Runner).to receive_message_chain(:down, :run)
 
-    it 'instruments the pending migration' do
-      expect(instrumentation).to receive(:observe).with(version: 2, name: 'test').and_yield
-
-      subject
+        subject
+      end
     end
 
-    it 'executes the pending migration' do
-      expect(ActiveRecord::Migrator).to receive_message_chain('new.run').with(:up, ctx.migrations, ctx.schema_migration, pending_migration.version).with(no_args)
+    describe 'legacy rake task' do
+      subject { run_rake_task('gitlab:db:migration_testing') }
 
-      subject
-    end
+      let(:runner) { double(Gitlab::Database::Migrations::Runner) }
 
-    it 'writes observations out to JSON file' do
-      subject
+      it 'delegates to the migration runner in legacy mode' do
+        expect(::Gitlab::Database::Migrations::Runner).to receive(:up).with(legacy_pipeline: true).and_return(runner)
+        expect(runner).to receive(:run)
 
-      expect(File.read(File.join(result_dir, filename))).to eq(observations.to_json)
+        subject
+      end
     end
   end