Add latest changes from gitlab-org/gitlab@14-1-stable-eev14.1.0-rc42

1 files changed, 100 insertions, 150 deletions

diff --git a/spec/lib/gitlab/database/partitioning_migration_helpers/foreign_key_helpers_spec.rb b/spec/lib/gitlab/database/partitioning_migration_helpers/foreign_key_helpers_spec.rb
index 83f2436043c..a524fe681e9 100644
--- a/spec/lib/gitlab/database/partitioning_migration_helpers/foreign_key_helpers_spec.rb
+++ b/spec/lib/gitlab/database/partitioning_migration_helpers/foreign_key_helpers_spec.rb
@@ -3,192 +3,142 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::ForeignKeyHelpers do
-  include Database::TriggerHelpers
+  include Database::TableSchemaHelpers
 
-  let(:model) do
-    ActiveRecord::Migration.new.extend(described_class)
+  let(:migration) do
+    ActiveRecord::Migration.new.extend(Gitlab::Database::PartitioningMigrationHelpers)
   end
 
-  let_it_be(:connection) { ActiveRecord::Base.connection }
-
-  let(:referenced_table) { :issues }
-  let(:function_name) { '_test_partitioned_foreign_keys_function' }
-  let(:trigger_name) { '_test_partitioned_foreign_keys_trigger' }
+  let(:source_table_name) { '_test_partitioned_table' }
+  let(:target_table_name) { '_test_referenced_table' }
+  let(:column_name) { "#{target_table_name}_id" }
+  let(:foreign_key_name) { '_test_partitioned_fk' }
+  let(:partition_schema) { 'gitlab_partitions_dynamic' }
+  let(:partition1_name) { "#{partition_schema}.#{source_table_name}_202001" }
+  let(:partition2_name) { "#{partition_schema}.#{source_table_name}_202002" }
+  let(:options) do
+    {
+      column: column_name,
+      name: foreign_key_name,
+      on_delete: :cascade,
+      validate: true
+    }
+  end
 
   before do
-    allow(model).to receive(:puts)
-    allow(model).to receive(:fk_function_name).and_return(function_name)
-    allow(model).to receive(:fk_trigger_name).and_return(trigger_name)
+    allow(migration).to receive(:puts)
+
+    connection.execute(<<~SQL)
+      CREATE TABLE #{target_table_name} (
+        id serial NOT NULL,
+        PRIMARY KEY (id)
+      );
+
+      CREATE TABLE #{source_table_name} (
+        id serial NOT NULL,
+        #{column_name} int NOT NULL,
+        created_at timestamptz NOT NULL,
+        PRIMARY KEY (id, created_at)
+      ) PARTITION BY RANGE (created_at);
+
+      CREATE TABLE #{partition1_name} PARTITION OF #{source_table_name}
+      FOR VALUES FROM ('2020-01-01') TO ('2020-02-01');
+
+      CREATE TABLE #{partition2_name} PARTITION OF #{source_table_name}
+      FOR VALUES FROM ('2020-02-01') TO ('2020-03-01');
+    SQL
   end
 
-  describe 'adding a foreign key' do
+  describe '#add_concurrent_partitioned_foreign_key' do
     before do
