1 files changed, 70 insertions, 36 deletions

diff --git a/spec/lib/gitlab/database/postgres_foreign_key_spec.rb b/spec/lib/gitlab/database/postgres_foreign_key_spec.rb
index a8dbc4be16f..ae56f66737d 100644
--- a/spec/lib/gitlab/database/postgres_foreign_key_spec.rb
+++ b/spec/lib/gitlab/database/postgres_foreign_key_spec.rb
@@ -6,26 +6,32 @@ RSpec.describe Gitlab::Database::PostgresForeignKey, type: :model, feature_categ
   # PostgresForeignKey does not `behaves_like 'a postgres model'` because it does not correspond 1-1 with a single entry
   # in pg_class
 
+  let(:table_prefix) { '_test_gitlab_main' }
+
   before do
     ApplicationRecord.connection.execute(<<~SQL)
-      CREATE TABLE public.referenced_table (
+      CREATE TABLE #{schema_table_name('referenced_table')} (
         id bigserial primary key not null,
         id_b bigserial not null,
         UNIQUE (id, id_b)
       );
 
-      CREATE TABLE public.other_referenced_table (
+      CREATE TABLE #{schema_table_name('other_referenced_table')} (
         id bigserial primary key not null
       );
 
-      CREATE TABLE public.constrained_table (
+      CREATE TABLE #{schema_table_name('constrained_table')} (
         id bigserial primary key not null,
         referenced_table_id bigint not null,
         referenced_table_id_b bigint not null,
         other_referenced_table_id bigint not null,
-        CONSTRAINT fk_constrained_to_referenced FOREIGN KEY(referenced_table_id, referenced_table_id_b) REFERENCES referenced_table(id, id_b) on delete restrict,
+        CONSTRAINT fk_constrained_to_referenced
+          FOREIGN KEY(referenced_table_id, referenced_table_id_b)
+          REFERENCES #{table_name('referenced_table')}(id, id_b)
+          ON DELETE restrict
+          ON UPDATE restrict,
         CONSTRAINT fk_constrained_to_other_referenced FOREIGN KEY(other_referenced_table_id)
-           REFERENCES other_referenced_table(id)
+           REFERENCES #{table_name('other_referenced_table')}(id)
       );
 
     SQL
@@ -33,13 +39,13 @@ RSpec.describe Gitlab::Database::PostgresForeignKey, type: :model, feature_categ
 
   describe '#by_referenced_table_identifier' do
     it 'throws an error when the identifier name is not fully qualified' do
-      expect { described_class.by_referenced_table_identifier('referenced_table') }.to raise_error(ArgumentError, /not fully qualified/)
+      expect { described_class.by_referenced_table_identifier(table_name("referenced_table")) }.to raise_error(ArgumentError, /not fully qualified/)
     end
 
     it 'finds the foreign keys for the referenced table' do
       expected = described_class.find_by!(name: 'fk_constrained_to_referenced')
 
-      expect(described_class.by_referenced_table_identifier('public.referenced_table')).to contain_exactly(expected)
+      expect(described_class.by_referenced_table_identifier(schema_table_name("referenced_table"))).to contain_exactly(expected)
     end
   end
 
@@ -47,19 +53,19 @@ RSpec.describe Gitlab::Database::PostgresForeignKey, type: :model, feature_categ
     it 'finds the foreign keys for the referenced table' do
       expected = described_class.find_by!(name: 'fk_constrained_to_referenced')
 
-      expect(described_class.by_referenced_table_name('referenced_table')).to contain_exactly(expected)
+      expect(described_class.by_referenced_table_name(table_name("referenced_table"))).to contain_exactly(expected)
     end
   end
 
   describe '#by_constrained_table_identifier' do
     it 'throws an error when the identifier name is not fully qualified' do
-      expect { described_class.by_constrained_table_identifier('constrained_table') }.to raise_error(ArgumentError, /not fully qualified/)
+      expect { described_class.by_constrained_table_identifier(table_name("constrained_table")) }.to raise_error(ArgumentError, /not fully qualified/)
     end
 
     it 'finds the foreign keys for the constrained table' do
       expected = described_class.where(name: %w[fk_constrained_to_referenced fk_constrained_to_other_referenced]).to_a
 
