Improves add_timestamps_with_timezone helper

This improves the `add_timestamps_with_timezone` helper by allowing the
column names to be configured. This has the advantage that unnecessary
columns can be avoided, saving space.

A helper for removing the columns is also provided, to be used in the
`down` method of migrations.

lib/gitlab/database/migration_helpers.rb

@@ -6,31 +6,45 @@ module Gitlab
       BACKGROUND_MIGRATION_BATCH_SIZE = 1000 # Number of rows to process per job
       BACKGROUND_MIGRATION_JOB_BUFFER_SIZE = 1000 # Number of jobs to bulk queue at a time
 
+      PERMITTED_TIMESTAMP_COLUMNS = %i[created_at updated_at deleted_at].to_set.freeze
+      DEFAULT_TIMESTAMP_COLUMNS = %i[created_at updated_at].freeze
+
       # Adds `created_at` and `updated_at` columns with timezone information.
       #
       # This method is an improved version of Rails' built-in method `add_timestamps`.
       #
+      # By default, adds `created_at` and `updated_at` columns, but these can be specified as:
+      #
+      #   add_timestamps_with_timezone(:my_table, columns: [:created_at, :deleted_at])
+      #
+      # This allows you to create just the timestamps you need, saving space.
+      #
       # Available options are:
-      # default - The default value for the column.
-      # null - When set to `true` the column will allow NULL values.
+      #  :default - The default value for the column.
+      #  :null - When set to `true` the column will allow NULL values.
       #        The default is to not allow NULL values.
+      #  :columns - the column names to create. Must be one
+      #             of `Gitlab::Database::MigrationHelpers::PERMITTED_TIMESTAMP_COLUMNS`.
+      #             Default value: `DEFAULT_TIMESTAMP_COLUMNS`
+      #
+      # All options are optional.
       def add_timestamps_with_timezone(table_name, options = {})
         options[:null] = false if options[:null].nil?
+        columns = options.fetch(:columns, DEFAULT_TIMESTAMP_COLUMNS)
+        default_value = options[:default]
 
-        [:created_at, :updated_at].each do |column_name|
-          if options[:default] && transaction_open?
-            raise '`add_timestamps_with_timezone` with default value cannot be run inside a transaction. ' \
-              'You can disable transactions by calling `disable_ddl_transaction!` ' \
-              'in the body of your migration class'
-          end
+        validate_not_in_transaction!(:add_timestamps_with_timezone, 'with default value') if default_value
+
+        columns.each do |column_name|
+          validate_timestamp_column_name!(column_name)
 
           # If default value is presented, use `add_column_with_default` method instead.
-          if options[:default]
+          if default_value
             add_column_with_default(
               table_name,
               column_name,
               :datetime_with_timezone,
-              default: options[:default],
+              default: default_value,
               allow_null: options[:null]
             )
           else
@@ -39,6 +53,21 @@ module Gitlab
         end
       end
 
+      # To be used in the `#down` method of migrations that
+      # use `#add_timestamps_with_timezone`.
+      #
+      # Available options are:
+      #  :columns - the column names to remove. Must be one
+      #             Default value: `DEFAULT_TIMESTAMP_COLUMNS`
+      #
+      # All options are optional.
+      def remove_timestamps(table_name, options = {})
+        columns = options.fetch(:columns, DEFAULT_TIMESTAMP_COLUMNS)
+        columns.each do |column_name|
+          remove_column(table_name, column_name)
+        end
+      end
+
       # Creates a new index, concurrently when supported
       #
       # On PostgreSQL this method creates an index concurrently, on MySQL this
@@ -1098,6 +1127,28 @@ into similar problems in the future (e.g. when new tables are created).
       def mysql_compatible_index_length
         Gitlab::Database.mysql? ? 20 : nil
       end
+
+      private
+
+      def validate_timestamp_column_name!(column_name)
+        return if PERMITTED_TIMESTAMP_COLUMNS.member?(column_name)
+
+        raise <<~MESSAGE
+          Illegal timestamp column name! Got #{column_name}.
+          Must be one of: #{PERMITTED_TIMESTAMP_COLUMNS.to_a}
+        MESSAGE
+      end
+
+      def validate_not_in_transaction!(method_name, modifier = nil)
+        return unless transaction_open?
+
+        raise <<~ERROR
+          #{["`#{method_name}`", modifier].compact.join(' ')} cannot be run inside a transaction.
+
+          You can disable transactions by calling `disable_ddl_transaction!` in the body of
+          your migration class
+        ERROR
+      end
     end
   end
 end

