Shortcut concurrent foreign key creation if already exists.ab-43887-concurrent-migration-helpers

Closes #43887.

6 files changed, 92 insertions, 51 deletions

diff --git a/db/migrate/20170530130129_project_foreign_keys_with_cascading_deletes.rb b/db/migrate/20170530130129_project_foreign_keys_with_cascading_deletes.rb
index af6d10b5158..1199073ed3a 100644
--- a/db/migrate/20170530130129_project_foreign_keys_with_cascading_deletes.rb
+++ b/db/migrate/20170530130129_project_foreign_keys_with_cascading_deletes.rb
@@ -154,7 +154,7 @@ class ProjectForeignKeysWithCascadingDeletes < ActiveRecord::Migration
   end
 
   def add_foreign_key_if_not_exists(source, target, column:)
-    return if foreign_key_exists?(source, column)
+    return if foreign_key_exists?(source, target, column: column)
 
     add_concurrent_foreign_key(source, target, column: column)
   end
@@ -175,12 +175,6 @@ class ProjectForeignKeysWithCascadingDeletes < ActiveRecord::Migration
   rescue ArgumentError
   end
 
-  def foreign_key_exists?(table, column)
-    foreign_keys(table).any? do |key|
-      key.options[:column] == column.to_s
-    end
-  end
-
   def connection
     # Rails memoizes connection objects, but this causes them to be shared
     # amongst threads; we don't want that.
diff --git a/db/migrate/20170703102400_add_stage_id_foreign_key_to_builds.rb b/db/migrate/20170703102400_add_stage_id_foreign_key_to_builds.rb
index 68b947583d3..a89d348b127 100644
--- a/db/migrate/20170703102400_add_stage_id_foreign_key_to_builds.rb
+++ b/db/migrate/20170703102400_add_stage_id_foreign_key_to_builds.rb
@@ -10,13 +10,13 @@ class AddStageIdForeignKeyToBuilds < ActiveRecord::Migration
       add_concurrent_index(:ci_builds, :stage_id)
     end
 
-    unless foreign_key_exists?(:ci_builds, :stage_id)
+    unless foreign_key_exists?(:ci_builds, :ci_stages, column: :stage_id)
       add_concurrent_foreign_key(:ci_builds, :ci_stages, column: :stage_id, on_delete: :cascade)
     end
   end
 
   def down
-    if foreign_key_exists?(:ci_builds, :stage_id)
+    if foreign_key_exists?(:ci_builds, column: :stage_id)
       remove_foreign_key(:ci_builds, column: :stage_id)
     end
 
@@ -24,12 +24,4 @@ class AddStageIdForeignKeyToBuilds < ActiveRecord::Migration
       remove_concurrent_index(:ci_builds, :stage_id)
     end
   end
-
-  private
-
-  def foreign_key_exists?(table, column)
-    foreign_keys(:ci_builds).any? do |key|
-      key.options[:column] == column.to_s
-    end
-  end
 end
diff --git a/db/migrate/20170713104829_add_foreign_key_to_merge_requests.rb b/db/migrate/20170713104829_add_foreign_key_to_merge_requests.rb
index c25d4fd5986..c409915ceed 100644
--- a/db/migrate/20170713104829_add_foreign_key_to_merge_requests.rb
+++ b/db/migrate/20170713104829_add_foreign_key_to_merge_requests.rb
@@ -23,23 +23,15 @@ class AddForeignKeyToMergeRequests < ActiveRecord::Migration
       merge_requests.update_all(head_pipeline_id: nil)
     end
 
-    unless foreign_key_exists?(:merge_requests, :head_pipeline_id)
+    unless foreign_key_exists?(:merge_requests, column: :head_pipeline_id)
       add_concurrent_foreign_key(:merge_requests, :ci_pipelines,
                                  column: :head_pipeline_id, on_delete: :nullify)
     end
   end
 
   def down
-    if foreign_key_exists?(:merge_requests, :head_pipeline_id)
+    if foreign_key_exists?(:merge_requests, column: :head_pipeline_id)
       remove_foreign_key(:merge_requests, column: :head_pipeline_id)
     end
   end
-
-  private
-
-  def foreign_key_exists?(table, column)
-    foreign_keys(table).any? do |key|
-      key.options[:column] == column.to_s
-    end
-  end
 end
