Improve UniquenessHelpers#ensure_unique_id

This commit will raise an error when detecting multiple
sequences owned by the same table. This will avoid weird situation
where the id was not generated correctly for some instances.

lib/gitlab/database/partitioning_migration_helpers/uniqueness_helpers.rb

@@ -7,20 +7,33 @@ module UniquenessHelpers
         include Gitlab::Database::MigrationHelpers
         include Gitlab::Database::SchemaHelpers
 
+        COL_NAME = :id
+        SequenceError = Class.new(StandardError)
+
         def ensure_unique_id(table_name)
           function_name = "assign_#{table_name}_id_value"
           trigger_name = "assign_#{table_name}_id_trigger"
+          sequences = existing_sequence(table_name, COL_NAME)
+
+          if sequences.many? || sequences.none?
+            raise(SequenceError, <<~MESSAGE)
+              Expected to find only one sequence for #{table_name}(#{COL_NAME}) but found #{sequences.size}.
+              Please ensure that there is only one sequence before proceeding.
+              Found sequences: #{sequences.map(&:seq_name)}
+            MESSAGE
+          end
 
           return if trigger_exists?(table_name, trigger_name)
 
-          change_column_default(table_name, :id, nil)
+          sequence_name = sequences.first.seq_name
+          change_column_default(table_name, COL_NAME, nil)
 
           create_trigger_function(function_name) do
             <<~SQL
               IF NEW."id" IS NOT NULL THEN
                 RAISE WARNING 'Manually assigning ids is not allowed, the value will be ignored';
               END IF;
-              NEW."id" := nextval(\'#{existing_sequence(table_name)}\'::regclass);
+              NEW."id" := nextval(\'#{sequence_name}\'::regclass);
               RETURN NEW;
             SQL
           end
@@ -30,8 +43,11 @@ def ensure_unique_id(table_name)
 
         private
 
-        def existing_sequence(table_name)
-          Gitlab::Database::PostgresSequence.by_table_name(table_name).first.seq_name
+        def existing_sequence(table_name, col_name)
+          @existing_sequence ||= Gitlab::Database::PostgresSequence
+            .by_table_name(table_name)
+            .by_col_name(col_name)
+            .to_a
         end
       end
     end

lib/gitlab/database/postgres_sequence.rb

@@ -7,6 +7,7 @@ class PostgresSequence < SharedModel
       self.primary_key = :seq_name
 
       scope :by_table_name, ->(table_name) { where(table_name: table_name) }
+      scope :by_col_name, ->(col_name) { where(col_name: col_name) }
     end
   end
 end

spec/lib/gitlab/database/partitioning_migration_helpers/uniqueness_helpers_spec.rb

@@ -66,6 +66,36 @@
           expect(migration.trigger_exists?(table_name, trigger_name)).to eq(true)
         end
       end
+
+      context 'when table does not have a sequence' do
+        before do
+          allow(migration).to receive(:existing_sequence).with(table_name, :id).and_return([])
+        end
+
+        it 'raises SequenceError' do
+          expect do
+            ensure_unique_id
+          end.to raise_error(described_class::SequenceError, /Expected to find only one sequence for/)
+        end
+      end
+
+      context 'when table has multiple sequences attached to it' do
+        before do
+          connection.execute(<<~SQL)
+            CREATE SEQUENCE second_sequence
+              START 0
+              INCREMENT 1
+              MINVALUE 0
+              OWNED BY _test_partitioned_table.id;
+          SQL
+        end
+
+        it 'raises SequenceError' do
+          expect do
+            ensure_unique_id
+          end.to raise_error(described_class::SequenceError, /Expected to find only one sequence/)
+        end
+      end
     end
   end
 end

spec/lib/gitlab/database/postgres_sequences_spec.rb

@@ -8,6 +8,7 @@
   let(:schema) { ActiveRecord::Base.connection.current_schema }
   let(:table_name) { '_test_table' }
   let(:table_name_without_sequence) { '_test_table_without_sequence' }
+  let(:col_name) { :id }
 
   before do
     ActiveRecord::Base.connection.execute(<<~SQL)
@@ -21,15 +22,23 @@
     SQL
   end
 
-  describe '#by_table_name' do
-    context 'when table does not have a sequence' do
-      it 'returns an empty collection' do
-        expect(described_class.by_table_name(table_name_without_sequence)).to be_empty
+  describe 'scopes' do
+    describe '#by_table_name' do
+      context 'when table does not have a sequence' do
+        it 'returns an empty collection' do
+          expect(described_class.by_table_name(table_name_without_sequence)).to be_empty
+        end
+      end
+
+      it 'returns the sequence for a given table' do
+        expect(described_class.by_table_name(table_name).first[:table_name]).to eq(table_name)
       end
     end
 
-    it 'returns the sequence for a given table' do
-      expect(described_class.by_table_name(table_name).first[:table_name]).to eq(table_name)
+    describe '#by_col_name' do
+      it 'returns the sequence for a col name' do
+        expect(described_class.by_col_name(col_name).first[:table_name]).to eq(table_name)
+      end
     end
   end
 end