1 files changed, 31 insertions, 4 deletions
diff --git a/lib/gitlab/database/migration_helpers.rb b/lib/gitlab/database/migration_helpers.rb
index e574422ce11..df40e3b3868 100644
--- a/lib/gitlab/database/migration_helpers.rb
+++ b/lib/gitlab/database/migration_helpers.rb
@@ -296,12 +296,11 @@ module Gitlab
 
           with_lock_retries do
             execute("LOCK TABLE #{target}, #{source} IN SHARE ROW EXCLUSIVE MODE") if reverse_lock_order
-
             execute <<-EOF.strip_heredoc
             ALTER TABLE #{source}
             ADD CONSTRAINT #{options[:name]}
-            FOREIGN KEY (#{options[:column]})
-            REFERENCES #{target} (#{target_column})
+            FOREIGN KEY (#{multiple_columns(options[:column])})
+            REFERENCES #{target} (#{multiple_columns(target_column)})
             #{on_delete_statement(options[:on_delete])}
             NOT VALID;
             EOF
@@ -355,7 +354,7 @@ module Gitlab
       # - For standard rails foreign keys the prefix is `fk_rails_`
       #
       def concurrent_foreign_key_name(table, column, prefix: 'fk_')
-        identifier = "#{table}_#{column}_fk"
+        identifier = "#{table}_#{multiple_columns(column, separator: '_')}_fk"
         hashed_identifier = Digest::SHA256.hexdigest(identifier).first(10)
 
         "#{prefix}#{hashed_identifier}"
@@ -1503,6 +1502,26 @@ into similar problems in the future (e.g. when new tables are created).
         SQL
       end
 
+      def drop_constraint(table_name, constraint_name, cascade: false)
+        execute <<~SQL
+          ALTER TABLE #{quote_table_name(table_name)} DROP CONSTRAINT #{quote_column_name(constraint_name)} #{cascade_statement(cascade)}
+        SQL
+      end
+
+      def add_primary_key_using_index(table_name, pk_name, index_to_use)
+        execute <<~SQL
+          ALTER TABLE #{quote_table_name(table_name)} ADD CONSTRAINT #{quote_table_name(pk_name)} PRIMARY KEY USING INDEX #{quote_table_name(index_to_use)}
+        SQL
+      end
+
+      def swap_primary_key(table_name, primary_key_name, index_to_use)
+        with_lock_retries(raise_on_exhaustion: true) do
+          drop_constraint(table_name, primary_key_name, cascade: true)
+          add_primary_key_using_index(table_name, primary_key_name, index_to_use)
+        end
+      end
+      alias_method :unswap_primary_key, :swap_primary_key
+
       def drop_sequence(table_name, column_name, sequence_name)
         execute <<~SQL
           ALTER TABLE #{quote_table_name(table_name)} ALTER COLUMN #{quote_column_name(column_name)} DROP DEFAULT;
@@ -1519,6 +1538,14 @@ into similar problems in the future (e.g. when new tables are created).
 
       private
 
+      def multiple_columns(columns, separator: ', ')
+        Array.wrap(columns).join(separator)
+      end
+
+      def cascade_statement(cascade)
+        cascade ? 'CASCADE' : ''
+      end
+
       def create_temporary_columns_and_triggers(table, columns, primary_key: :id, data_type: :bigint)
         unless table_exists?(table)
           raise "Table #{table} does not exist"