diff --git a/db/post_migrate/20180223124427_build_user_interacted_projects_table.rb b/db/post_migrate/20180223124427_build_user_interacted_projects_table.rb
index d1a29a5c71b..9addd36dca6 100644
--- a/db/post_migrate/20180223124427_build_user_interacted_projects_table.rb
+++ b/db/post_migrate/20180223124427_build_user_interacted_projects_table.rb
@@ -26,11 +26,11 @@ class BuildUserInteractedProjectsTable < ActiveRecord::Migration
   def down
     execute "TRUNCATE user_interacted_projects"
 
-    if foreign_key_exists?(:user_interacted_projects, :user_id)
+    if foreign_key_exists?(:user_interacted_projects, :users)
       remove_foreign_key :user_interacted_projects, :users
     end
 
-    if foreign_key_exists?(:user_interacted_projects, :project_id)
+    if foreign_key_exists?(:user_interacted_projects, :projects)
       remove_foreign_key :user_interacted_projects, :projects
     end
 
@@ -115,7 +115,7 @@ class BuildUserInteractedProjectsTable < ActiveRecord::Migration
     end
 
     def create_fk(table, target, column)
-      return if foreign_key_exists?(table, column)
+      return if foreign_key_exists?(table, target, column: column)
 
       add_foreign_key table, target, column: column, on_delete: :cascade
     end
@@ -158,11 +158,11 @@ class BuildUserInteractedProjectsTable < ActiveRecord::Migration
         add_concurrent_index :user_interacted_projects, [:project_id, :user_id], unique: true, name: UNIQUE_INDEX_NAME
       end
 
-      unless foreign_key_exists?(:user_interacted_projects, :user_id)
+      unless foreign_key_exists?(:user_interacted_projects, :users, column: :user_id)
         add_concurrent_foreign_key :user_interacted_projects, :users, column: :user_id, on_delete: :cascade
       end
 
-      unless foreign_key_exists?(:user_interacted_projects, :project_id)
+      unless foreign_key_exists?(:user_interacted_projects, :projects, column: :project_id)
         add_concurrent_foreign_key :user_interacted_projects, :projects, column: :project_id, on_delete: :cascade
       end
     end
diff --git a/lib/gitlab/database/migration_helpers.rb b/lib/gitlab/database/migration_helpers.rb
index 55160ca8708..44ca434056f 100644
--- a/lib/gitlab/database/migration_helpers.rb
+++ b/lib/gitlab/database/migration_helpers.rb
@@ -155,6 +155,13 @@ module Gitlab
         # of PostgreSQL's "VALIDATE CONSTRAINT". As a result we'll just fall
         # back to the normal foreign key procedure.
         if Database.mysql?
+          if foreign_key_exists?(source, target, column: column)
+            Rails.logger.warn "Foreign key not created because it exists already " \
+              "(this may be due to an aborted migration or similar): " \
+              "source: #{source}, target: #{target}, column: #{column}"
+            return
+          end
+
           return add_foreign_key(source, target,
                                  column: column,
                                  on_delete: on_delete)
@@ -166,25 +173,43 @@ module Gitlab
 
         key_name = concurrent_foreign_key_name(source, column)
 
