Add latest changes from gitlab-org/gitlab@master

2 files changed, 96 insertions, 25 deletions

diff --git a/lib/gitlab/database/migrations/runner.rb b/lib/gitlab/database/migrations/runner.rb
index 4404b5bf961..85dc6051c7c 100644
--- a/lib/gitlab/database/migrations/runner.rb
+++ b/lib/gitlab/database/migrations/runner.rb
@@ -6,40 +6,68 @@ module Gitlab
       class Runner
         BASE_RESULT_DIR = Rails.root.join('tmp', 'migration-testing').freeze
         METADATA_FILENAME = 'metadata.json'
-        SCHEMA_VERSION = 3 # Version of the output format produced by the runner
+        SCHEMA_VERSION = 4 # Version of the output format produced by the runner
 
         class << self
-          def up
-            Runner.new(direction: :up, migrations: migrations_for_up, result_dir: BASE_RESULT_DIR.join('up'))
+          def up(database:, legacy_mode: false)
+            within_context_for_database(database) do
+              Runner.new(direction: :up, database: database, migrations: migrations_for_up(database), legacy_mode: legacy_mode)
+            end
           end
 
-          def down
-            Runner.new(direction: :down, migrations: migrations_for_down, result_dir: BASE_RESULT_DIR.join('down'))
+          def down(database:, legacy_mode: false)
+            within_context_for_database(database) do
+              Runner.new(direction: :down, database: database, migrations: migrations_for_down(database), legacy_mode: legacy_mode)
+            end
           end
 
           def background_migrations
             TestBackgroundRunner.new(result_dir: BASE_RESULT_DIR.join('background_migrations'))
           end
 
-          def batched_background_migrations(for_database:)
+          def batched_background_migrations(for_database:, legacy_mode: false)
             runner = nil
 
+            result_dir = if legacy_mode
+                           BASE_RESULT_DIR.join('background_migrations')
+                         else
+                           BASE_RESULT_DIR.join(for_database.to_s, 'background_migrations')
+                         end
+
             # Only one loop iteration since we pass `only:` here
             Gitlab::Database::EachDatabase.each_database_connection(only: for_database) do |connection|
               runner = Gitlab::Database::Migrations::TestBatchedBackgroundRunner
-                         .new(result_dir: BASE_RESULT_DIR.join('background_migrations'), connection: connection)
+                         .new(result_dir: result_dir, connection: connection)
             end
 
             runner
           end
 
           def migration_context
-            @migration_context ||= ApplicationRecord.connection.migration_context
+            # We're mirroring rails internal migration code, which requires that
+            # ActiveRecord::Base has connected to the current database. The correct database is chosen by
+            # within_context_for_database
+            ActiveRecord::Base.connection.migration_context # rubocop:disable Database/MultipleDatabases
+          end
+
+          # rubocop:disable Database/MultipleDatabases
+          def within_context_for_database(database)
+            original_db_config = ActiveRecord::Base.connection_db_config
+            # The config only works if passed a string
+            db_config = ActiveRecord::Base.configurations.configs_for(env_name: Rails.env, name: database.to_s)
+            raise ArgumentError, "Cannot find a database configuration for #{database}" unless db_config
+
+            ActiveRecord::Base.establish_connection(db_config) # rubocop:disable Database/EstablishConnection
+
+            yield
+          ensure
+            ActiveRecord::Base.establish_connection(original_db_config) # rubocop:disable Database/EstablishConnection
           end
+          # rubocop:enable Database/MultipleDatabases
 
           private
 
-          def migrations_for_up
+          def migrations_for_up(database)
             existing_versions = migration_context.get_all_versions.to_set
 
             migration_context.migrations.reject do |migration|
@@ -51,7 +79,7 @@ module Gitlab
             `git diff --name-only origin/HEAD...HEAD db/post_migrate db/migrate`.split("\n")
           end
 
-          def migrations_for_down
+          def migrations_for_down(database)
             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
@@ -65,14 +93,21 @@ module Gitlab
 
         attr_reader :direction, :result_dir, :migrations
 
-        delegate :migration_context, to: :class
+        delegate :migration_context, :within_context_for_database, to: :class
 
