1 files changed, 20 insertions, 7 deletions
diff --git a/lib/gitlab/database/migration.rb b/lib/gitlab/database/migration.rb
index b2248b0f4eb..dc695a74a4b 100644
--- a/lib/gitlab/database/migration.rb
+++ b/lib/gitlab/database/migration.rb
@@ -33,20 +33,33 @@ module Gitlab
       # We use major version bumps to indicate significant changes and minor version bumps
       # to indicate backwards-compatible or otherwise minor changes (e.g. a Rails version bump).
       # However, this hasn't been strictly formalized yet.
-      MIGRATION_CLASSES = {
-        1.0 => Class.new(ActiveRecord::Migration[6.1]) do
-          include LockRetriesConcern
-          include Gitlab::Database::MigrationHelpers::V2
+
+      class V1_0 < ActiveRecord::Migration[6.1] # rubocop:disable Naming/ClassAndModuleCamelCase
+        include LockRetriesConcern
+        include Gitlab::Database::MigrationHelpers::V2
+      end
+
+      class V2_0 < V1_0 # rubocop:disable Naming/ClassAndModuleCamelCase
+        include Gitlab::Database::MigrationHelpers::RestrictGitlabSchema
+
+        # When running migrations, the `db:migrate` switches connection of
+        # ActiveRecord::Base depending where the migration runs.
+        # This helper class is provided to avoid confusion using `ActiveRecord::Base`
+        class MigrationRecord < ActiveRecord::Base
         end
-      }.freeze
+      end
 
       def self.[](version)
-        MIGRATION_CLASSES[version] || raise(ArgumentError, "Unknown migration version: #{version}")
+        version = version.to_s
+        name = "V#{version.tr('.', '_')}"
+        raise ArgumentError, "Unknown migration version: #{version}" unless const_defined?(name, false)
+
+        const_get(name, false)
       end
 
       # The current version to be used in new migrations
       def self.current_version
-        1.0
+        2.0
       end
     end
   end