Merge branch 'pb-order-locks-when-adding-foreign-keys' into 'master'

Add option to order locks on add_concurrent_foreign_key

See merge request gitlab-org/gitlab!67448

lib/gitlab/database/migration_helpers.rb

@@ -231,8 +231,13 @@ module Gitlab
       # on_delete - The action to perform when associated data is removed,
       #             defaults to "CASCADE".
       # name - The name of the foreign key.
+      # validate - Flag that controls whether the new foreign key will be validated after creation.
+      #            If the flag is not set, the constraint will only be enforced for new data.
+      # reverse_lock_order - Flag that controls whether we should attempt to acquire locks in the reverse
+      #                      order of the ALTER TABLE. This can be useful in situations where the foreign
+      #                      key creation could deadlock with another process.
       #
-      def add_concurrent_foreign_key(source, target, column:, on_delete: :cascade, target_column: :id, name: nil, validate: true)
+      def add_concurrent_foreign_key(source, target, column:, on_delete: :cascade, target_column: :id, name: nil, validate: true, reverse_lock_order: false)
         # Transactions would result in ALTER TABLE locks being held for the
         # duration of the transaction, defeating the purpose of this method.
         if transaction_open?
@@ -260,6 +265,8 @@ module Gitlab
           # data.
 
           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]}

spec/lib/gitlab/database/migration_helpers_spec.rb

@@ -586,6 +586,22 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
           it_behaves_like 'performs validation', {}
         end
       end
+
+      context 'when the reverse_lock_order flag is set' do
+        it 'explicitly locks the tables in target-source order', :aggregate_failures do
+          expect(model).to receive(:with_lock_retries).and_call_original
+          expect(model).to receive(:disable_statement_timeout).and_call_original
+          expect(model).to receive(:statement_timeout_disabled?).and_return(false)
+          expect(model).to receive(:execute).with(/statement_timeout/)
+          expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
+          expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+
+          expect(model).to receive(:execute).with('LOCK TABLE users, projects IN SHARE ROW EXCLUSIVE MODE')
+          expect(model).to receive(:execute).with(/REFERENCES users \(id\)/)
+
+          model.add_concurrent_foreign_key(:projects, :users, column: :user_id, reverse_lock_order: true)
+        end
+      end
     end
   end