-        def initialize(direction:, migrations:, result_dir:)
+        def initialize(direction:, database:, migrations:, legacy_mode: false)
           raise "Direction must be up or down" unless %i[up down].include?(direction)
 
           @direction = direction
           @migrations = migrations
-          @result_dir = result_dir
+          @result_dir = if legacy_mode
+                          BASE_RESULT_DIR.join(direction.to_s)
+                        else
+                          BASE_RESULT_DIR.join(database.to_s, direction.to_s)
+                        end
+
+          @database = database
+          @legacy_mode = legacy_mode
         end
 
         def run
@@ -86,14 +121,22 @@ module Gitlab
 
           instrumentation = Instrumentation.new(result_dir: result_dir)
 
-          sorted_migrations.each do |migration|
-            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
+          within_context_for_database(@database) do
+            sorted_migrations.each do |migration|
+              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
           end
         ensure
           metadata_filename = File.join(result_dir, METADATA_FILENAME)
-          File.write(metadata_filename, { version: SCHEMA_VERSION }.to_json)
+          version = if @legacy_mode
+                      3
+                    else
+                      SCHEMA_VERSION
+                    end
+
+          File.write(metadata_filename, { database: @database.to_s, version: version }.to_json)
 
           # 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
diff --git a/lib/tasks/gitlab/db.rake b/lib/tasks/gitlab/db.rake
index 30e0e3e72ff..4ef0c396f4a 100644
--- a/lib/tasks/gitlab/db.rake
+++ b/lib/tasks/gitlab/db.rake
@@ -304,14 +304,30 @@ namespace :gitlab do
     end
 
     namespace :migration_testing do
-      desc 'Run migrations with instrumentation'
+      # Not possible to import Gitlab::Database::DATABASE_NAMES here
+      # Specs verify that a task exists for each entry in that array.
+      all_databases = %i[main ci]
+
       task up: :environment do
-        Gitlab::Database::Migrations::Runner.up.run
+        Gitlab::Database::Migrations::Runner.up(database: 'main', legacy_mode: true).run
+      end
+
+      namespace :up do
+        all_databases.each do |db|
+          desc "Run migrations on #{db} with instrumentation"
+          task db => :environment do
+            Gitlab::Database::Migrations::Runner.up(database: db).run
+          end
+        end
       end
 
-      desc 'Run down migrations in current branch with instrumentation'
-      task down: :environment do
-        Gitlab::Database::Migrations::Runner.down.run
+      namespace :down do
+        all_databases.each do |db|
+          desc "Run down migrations on #{db} in current branch with instrumentation"
+          task db => :environment do
+            Gitlab::Database::Migrations::Runner.down(database: db).run
+          end
+        end
       end
 
       desc 'Sample traditional background migrations with instrumentation'
@@ -321,12 +337,24 @@ namespace :gitlab do
         Gitlab::Database::Migrations::Runner.background_migrations.run_jobs(for_duration: duration)
       end
 
-      desc 'Sample batched background migrations with instrumentation'
+      namespace :sample_batched_background_migrations do
+        all_databases.each do |db|
+          desc "Sample batched background migrations on #{db} with instrumentation"
+          task db, [:duration_s] => [:environment] do |_t, args|
+            duration = args[:duration_s]&.to_i&.seconds || 30.minutes # Default of 30 minutes
+
+            Gitlab::Database::Migrations::Runner.batched_background_migrations(for_database: db)
+                                                .run_jobs(for_duration: duration)
+          end
+        end
+      end
+
+      desc "Sample batched background migrations with instrumentation (legacy)"
       task :sample_batched_background_migrations, [:database, :duration_s] => [:environment] do |_t, args|
-        database_name = args[:database] || 'main'
         duration = args[:duration_s]&.to_i&.seconds || 30.minutes # Default of 30 minutes
 
-        Gitlab::Database::Migrations::Runner.batched_background_migrations(for_database: database_name)
+        database = args[:database] || 'main'
+        Gitlab::Database::Migrations::Runner.batched_background_migrations(for_database: database, legacy_mode: true)
                                             .run_jobs(for_duration: duration)
       end
     end