-      expect(described_class.by_constrained_table_identifier('public.constrained_table')).to match_array(expected)
+      expect(described_class.by_constrained_table_identifier(schema_table_name('constrained_table'))).to match_array(expected)
     end
   end
 
@@ -67,7 +73,7 @@ RSpec.describe Gitlab::Database::PostgresForeignKey, type: :model, feature_categ
     it 'finds the foreign keys for the constrained table' do
       expected = described_class.where(name: %w[fk_constrained_to_referenced fk_constrained_to_other_referenced]).to_a
 
-      expect(described_class.by_constrained_table_name('constrained_table')).to match_array(expected)
+      expect(described_class.by_constrained_table_name(table_name("constrained_table"))).to match_array(expected)
     end
   end
 
@@ -80,7 +86,7 @@ RSpec.describe Gitlab::Database::PostgresForeignKey, type: :model, feature_categ
   context 'when finding columns for foreign keys' do
     using RSpec::Parameterized::TableSyntax
 
-    let(:fks) { described_class.by_constrained_table_name('constrained_table') }
+    let(:fks) { described_class.by_constrained_table_name(table_name("constrained_table")) }
 
     where(:fk, :expected_constrained, :expected_referenced) do
       lazy { described_class.find_by(name: 'fk_constrained_to_referenced') } | %w[referenced_table_id referenced_table_id_b] | %w[id id_b]
@@ -110,25 +116,34 @@ RSpec.describe Gitlab::Database::PostgresForeignKey, type: :model, feature_categ
     end
   end
 
-  describe '#on_delete_action' do
+  describe '#on_delete_action and #on_update_action' do
     before do
       ApplicationRecord.connection.execute(<<~SQL)
-        create table public.referenced_table_all_on_delete_actions (
+        create table #{schema_table_name('referenced_table_all_on_delete_actions')} (
           id bigserial primary key not null
         );
 
-        create table public.constrained_table_all_on_delete_actions (
+        create table #{schema_table_name('constrained_table_all_on_delete_actions')} (
           id bigserial primary key not null,
-          ref_id_no_action bigint not null constraint fk_no_action references referenced_table_all_on_delete_actions(id),
-          ref_id_restrict bigint not null constraint fk_restrict references referenced_table_all_on_delete_actions(id) on delete restrict,
-          ref_id_nullify bigint not null constraint fk_nullify references referenced_table_all_on_delete_actions(id) on delete set null,
-          ref_id_cascade bigint not null constraint fk_cascade references referenced_table_all_on_delete_actions(id) on delete cascade,
-          ref_id_set_default bigint not null constraint fk_set_default references referenced_table_all_on_delete_actions(id) on delete set default
+          ref_id_no_action bigint not null constraint fk_no_action
+            references #{table_name('referenced_table_all_on_delete_actions')}(id),
+          ref_id_restrict bigint not null constraint fk_restrict
+            references #{table_name('referenced_table_all_on_delete_actions')}(id)
+            on delete restrict on update restrict,
+          ref_id_nullify bigint not null constraint fk_nullify
+            references #{table_name('referenced_table_all_on_delete_actions')}(id)
+            on delete set null on update set null,
+          ref_id_cascade bigint not null constraint fk_cascade
+            references #{table_name('referenced_table_all_on_delete_actions')}(id)
+            on delete cascade on update cascade,
+          ref_id_set_default bigint not null constraint fk_set_default
+            references #{table_name('referenced_table_all_on_delete_actions')}(id)
+            on delete set default on update set default
         )
       SQL
     end
 
-    let(:fks) { described_class.by_constrained_table_name('constrained_table_all_on_delete_actions') }
+    let(:fks) { described_class.by_constrained_table_name(table_name('constrained_table_all_on_delete_actions')) }
 
     context 'with an invalid on_delete_action' do
       it 'raises an error' do
