Add support for multiple databases for migration observers

- MigrationObserver starts to accept a connection as an argument
- Migrations::Runner passes a connection when observes a migration

Changelog: added
Related to https://gitlab.com/gitlab-org/gitlab/-/merge_requests/76681

lib/gitlab/database/migrations/instrumentation.rb

@@ -14,11 +14,11 @@ module Gitlab
           @result_dir = result_dir
         end
 
-        def observe(version:, name:, &block)
+        def observe(version:, name:, connection:, &block)
           observation = Observation.new(version, name)
           observation.success = true
 
-          observers = observer_classes.map { |c| c.new(observation, @result_dir) }
+          observers = observer_classes.map { |c| c.new(observation, @result_dir, connection) }
 
           exception = nil
 

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

@@ -7,8 +7,8 @@ module Gitlab
         class MigrationObserver
           attr_reader :connection, :observation, :output_dir
 
-          def initialize(observation, output_dir)
-            @connection = ActiveRecord::Base.connection
+          def initialize(observation, output_dir, connection)
+            @connection = connection
             @observation = observation
             @output_dir = output_dir
           end

lib/gitlab/database/migrations/runner.rb

@@ -69,7 +69,7 @@ module Gitlab
           instrumentation = Instrumentation.new(result_dir: result_dir)
 
           sorted_migrations.each do |migration|
-            instrumentation.observe(version: migration.version, name: migration.name) do
+            instrumentation.observe(version: migration.version, name: migration.name, connection: ActiveRecord::Migration.connection) do
               ActiveRecord::Migrator.new(direction, migration_context.migrations, migration_context.schema_migration, migration.version).run
             end
           end

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

@@ -3,6 +3,7 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::Instrumentation do
   let(:result_dir) { Dir.mktmpdir }
+  let(:connection) { ActiveRecord::Migration.connection }
 
   after do
     FileUtils.rm_rf(result_dir)
@@ -14,11 +15,11 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
     let(:migration_version) { '12345' }
 
     it 'executes the given block' do
-      expect { |b| subject.observe(version: migration_version, name: migration_name, &b) }.to yield_control
+      expect { |b| subject.observe(version: migration_version, name: migration_name, connection: connection, &b) }.to yield_control
     end
 
     context 'behavior with observers' do
-      subject { described_class.new(observer_classes: [Gitlab::Database::Migrations::Observers::MigrationObserver], result_dir: result_dir).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, connection: connection) {} }
 
       let(:observer) { instance_double('Gitlab::Database::Migrations::Observers::MigrationObserver', before: nil, after: nil, record: nil) }
 
@@ -29,7 +30,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), anything) { |observation| expect(observation.version).to eq(migration_version) }
+          .with(instance_of(Gitlab::Database::Migrations::Observation), anything, connection) { |observation| expect(observation.version).to eq(migration_version) }
           .and_return(observer)
 
         subject
@@ -63,7 +64,7 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
     end
 
     context 'on successful execution' do
-      subject { described_class.new(result_dir: result_dir).observe(version: migration_version, name: migration_name) {} }
+      subject { described_class.new(result_dir: result_dir).observe(version: migration_version, name: migration_name, connection: connection) {} }
 
       it 'records walltime' do
         expect(subject.walltime).not_to be_nil
@@ -83,7 +84,7 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
     end
 
     context 'upon failure' do
-      subject { described_class.new(result_dir: result_dir).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, connection: connection) { raise 'something went wrong' } }
 
       it 'raises the exception' do
         expect { subject }.to raise_error(/something went wrong/)
@@ -93,7 +94,7 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
         subject { instance.observations.first }
 
         before do
-          instance.observe(version: migration_version, name: migration_name) { raise 'something went wrong' }
+          instance.observe(version: migration_version, name: migration_name, connection: connection) { raise 'something went wrong' }
         rescue StandardError
           # ignore
         end
@@ -125,8 +126,8 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
       let(:migration2) { double('migration2', call: nil) }
 
       it 'records observations for all migrations' do
