Instrument database migrations

Relates to
https://gitlab.com/gitlab-org/database-team/team-tasks/-/issues/137

lib/gitlab/database/migrations/instrumentation.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Migrations
+      Observation = Struct.new(
+        :migration,
+        :walltime,
+        :success
+      )
+
+      class Instrumentation
+        attr_reader :observations
+
+        def initialize
+          @observations = []
+        end
+
+        def observe(migration, &block)
+          observation = Observation.new(migration)
+          observation.success = true
+
+          exception = nil
+
+          observation.walltime = Benchmark.realtime do
+            yield
+          rescue => e
+            exception = e
+            observation.success = false
+          end
+
+          record_observation(observation)
+
+          raise exception if exception
+
+          observation
+        end
+
+        private
+
+        def record_observation(observation)
+          @observations << observation
+        end
+      end
+    end
+  end
+end

lib/tasks/gitlab/db.rake

@@ -231,5 +231,37 @@ namespace :gitlab do
       puts "Found user created projects. Database active"
       exit 0
     end
+
+    desc 'Run migrations with instrumentation'
+    task :migration_testing, [:result_file] => :environment do |_, args|
+      result_file = args[:result_file] || raise("Please specify result_file argument")
+      raise "File exists already, won't overwrite: #{result_file}" if File.exist?(result_file)
+
+      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)
+      end
+
+      instrumentation = Gitlab::Database::Migrations::Instrumentation.new
+
+      pending_migrations.each do |migration|
+        instrumentation.observe(migration.version) do
+          ActiveRecord::Migrator.new(:up, ctx.migrations, ctx.schema_migration, migration.version).run
+        end
+      end
+    ensure
+      if instrumentation
+        File.open(result_file, 'wb+') do |io|
+          io << instrumentation.observations.to_json
+        end
+      end
+
+      ActiveRecord::Base.clear_cache!
+      ActiveRecord::Migration.verbose = verbose_was
+    end
   end
 end

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

+# frozen_string_literal: true
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Migrations::Instrumentation do
+  describe '#observe' do
+    subject { described_class.new }
+
+    let(:migration) { 1234 }
+
+    it 'executes the given block' do
+      expect { |b| subject.observe(migration, &b) }.to yield_control
+    end
+
+    context 'on successful execution' do
+      subject { described_class.new.observe(migration) {} }
+
+      it 'records walltime' do
+        expect(subject.walltime).not_to be_nil
+      end
+
+      it 'records success' do
+        expect(subject.success).to be_truthy
+      end
+
+      it 'records the migration version' do
+        expect(subject.migration).to eq(migration)
+      end
+    end
+
+    context 'upon failure' do
+      subject { described_class.new.observe(migration) { raise 'something went wrong' } }
+
+      it 'raises the exception' do
+        expect { subject }.to raise_error(/something went wrong/)
+      end
+
+      context 'retrieving observations' do
+        subject { instance.observations.first }
+
+        before do
+          instance.observe(migration) { raise 'something went wrong' }
+        rescue
+          # ignore
+        end
+
+        let(:instance) { described_class.new }
+
+        it 'records walltime' do
+          expect(subject.walltime).not_to be_nil
+        end
+
+        it 'records failure' do
+          expect(subject.success).to be_falsey
+        end
+
+        it 'records the migration version' do
+          expect(subject.migration).to eq(migration)
+        end
+      end
+    end
+
+    context 'sequence of migrations with failures' do
+      subject { described_class.new }
+
+      let(:migration1) { double('migration1', call: nil) }
+      let(:migration2) { double('migration2', call: nil) }
+
+      it 'records observations for all migrations' do
+        subject.observe('migration1') {}
+        subject.observe('migration2') { raise 'something went wrong' } rescue nil
+
+        expect(subject.observations.size).to eq(2)
+      end
+    end
+  end
+end

spec/tasks/gitlab/db_rake_spec.rb

@@ -297,6 +297,57 @@ RSpec.describe 'gitlab:db namespace rake task' do
     end
   end
 
+  describe '#migrate_with_instrumentation' do
+    subject { run_rake_task('gitlab:db:migration_testing', "[#{filename}]") }
+
+    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), pending_migration] }
+    let(:pending_migration) { double('migration2', version: 2) }
+    let(:filename) { 'results-file.json'}
+    let(:buffer) { StringIO.new }
+    let(:observations) { %w[some data] }
+
+    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
+
+      allow(File).to receive(:open).with(filename, 'wb+').and_yield(buffer)
+    end
+
+    it 'fails when given no filename argument' do
+      expect { run_rake_task('gitlab:db:migration_testing') }.to raise_error(/specify result_file/)
+    end
+
+    it 'fails when the given file already exists' do
+      expect(File).to receive(:exist?).with(filename).and_return(true)
+
+      expect { subject }.to raise_error(/File exists/)
+    end
+
+    it 'instruments the pending migration' do
+      expect(instrumentation).to receive(:observe).with(2).and_yield
+
+      subject
+    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)
+
+      subject
+    end
+
+    it 'writes observations out to JSON file' do
+      subject
+
+      expect(buffer.string).to eq(observations.to_json)
+    end
+  end
+
   def run_rake_task(task_name, arguments = '')
     Rake::Task[task_name].reenable
     Rake.application.invoke_task("#{task_name}#{arguments}")