@@ -137,7 +152,7 @@ RSpec.describe Gitlab::Database::PostgresForeignKey, type: :model, feature_categ
       end
     end
 
-    where(:fk_name, :expected_on_delete_action) do
+    where(:fk_name, :expected_action) do
       [
         %w[fk_no_action no_action],
         %w[fk_restrict restrict],
@@ -151,12 +166,22 @@ RSpec.describe Gitlab::Database::PostgresForeignKey, type: :model, feature_categ
       subject(:fk) { fks.find_by(name: fk_name) }
 
       it 'has the appropriate on delete action' do
-        expect(fk.on_delete_action).to eq(expected_on_delete_action)
+        expect(fk.on_delete_action).to eq(expected_action)
+      end
+
+      it 'has the appropriate on update action' do
+        expect(fk.on_update_action).to eq(expected_action)
       end
 
       describe '#by_on_delete_action' do
         it 'finds the key by on delete action' do
-          expect(fks.by_on_delete_action(expected_on_delete_action)).to contain_exactly(fk)
+          expect(fks.by_on_delete_action(expected_action)).to contain_exactly(fk)
+        end
+      end
+
+      describe '#by_on_update_action' do
+        it 'finds the key by on update action' do
+          expect(fks.by_on_update_action(expected_action)).to contain_exactly(fk)
         end
       end
     end
@@ -167,16 +192,17 @@ RSpec.describe Gitlab::Database::PostgresForeignKey, type: :model, feature_categ
       skip('not supported before postgres 12') if ApplicationRecord.database.version.to_f < 12
 
       ApplicationRecord.connection.execute(<<~SQL)
-      create table public.parent (
-        id bigserial primary key not null
-      ) partition by hash(id);
+        create table #{schema_table_name('parent')} (
+          id bigserial primary key not null
+        ) partition by hash(id);
 
-      create table public.child partition of parent for values with (modulus 2, remainder 1);
+        create table #{schema_table_name('child')} partition of #{table_name('parent')}
+          for values with (modulus 2, remainder 1);
 
-      create table public.referencing_partitioned (
-        id bigserial not null primary key,
-        constraint fk_inherited foreign key (id) references parent(id)
-      )
+        create table #{schema_table_name('referencing_partitioned')} (
+          id bigserial not null primary key,
+          constraint fk_inherited foreign key (id) references #{table_name('parent')}(id)
+        )
       SQL
     end
 
@@ -185,7 +211,7 @@ RSpec.describe Gitlab::Database::PostgresForeignKey, type: :model, feature_categ
 
       where(:fk, :inherited) do
         lazy { described_class.find_by(name: 'fk_inherited') } | false
-        lazy { described_class.by_referenced_table_identifier('public.child').first! } | true
+        lazy { described_class.by_referenced_table_identifier(schema_table_name('child')).first! } | true
         lazy { described_class.find_by(name: 'fk_constrained_to_referenced') } | false
       end
 
@@ -197,12 +223,20 @@ RSpec.describe Gitlab::Database::PostgresForeignKey, type: :model, feature_categ
     end
 
     describe '#not_inherited' do
-      let(:fks) { described_class.by_constrained_table_identifier('public.referencing_partitioned') }
+      let(:fks) { described_class.by_constrained_table_identifier(schema_table_name('referencing_partitioned')) }
 
       it 'lists all non-inherited foreign keys' do
-        expect(fks.pluck(:referenced_table_name)).to contain_exactly('parent', 'child')
-        expect(fks.not_inherited.pluck(:referenced_table_name)).to contain_exactly('parent')
+        expect(fks.pluck(:referenced_table_name)).to contain_exactly(table_name('parent'), table_name('child'))
+        expect(fks.not_inherited.pluck(:referenced_table_name)).to contain_exactly(table_name('parent'))
       end
     end
   end
+
+  def schema_table_name(name)
+    "public.#{table_name(name)}"
+  end
+
+  def table_name(name)
+    "#{table_prefix}_#{name}"
+  end
 end