spec/lib/gitlab/database/migration_helpers_spec.rb

@@ -9,9 +9,27 @@ describe Gitlab::Database::MigrationHelpers do
     allow(model).to receive(:puts)
   end
 
+  describe '#remove_timestamps' do
+    it 'can remove the default timestamps' do
+      Gitlab::Database::MigrationHelpers::DEFAULT_TIMESTAMP_COLUMNS.each do |column_name|
+        expect(model).to receive(:remove_column).with(:foo, column_name)
+      end
+
+      model.remove_timestamps(:foo)
+    end
+
+    it 'can remove custom timestamps' do
+      expect(model).to receive(:remove_column).with(:foo, :bar)
+
+      model.remove_timestamps(:foo, columns: [:bar])
+    end
+  end
+
   describe '#add_timestamps_with_timezone' do
+    let(:in_transaction) { false }
+
     before do
-      allow(model).to receive(:transaction_open?).and_return(false)
+      allow(model).to receive(:transaction_open?).and_return(in_transaction)
     end
 
     context 'using PostgreSQL' do
@@ -21,11 +39,64 @@ describe Gitlab::Database::MigrationHelpers do
       end
 
       it 'adds "created_at" and "updated_at" fields with the "datetime_with_timezone" data type' do
-        expect(model).to receive(:add_column).with(:foo, :created_at, :datetime_with_timezone, { null: false })
-        expect(model).to receive(:add_column).with(:foo, :updated_at, :datetime_with_timezone, { null: false })
+        Gitlab::Database::MigrationHelpers::DEFAULT_TIMESTAMP_COLUMNS.each do |column_name|
+          expect(model).to receive(:add_column).with(:foo, column_name, :datetime_with_timezone, { null: false })
+        end
 
         model.add_timestamps_with_timezone(:foo)
       end
+
+      it 'can disable the NOT NULL constraint' do
+        Gitlab::Database::MigrationHelpers::DEFAULT_TIMESTAMP_COLUMNS.each do |column_name|
+          expect(model).to receive(:add_column).with(:foo, column_name, :datetime_with_timezone, { null: true })
+        end
+
+        model.add_timestamps_with_timezone(:foo, null: true)
+      end
+
+      it 'can add just one column' do
+        expect(model).to receive(:add_column).with(:foo, :created_at, :datetime_with_timezone, anything)
+        expect(model).not_to receive(:add_column).with(:foo, :updated_at, :datetime_with_timezone, anything)
+
+        model.add_timestamps_with_timezone(:foo, columns: [:created_at])
+      end
+
+      it 'can add choice of acceptable columns' do
+        expect(model).to receive(:add_column).with(:foo, :created_at, :datetime_with_timezone, anything)
+        expect(model).to receive(:add_column).with(:foo, :deleted_at, :datetime_with_timezone, anything)
+        expect(model).not_to receive(:add_column).with(:foo, :updated_at, :datetime_with_timezone, anything)
+
+        model.add_timestamps_with_timezone(:foo, columns: [:created_at, :deleted_at])
+      end
+
+      it 'cannot add unacceptable column names' do
+        expect do
+          model.add_timestamps_with_timezone(:foo, columns: [:bar])
+        end.to raise_error %r/Illegal timestamp column name/
+      end
+
+      context 'in a transaction' do
+        let(:in_transaction) { true }
+
+        before do
+          allow(model).to receive(:add_column).with(any_args).and_call_original
+          allow(model).to receive(:add_column)
+            .with(:foo, anything, :datetime_with_timezone, anything)
+            .and_return(nil)
+        end
+
+        it 'cannot add a default value' do
+          expect do
+            model.add_timestamps_with_timezone(:foo, default: :i_cause_an_error)
+          end.to raise_error %r/add_timestamps_with_timezone/
+        end
+
+        it 'can add columns without defaults' do
+          expect do
+            model.add_timestamps_with_timezone(:foo)
+          end.not_to raise_error
+        end
+      end
     end
 
     context 'using MySQL' do