-        # Using NOT VALID allows us to create a key without immediately
-        # validating it. This means we keep the ALTER TABLE lock only for a
-        # short period of time. The key _is_ enforced for any newly created
-        # data.
-        execute <<-EOF.strip_heredoc
-        ALTER TABLE #{source}
-        ADD CONSTRAINT #{key_name}
-        FOREIGN KEY (#{column})
-        REFERENCES #{target} (id)
-        #{on_delete ? "ON DELETE #{on_delete.upcase}" : ''}
-        NOT VALID;
-        EOF
+        unless foreign_key_exists?(source, target, column: column)
+          Rails.logger.warn "Foreign key not created because it exists already " \
+            "(this may be due to an aborted migration or similar): " \
+            "source: #{source}, target: #{target}, column: #{column}"
+
+          # Using NOT VALID allows us to create a key without immediately
+          # validating it. This means we keep the ALTER TABLE lock only for a
+          # short period of time. The key _is_ enforced for any newly created
+          # data.
+          execute <<-EOF.strip_heredoc
+          ALTER TABLE #{source}
+          ADD CONSTRAINT #{key_name}
+          FOREIGN KEY (#{column})
+          REFERENCES #{target} (id)
+          #{on_delete ? "ON DELETE #{on_delete.upcase}" : ''}
+          NOT VALID;
+          EOF
+        end
 
         # Validate the existing constraint. This can potentially take a very
         # long time to complete, but fortunately does not lock the source table
         # while running.
+        #
+        # Note this is a no-op in case the constraint is VALID already
         execute("ALTER TABLE #{source} VALIDATE CONSTRAINT #{key_name};")
       end
 
+      def foreign_key_exists?(source, target = nil, column: nil)
+        foreign_keys(source).any? do |key|
+          if column
+            key.options[:column].to_s == column.to_s
+          else
+            key.to_table.to_s == target.to_s
+          end
+        end
+      end
+
       # Returns the name for a concurrent foreign key.
       #
       # PostgreSQL constraint names have a limit of 63 bytes. The logic used
@@ -875,12 +900,6 @@ into similar problems in the future (e.g. when new tables are created).
         end
       end
 
-      def foreign_key_exists?(table, column)
-        foreign_keys(table).any? do |key|
-          key.options[:column] == column.to_s
-        end
-      end
-
       # Rails' index_exists? doesn't work when you only give it a table and index
       # name. As such we have to use some extra code to check if an index exists for
       # a given name.
diff --git a/spec/lib/gitlab/database/migration_helpers_spec.rb b/spec/lib/gitlab/database/migration_helpers_spec.rb
index 9074e17ae80..a41b7f4e104 100644
--- a/spec/lib/gitlab/database/migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migration_helpers_spec.rb
@@ -183,6 +183,10 @@ describe Gitlab::Database::MigrationHelpers do
   end
 
   describe '#add_concurrent_foreign_key' do
+    before do
+      allow(model).to receive(:foreign_key_exists?).and_return(false)
+    end
+
     context 'inside a transaction' do
       it 'raises an error' do
         expect(model).to receive(:transaction_open?).and_return(true)
@@ -199,14 +203,23 @@ describe Gitlab::Database::MigrationHelpers do
       end
 
       context 'using MySQL' do
-        it 'creates a regular foreign key' do
+        before do
           allow(Gitlab::Database).to receive(:mysql?).and_return(true)
+        end
 
+        it 'creates a regular foreign key' do
           expect(model).to receive(:add_foreign_key)
             .with(:projects, :users, column: :user_id, on_delete: :cascade)
 
           model.add_concurrent_foreign_key(:projects, :users, column: :user_id)
         end
+
+        it 'does not create a foreign key if it exists already' do
+          expect(model).to receive(:foreign_key_exists?).with(:projects, :users, column: :user_id).and_return(true)
+          expect(model).not_to receive(:add_foreign_key)
+
+          model.add_concurrent_foreign_key(:projects, :users, column: :user_id)
+        end
       end
 
       context 'using PostgreSQL' do
@@ -231,6 +244,14 @@ describe Gitlab::Database::MigrationHelpers do
                                            column: :user_id,
                                            on_delete: :nullify)
         end
+
+        it 'does not create a foreign key if it exists already' do
+          expect(model).to receive(:foreign_key_exists?).with(:projects, :users, column: :user_id).and_return(true)
+          expect(model).not_to receive(:execute).with(/ADD CONSTRAINT/)
+          expect(model).to receive(:execute).with(/VALIDATE CONSTRAINT/)
+
+          model.add_concurrent_foreign_key(:projects, :users, column: :user_id)
+        end
       end
     end
   end
@@ -245,6 +266,29 @@ describe Gitlab::Database::MigrationHelpers do
     end
   end
 
+  describe '#foreign_key_exists?' do
+    before do
+      key = ActiveRecord::ConnectionAdapters::ForeignKeyDefinition.new(:projects, :users, { column: :non_standard_id })
+      allow(model).to receive(:foreign_keys).with(:projects).and_return([key])
+    end
+
+    it 'finds existing foreign keys by column' do
+      expect(model.foreign_key_exists?(:projects, :users, column: :non_standard_id)).to be_truthy
+    end
+
+    it 'finds existing foreign keys by target table only' do
+      expect(model.foreign_key_exists?(:projects, :users)).to be_truthy
+    end
+
+    it 'compares by column name if given' do
+      expect(model.foreign_key_exists?(:projects, :users, column: :user_id)).to be_falsey
+    end
+
+    it 'compares by target if no column given' do
+      expect(model.foreign_key_exists?(:projects, :other_table)).to be_falsey
+    end
+  end
+
   describe '#disable_statement_timeout' do
     context 'using PostgreSQL' do
       it 'disables statement timeouts' do