-      allow(model).to receive(:transaction_open?).and_return(false)
-    end
-
-    context 'when the table has no foreign keys' do
-      it 'creates a trigger function to handle the single cascade' do
-        model.add_partitioned_foreign_key :issue_assignees, referenced_table
-
-        expect_function_to_contain(function_name, 'delete from issue_assignees where issue_id = old.id')
-        expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
-      end
-    end
-
-    context 'when the table already has foreign keys' do
-      context 'when the foreign key is from a different table' do
-        before do
-          model.add_partitioned_foreign_key :issue_assignees, referenced_table
-        end
-
-        it 'creates a trigger function to handle the multiple cascades' do
-          model.add_partitioned_foreign_key :epic_issues, referenced_table
-
-          expect_function_to_contain(function_name,
-            'delete from issue_assignees where issue_id = old.id',
-            'delete from epic_issues where issue_id = old.id')
-          expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
-        end
-      end
-
-      context 'when the foreign key is from the same table' do
-        before do
-          model.add_partitioned_foreign_key :issues, referenced_table, column: :moved_to_id
-        end
-
-        context 'when the foreign key is from a different column' do
-          it 'creates a trigger function to handle the multiple cascades' do
-            model.add_partitioned_foreign_key :issues, referenced_table, column: :duplicated_to_id
-
-            expect_function_to_contain(function_name,
-              'delete from issues where moved_to_id = old.id',
-              'delete from issues where duplicated_to_id = old.id')
-            expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
-          end
-        end
-
-        context 'when the foreign key is from the same column' do
-          it 'ignores the duplicate and properly recreates the trigger function' do
-            model.add_partitioned_foreign_key :issues, referenced_table, column: :moved_to_id
-
-            expect_function_to_contain(function_name, 'delete from issues where moved_to_id = old.id')
-            expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
-          end
-        end
-      end
-    end
+      allow(migration).to receive(:foreign_key_exists?)
+        .with(source_table_name, target_table_name, anything)
+        .and_return(false)
 
-    context 'when the foreign key is set to nullify' do
-      it 'creates a trigger function that nullifies the foreign key' do
-        model.add_partitioned_foreign_key :issue_assignees, referenced_table, on_delete: :nullify
-
-        expect_function_to_contain(function_name, 'update issue_assignees set issue_id = null where issue_id = old.id')
-        expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
-      end
+      allow(migration).to receive(:with_lock_retries).and_yield
     end
 
-    context 'when the referencing column is a custom value' do
-      it 'creates a trigger function with the correct column name' do
-        model.add_partitioned_foreign_key :issues, referenced_table, column: :duplicated_to_id
+    context 'when the foreign key does not exist on the parent table' do
+      it 'creates the foreign key on each partition, and the parent table' do
+        expect(migration).to receive(:foreign_key_exists?)
+          .with(source_table_name, target_table_name, **options)
+          .and_return(false)
 
-        expect_function_to_contain(function_name, 'delete from issues where duplicated_to_id = old.id')
-        expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
-      end
-    end
+        expect(migration).to receive(:concurrent_partitioned_foreign_key_name).and_return(foreign_key_name)
 
-    context 'when the referenced column is a custom value' do
-      let(:referenced_table) { :user_details }
+        expect_add_concurrent_fk_and_call_original(partition1_name, target_table_name, **options)
+        expect_add_concurrent_fk_and_call_original(partition2_name, target_table_name, **options)
 
-      it 'creates a trigger function with the correct column name' do
-        model.add_partitioned_foreign_key :user_preferences, referenced_table, column: :user_id, primary_key: :user_id
+        expect(migration).to receive(:with_lock_retries).ordered.and_yield
+        expect(migration).to receive(:add_foreign_key)
+          .with(source_table_name, target_table_name, **options)
+          .ordered
+          .and_call_original
 
-        expect_function_to_contain(function_name, 'delete from user_preferences where user_id = old.user_id')
-        expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
-      end
-    end
+        migration.add_concurrent_partitioned_foreign_key(source_table_name, target_table_name, column: column_name)
 
-    context 'when the given key definition is invalid' do
-      it 'raises an error with the appropriate message' do
-        expect do
-          model.add_partitioned_foreign_key :issue_assignees, referenced_table, column: :not_a_real_issue_id
-        end.to raise_error(/From column must be a valid column/)
+        expect_foreign_key_to_exist(source_table_name, foreign_key_name)
       end
-    end
-
-    context 'when run inside a transaction' do
-      it 'raises an error' do
-        expect(model).to receive(:transaction_open?).and_return(true)
 
-        expect do
-          model.add_partitioned_foreign_key :issue_assignees, referenced_table
-        end.to raise_error(/can not be run inside a transaction/)
+      def expect_add_concurrent_fk_and_call_original(source_table_name, target_table_name, options)
+        expect(migration).to receive(:add_concurrent_foreign_key)
+          .ordered
+          .with(source_table_name, target_table_name, options)
+          .and_wrap_original do |_, source_table_name, target_table_name, options|
+            connection.add_foreign_key(source_table_name, target_table_name, **options)
+          end
       end
     end
-  end
 
