Add new methods to support the PK migration - STEP 3

-restore_conversion_of_integer_to_bigint adds back tmp columns/triggers

Related to https://gitlab.com/gitlab-org/gitlab/-/issues/339201

Changelog: added

lib/gitlab/database/migration_helpers.rb

@@ -968,42 +968,7 @@ module Gitlab
       # columns - The name, or array of names, of the column(s) that we want to convert to bigint.
       # primary_key - The name of the primary key column (most often :id)
       def initialize_conversion_of_integer_to_bigint(table, columns, primary_key: :id)
-        unless table_exists?(table)
-          raise "Table #{table} does not exist"
-        end
-
-        unless column_exists?(table, primary_key)
-          raise "Column #{primary_key} does not exist on #{table}"
-        end
-
-        columns = Array.wrap(columns)
-        columns.each do |column|
-          next if column_exists?(table, column)
-
-          raise ArgumentError, "Column #{column} does not exist on #{table}"
-        end
-
-        check_trigger_permissions!(table)
-
-        conversions = columns.to_h { |column| [column, convert_to_bigint_column(column)] }
-
-        with_lock_retries do
-          conversions.each do |(source_column, temporary_name)|
-            column = column_for(table, source_column)
-
-            if (column.name.to_s == primary_key.to_s) || !column.null
-              # If the column to be converted is either a PK or is defined as NOT NULL,
-              # set it to `NOT NULL DEFAULT 0` and we'll copy paste the correct values bellow
-              # That way, we skip the expensive validation step required to add
-              #  a NOT NULL constraint at the end of the process
-              add_column(table, temporary_name, :bigint, default: column.default || 0, null: false)
-            else
-              add_column(table, temporary_name, :bigint, default: column.default)
-            end
-          end
-
-          install_rename_triggers(table, conversions.keys, conversions.values)
-        end
+        create_temporary_columns_and_triggers(table, columns, primary_key: primary_key, data_type: :bigint)
       end
 
       # Reverts `initialize_conversion_of_integer_to_bigint`
@@ -1020,6 +985,16 @@ module Gitlab
         temporary_columns.each { |column| remove_column(table, column) }
       end
 
+      # Reverts `cleanup_conversion_of_integer_to_bigint`
+      #
+      # table - The name of the database table containing the columns
+      # columns - The name, or array of names, of the column(s) that we have converted to bigint.
+      # primary_key - The name of the primary key column (most often :id)
+
+      def restore_conversion_of_integer_to_bigint(table, columns, primary_key: :id)
+        create_temporary_columns_and_triggers(table, columns, primary_key: primary_key, data_type: :int)
+      end
+
       # Backfills the new columns used in an integer-to-bigint conversion using background migrations.
       #
       # - This helper should be called from a post-deployment migration.
@@ -1649,6 +1624,45 @@ into similar problems in the future (e.g. when new tables are created).
 
       private
 
+      def create_temporary_columns_and_triggers(table, columns, primary_key: :id, data_type: :bigint)
+        unless table_exists?(table)
+          raise "Table #{table} does not exist"
+        end
+
+        unless column_exists?(table, primary_key)
+          raise "Column #{primary_key} does not exist on #{table}"
+        end
+
+        columns = Array.wrap(columns)
+        columns.each do |column|
+          next if column_exists?(table, column)
+
+          raise ArgumentError, "Column #{column} does not exist on #{table}"
+        end
+
+        check_trigger_permissions!(table)
+
+        conversions = columns.to_h { |column| [column, convert_to_bigint_column(column)] }
+
+        with_lock_retries do
+          conversions.each do |(source_column, temporary_name)|
+            column = column_for(table, source_column)
+
+            if (column.name.to_s == primary_key.to_s) || !column.null
+              # If the column to be converted is either a PK or is defined as NOT NULL,
+              # set it to `NOT NULL DEFAULT 0` and we'll copy paste the correct values bellow
+              # That way, we skip the expensive validation step required to add
+              #  a NOT NULL constraint at the end of the process
+              add_column(table, temporary_name, data_type, default: column.default || 0, null: false)
+            else
+              add_column(table, temporary_name, data_type, default: column.default)
+            end
+          end
+
+          install_rename_triggers(table, conversions.keys, conversions.values)
+        end
+      end
+
       def validate_check_constraint_name!(constraint_name)
         if constraint_name.to_s.length > MAX_IDENTIFIER_NAME_LENGTH
           raise "The maximum allowed constraint name is #{MAX_IDENTIFIER_NAME_LENGTH} characters"

spec/lib/gitlab/database/migration_helpers_spec.rb

@@ -1886,6 +1886,61 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
     end
   end
 
+  describe '#restore_conversion_of_integer_to_bigint' do
+    let(:table) { :test_table }
+    let(:column) { :id }
+    let(:tmp_column) { model.convert_to_bigint_column(column) }
+
+    before do
+      model.create_table table, id: false do |t|
+        t.bigint :id, primary_key: true
+        t.bigint :build_id, null: false
+        t.timestamps
+      end
+    end
+
+    context 'when the target table does not exist' do
+      it 'raises an error' do
+        expect { model.restore_conversion_of_integer_to_bigint(:this_table_is_not_real, column) }
+          .to raise_error('Table this_table_is_not_real does not exist')
+      end
+    end
+
+    context 'when the column to migrate does not exist' do
+      it 'raises an error' do
+        expect { model.restore_conversion_of_integer_to_bigint(table, :this_column_is_not_real) }
+          .to raise_error(ArgumentError, "Column this_column_is_not_real does not exist on #{table}")
+      end
+    end
+
+    context 'when a single column is given' do
+      let(:column_to_convert) { 'id' }
+      let(:temporary_column) { model.convert_to_bigint_column(column_to_convert) }
+
+      it 'creates the correct columns and installs the trigger' do
+        expect(model).to receive(:add_column).with(table, temporary_column, :int, default: 0, null: false)
+
+        expect(model).to receive(:install_rename_triggers).with(table, [column_to_convert], [temporary_column])
+
+        model.restore_conversion_of_integer_to_bigint(table, column_to_convert)
+      end
+    end
+
+    context 'when multiple columns are given' do
+      let(:columns_to_convert) { %i[id build_id] }
+      let(:temporary_columns) { columns_to_convert.map { |column| model.convert_to_bigint_column(column) } }
+
+      it 'creates the correct columns and installs the trigger' do
+        expect(model).to receive(:add_column).with(table, temporary_columns[0], :int, default: 0, null: false)
+        expect(model).to receive(:add_column).with(table, temporary_columns[1], :int, default: 0, null: false)
+
+        expect(model).to receive(:install_rename_triggers).with(table, columns_to_convert, temporary_columns)
+
+        model.restore_conversion_of_integer_to_bigint(table, columns_to_convert)
+      end
+    end
+  end
+
   describe '#revert_initialize_conversion_of_integer_to_bigint' do
     let(:table) { :test_table }