-        subject.observe(version: migration_version, name: migration_name) {}
-        subject.observe(version: migration_version, name: migration_name) { raise 'something went wrong' } rescue nil
+        subject.observe(version: migration_version, name: migration_name, connection: connection) {}
+        subject.observe(version: migration_version, name: migration_name, connection: connection) { raise 'something went wrong' } rescue nil
 
         expect(subject.observations.size).to eq(2)
       end

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

@@ -2,10 +2,10 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::Observers::QueryDetails do
-  subject { described_class.new(observation, directory_path) }
+  subject { described_class.new(observation, directory_path, connection) }
 
+  let(:connection) { ActiveRecord::Migration.connection }
   let(:observation) { Gitlab::Database::Migrations::Observation.new(migration_version, migration_name) }
-  let(:connection) { ActiveRecord::Base.connection }
   let(:query) { "select date_trunc('day', $1::timestamptz) + $2 * (interval '1 hour')" }
   let(:query_binds) { [Time.current, 3] }
   let(:directory_path) { Dir.mktmpdir }

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

@@ -2,10 +2,10 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::Observers::QueryLog do
-  subject { described_class.new(observation, directory_path) }
+  subject { described_class.new(observation, directory_path, connection) }
 
   let(:observation) { Gitlab::Database::Migrations::Observation.new(migration_version, migration_name) }
-  let(:connection) { ActiveRecord::Base.connection }
+  let(:connection) { ActiveRecord::Migration.connection }
   let(:query) { 'select 1' }
   let(:directory_path) { Dir.mktmpdir }
   let(:migration_version) { 20210422152437 }

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

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

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

@@ -2,10 +2,10 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::Observers::TotalDatabaseSizeChange do
-  subject { described_class.new(observation, double('unused path')) }
+  subject { described_class.new(observation, double('unused path'), connection) }
 
   let(:observation) { Gitlab::Database::Migrations::Observation.new }
-  let(:connection) { ActiveRecord::Base.connection }
+  let(:connection) { ActiveRecord::Migration.connection }
   let(:query) { 'select pg_database_size(current_database())' }
 
   it 'records the size change' do

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

@@ -2,8 +2,9 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::Observers::TransactionDuration do
-  subject(:transaction_duration_observer) { described_class.new(observation, directory_path) }
+  subject(:transaction_duration_observer) { described_class.new(observation, directory_path, connection) }
 
+  let(:connection) { ActiveRecord::Migration.connection }
   let(:observation) { Gitlab::Database::Migrations::Observation.new(migration_version, migration_name) }
   let(:directory_path) { Dir.mktmpdir }
   let(:log_file) { "#{directory_path}/#{migration_version}_#{migration_name}-transaction-duration.json" }
@@ -78,17 +79,17 @@ RSpec.describe Gitlab::Database::Migrations::Observers::TransactionDuration do
   end
 
   def run_real_transactions
-    ActiveRecord::Base.transaction do
+    ApplicationRecord.transaction do
     end
   end
 
   def run_sub_transactions
-    ActiveRecord::Base.transaction(requires_new: true) do
+    ApplicationRecord.transaction(requires_new: true) do
     end
   end
 
   def run_transaction
-    ActiveRecord::Base.connection_pool.with_connection do |connection|
+    ApplicationRecord.connection_pool.with_connection do |connection|
       Gitlab::Database::SharedModel.using_connection(connection) do
         Gitlab::Database::SharedModel.transaction do
           Gitlab::Database::SharedModel.transaction(requires_new: true) do

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

@@ -76,7 +76,7 @@ RSpec.describe Gitlab::Database::Migrations::Runner do
       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(&:dir)).to match_array([:up, :up])
         expect(migration_runs.map(&:version_to_migrate)).to eq(pending_migrations.map(&:version))
       end
     end
@@ -101,7 +101,7 @@ RSpec.describe Gitlab::Database::Migrations::Runner 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(&:dir)).to match_array([:down, :down])
         expect(migration_runs.map(&:version_to_migrate)).to eq(applied_migrations_this_branch.reverse.map(&:version))
       end
     end