-  context 'removing a foreign key' do
-    before do
-      allow(model).to receive(:transaction_open?).and_return(false)
-    end
+    context 'when the foreign key exists on the parent table' do
+      it 'does not attempt to create any foreign keys' do
+        expect(migration).to receive(:concurrent_partitioned_foreign_key_name).and_return(foreign_key_name)
 
-    context 'when the table has multiple foreign keys' do
-      before do
-        model.add_partitioned_foreign_key :issue_assignees, referenced_table
-        model.add_partitioned_foreign_key :epic_issues, referenced_table
-      end
+        expect(migration).to receive(:foreign_key_exists?)
+          .with(source_table_name, target_table_name, **options)
+          .and_return(true)
 
-      it 'creates a trigger function without the removed cascade' do
-        expect_function_to_contain(function_name,
-          'delete from issue_assignees where issue_id = old.id',
-          'delete from epic_issues where issue_id = old.id')
-        expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
+        expect(migration).not_to receive(:add_concurrent_foreign_key)
+        expect(migration).not_to receive(:with_lock_retries)
+        expect(migration).not_to receive(:add_foreign_key)
 
-        model.remove_partitioned_foreign_key :issue_assignees, referenced_table
+        migration.add_concurrent_partitioned_foreign_key(source_table_name, target_table_name, column: column_name)
 
-        expect_function_to_contain(function_name, 'delete from epic_issues where issue_id = old.id')
-        expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
+        expect_foreign_key_not_to_exist(source_table_name, foreign_key_name)
       end
     end
 
-    context 'when the table has only one remaining foreign key' do
-      before do
-        model.add_partitioned_foreign_key :issue_assignees, referenced_table
+    context 'when additional foreign key options are given' do
+      let(:options) do
+        {
+          column: column_name,
+          name: '_my_fk_name',
+          on_delete: :restrict,
+          validate: true
+        }
       end
 
-      it 'removes the trigger function altogether' do
-        expect_function_to_contain(function_name, 'delete from issue_assignees where issue_id = old.id')
-        expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
-
-        model.remove_partitioned_foreign_key :issue_assignees, referenced_table
-
-        expect_function_not_to_exist(function_name)
-        expect_trigger_not_to_exist(referenced_table, trigger_name)
-      end
-    end
+      it 'forwards them to the foreign key helper methods' do
+        expect(migration).to receive(:foreign_key_exists?)
+          .with(source_table_name, target_table_name, **options)
+          .and_return(false)
 
-    context 'when the foreign key does not exist' do
-      before do
-        model.add_partitioned_foreign_key :issue_assignees, referenced_table
-      end
+        expect(migration).not_to receive(:concurrent_partitioned_foreign_key_name)
 
-      it 'ignores the invalid key and properly recreates the trigger function' do
-        expect_function_to_contain(function_name, 'delete from issue_assignees where issue_id = old.id')
-        expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
+        expect_add_concurrent_fk(partition1_name, target_table_name, **options)
+        expect_add_concurrent_fk(partition2_name, target_table_name, **options)
 
-        model.remove_partitioned_foreign_key :issues, referenced_table, column: :moved_to_id
+        expect(migration).to receive(:with_lock_retries).ordered.and_yield
+        expect(migration).to receive(:add_foreign_key).with(source_table_name, target_table_name, **options).ordered
 
-        expect_function_to_contain(function_name, 'delete from issue_assignees where issue_id = old.id')
-        expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
+        migration.add_concurrent_partitioned_foreign_key(source_table_name, target_table_name,
+          column: column_name, name: '_my_fk_name', on_delete: :restrict)
       end
-    end
-
-    context 'when run outside a transaction' do
-      it 'raises an error' do
-        expect(model).to receive(:transaction_open?).and_return(true)
 
-        expect do
-          model.remove_partitioned_foreign_key :issue_assignees, referenced_table
-        end.to raise_error(/can not be run inside a transaction/)
+      def expect_add_concurrent_fk(source_table_name, target_table_name, options)
+        expect(migration).to receive(:add_concurrent_foreign_key)
+          .ordered
+          .with(source_table_name, target_table_name, options)
       end
     end
   end