Add latest changes from gitlab-org/gitlab@13-0-stable-ee

9 files changed, 346 insertions, 77 deletions

diff --git a/lib/gitlab/database/batch_count.rb b/lib/gitlab/database/batch_count.rb
index 2359dceae48..ab069ce1da1 100644
--- a/lib/gitlab/database/batch_count.rb
+++ b/lib/gitlab/database/batch_count.rb
@@ -91,11 +91,17 @@ module Gitlab
 
       def batch_fetch(start, finish, mode)
         # rubocop:disable GitlabSecurity/PublicSend
-        @relation.select(@column).public_send(mode).where(@column => start..(finish - 1)).count
+        @relation.select(@column).public_send(mode).where(between_condition(start, finish)).count
       end
 
       private
 
+      def between_condition(start, finish)
+        return @column.between(start..(finish - 1)) if @column.is_a?(Arel::Attributes::Attribute)
+
+        { @column => start..(finish - 1) }
+      end
+
       def actual_start(start)
         start || @relation.minimum(@column) || 0
       end
diff --git a/lib/gitlab/database/count/reltuples_count_strategy.rb b/lib/gitlab/database/count/reltuples_count_strategy.rb
index 6cd90c01ab2..e226ed7613a 100644
--- a/lib/gitlab/database/count/reltuples_count_strategy.rb
+++ b/lib/gitlab/database/count/reltuples_count_strategy.rb
@@ -72,7 +72,7 @@ module Gitlab
         # @param [Array] table names
         # @returns [Hash] Table name to count mapping (e.g. { 'projects' => 5, 'users' => 100 })
         def get_statistics(table_names, check_statistics: true)
-          time = 1.hour.ago
+          time = 6.hours.ago
 
           query = PgClass.joins("LEFT JOIN pg_stat_user_tables USING (relname)")
             .where(relname: table_names)
diff --git a/lib/gitlab/database/migration_helpers.rb b/lib/gitlab/database/migration_helpers.rb
index cf5ff8ddb7b..96be057f77e 100644
--- a/lib/gitlab/database/migration_helpers.rb
+++ b/lib/gitlab/database/migration_helpers.rb
@@ -265,12 +265,19 @@ module Gitlab
       # or `RESET ALL` is executed
       def disable_statement_timeout
         if block_given?
-          begin
-            execute('SET statement_timeout TO 0')
-
+          if statement_timeout_disabled?
+            # Don't do anything if the statement_timeout is already disabled
+            # Allows for nested calls of disable_statement_timeout without
+            # resetting the timeout too early (before the outer call ends)
             yield
-          ensure
-            execute('RESET ALL')
+          else
+            begin
+              execute('SET statement_timeout TO 0')
+
+              yield
+            ensure
+              execute('RESET ALL')
+            end
           end
         else
           unless transaction_open?
@@ -378,7 +385,7 @@ module Gitlab
       # make things _more_ complex).
       #
       # `batch_column_name` option is for tables without primary key, in this
-      # case an other unique integer column can be used. Example: :user_id
+      # case another unique integer column can be used. Example: :user_id
       #
       # rubocop: disable Metrics/AbcSize
       def update_column_in_batches(table, column, value, batch_size: nil, batch_column_name: :id)
@@ -444,66 +451,13 @@ module Gitlab
 
       # Adds a column with a default value without locking an entire table.
       #
-      # This method runs the following steps:
-      #
-      # 1. Add the column with a default value of NULL.
-      # 2. Change the default value of the column to the specified value.
-      # 3. Update all existing rows in batches.
-      # 4. Set a `NOT NULL` constraint on the column if desired (the default).
-      #
-      # These steps ensure a column can be added to a large and commonly used
-      # table without locking the entire table for the duration of the table
-      # modification.
-      #
-      # table - The name of the table to update.
-      # column - The name of the column to add.
-      # type - The column type (e.g. `:integer`).
-      # default - The default value for the column.
-      # limit - Sets a column limit. For example, for :integer, the default is
-      #         4-bytes. Set `limit: 8` to allow 8-byte integers.
-      # allow_null - When set to `true` the column will allow NULL values, the
-      #              default is to not allow NULL values.
-      #
-      # This method can also take a block which is passed directly to the
-      # `update_column_in_batches` method.
-      def add_column_with_default(table, column, type, default:, limit: nil, allow_null: false, update_column_in_batches_args: {}, &block)
-        if transaction_open?
-          raise 'add_column_with_default can not be run inside a transaction, ' \
-            'you can disable transactions by calling disable_ddl_transaction! ' \
-            'in the body of your migration class'
-        end
-
-        disable_statement_timeout do
-          transaction do
-            if limit
-              add_column(table, column, type, default: nil, limit: limit)
-            else
-              add_column(table, column, type, default: nil)
-            end
-
-            # Changing the default before the update ensures any newly inserted
-            # rows already use the proper default value.
-            change_column_default(table, column, default)
-          end
-
-          begin
-            default_after_type_cast = connection.type_cast(default, column_for(table, column))
+      # @deprecated With PostgreSQL 11, adding columns with a default does not lead to a table rewrite anymore.
+      #             As such, this method is not needed anymore and the default `add_column` helper should be used.
+      #             This helper is subject to be removed in a >13.0 release.
+      def add_column_with_default(table, column, type, default:, limit: nil, allow_null: false)
+        raise 'Deprecated: add_column_with_default does not support being passed blocks anymore' if block_given?
 
-            if update_column_in_batches_args.any?
-              update_column_in_batches(table, column, default_after_type_cast, **update_column_in_batches_args, &block)
-            else
-              update_column_in_batches(table, column, default_after_type_cast, &block)
-            end
-
-            change_column_null(table, column, false) unless allow_null
-          # We want to rescue _all_ exceptions here, even those that don't inherit
-          # from StandardError.
-          rescue Exception => error # rubocop: disable all
-            remove_column(table, column)
-
-            raise error
-          end
-        end
+        add_column(table, column, type, default: default, limit: limit, null: allow_null)
       end
 
       # Renames a column without requiring downtime.
@@ -519,14 +473,20 @@ module Gitlab
       # new - The new column name.
       # type - The type of the new column. If no type is given the old column's
       #        type is used.
-      def rename_column_concurrently(table, old, new, type: nil)
+      # batch_column_name - option is for tables without primary key, in this
+      #        case another unique integer column can be used. Example: :user_id
+      def rename_column_concurrently(table, old, new, type: nil, batch_column_name: :id)
+        unless column_exists?(table, batch_column_name)
+          raise "Column #{batch_column_name} does not exist on #{table}"
+        end
+
         if transaction_open?
           raise 'rename_column_concurrently can not be run inside a transaction'
         end
 
         check_trigger_permissions!(table)
 
-        create_column_from(table, old, new, type: type)
+        create_column_from(table, old, new, type: type, batch_column_name: batch_column_name)
 
         install_rename_triggers(table, old, new)
       end
@@ -626,14 +586,20 @@ module Gitlab
       # new - The new column name.
       # type - The type of the old column. If no type is given the new column's
       #        type is used.
-      def undo_cleanup_concurrent_column_rename(table, old, new, type: nil)
+      # batch_column_name - option is for tables without primary key, in this
+      #        case another unique integer column can be used. Example: :user_id
+      def undo_cleanup_concurrent_column_rename(table, old, new, type: nil, batch_column_name: :id)
+        unless column_exists?(table, batch_column_name)
+          raise "Column #{batch_column_name} does not exist on #{table}"
+        end
+
         if transaction_open?
           raise 'undo_cleanup_concurrent_column_rename can not be run inside a transaction'
         end
 
         check_trigger_permissions!(table)
 
-        create_column_from(table, new, old, type: type)
+        create_column_from(table, new, old, type: type, batch_column_name: batch_column_name)
 
         install_rename_triggers(table, old, new)
       end
@@ -1063,6 +1029,8 @@ into similar problems in the future (e.g. when new tables are created).
       # batch_size - The maximum number of rows per job
       # other_arguments - Other arguments to send to the job
       #
+      # *Returns the final migration delay*
+      #
       # Example:
       #
       #     class Route < ActiveRecord::Base
@@ -1079,7 +1047,7 @@ into similar problems in the future (e.g. when new tables are created).
       #         # do something
       #       end
       #     end
-      def queue_background_migration_jobs_by_range_at_intervals(model_class, job_class_name, delay_interval, batch_size: BACKGROUND_MIGRATION_BATCH_SIZE, other_arguments: [])
+      def queue_background_migration_jobs_by_range_at_intervals(model_class, job_class_name, delay_interval, batch_size: BACKGROUND_MIGRATION_BATCH_SIZE, other_job_arguments: [], initial_delay: 0)
         raise "#{model_class} does not have an ID to use for batch ranges" unless model_class.column_names.include?('id')
 
         # To not overload the worker too much we enforce a minimum interval both
@@ -1088,14 +1056,19 @@ into similar problems in the future (e.g. when new tables are created).
           delay_interval = BackgroundMigrationWorker.minimum_interval
         end
 
+        final_delay = 0
+
         model_class.each_batch(of: batch_size) do |relation, index|
           start_id, end_id = relation.pluck(Arel.sql('MIN(id), MAX(id)')).first
 
           # `BackgroundMigrationWorker.bulk_perform_in` schedules all jobs for
           # the same time, which is not helpful in most cases where we wish to
           # spread the work over time.
-          migrate_in(delay_interval * index, job_class_name, [start_id, end_id] + other_arguments)
+          final_delay = initial_delay + delay_interval * index
+          migrate_in(final_delay, job_class_name, [start_id, end_id] + other_job_arguments)
         end
+
+        final_delay
       end
 
       # Fetches indexes on a column by name for postgres.
@@ -1315,12 +1288,73 @@ into similar problems in the future (e.g. when new tables are created).
         check_constraint_exists?(table, text_limit_name(table, column, name: constraint_name))
       end
 
+      # Migration Helpers for managing not null constraints
+      def add_not_null_constraint(table, column, constraint_name: nil, validate: true)
+        if column_is_nullable?(table, column)
+          add_check_constraint(
+            table,
+            "#{column} IS NOT NULL",
+            not_null_constraint_name(table, column, name: constraint_name),
+            validate: validate
+          )
+        else
+          warning_message = <<~MESSAGE
+            NOT NULL check constraint was not created:
+            column #{table}.#{column} is already defined as `NOT NULL`
+          MESSAGE
+
+          Rails.logger.warn warning_message
+        end
+      end
+
+      def validate_not_null_constraint(table, column, constraint_name: nil)
+        validate_check_constraint(
+          table,
+          not_null_constraint_name(table, column, name: constraint_name)
+        )
+      end
+
+      def remove_not_null_constraint(table, column, constraint_name: nil)
+        remove_check_constraint(
+          table,
+          not_null_constraint_name(table, column, name: constraint_name)
+        )
+      end
+
+      def check_not_null_constraint_exists?(table, column, constraint_name: nil)
+        check_constraint_exists?(
+          table,
+          not_null_constraint_name(table, column, name: constraint_name)
+        )
+      end
+
       private
 
+      def statement_timeout_disabled?
+        # This is a string of the form "100ms" or "0" when disabled
+        connection.select_value('SHOW statement_timeout') == "0"
+      end
+
+      def column_is_nullable?(table, column)
+        # Check if table.column has not been defined with NOT NULL
+        check_sql = <<~SQL
+          SELECT c.is_nullable
+          FROM information_schema.columns c
+          WHERE c.table_name = '#{table}'
+          AND c.column_name = '#{column}'
+        SQL
+
+        connection.select_value(check_sql) == 'YES'
+      end
+
       def text_limit_name(table, column, name: nil)
         name.presence || check_constraint_name(table, column, 'max_length')
       end
 
+      def not_null_constraint_name(table, column, name: nil)
+        name.presence || check_constraint_name(table, column, 'not_null')
+      end
+
       def missing_schema_object_message(table, type, name)
         <<~MESSAGE
           Could not find #{type} "#{name}" on table "#{table}" which was referenced during the migration.
@@ -1348,7 +1382,7 @@ into similar problems in the future (e.g. when new tables are created).
         "ON DELETE #{on_delete.upcase}"
       end
 
-      def create_column_from(table, old, new, type: nil)
+      def create_column_from(table, old, new, type: nil, batch_column_name: :id)
         old_col = column_for(table, old)
         new_type = type || old_col.type
 
@@ -1362,9 +1396,9 @@ into similar problems in the future (e.g. when new tables are created).
         # necessary since we copy over old values further down.
         change_column_default(table, new, old_col.default) unless old_col.default.nil?
 
-        update_column_in_batches(table, new, Arel::Table.new(table)[old])
+        update_column_in_batches(table, new, Arel::Table.new(table)[old], batch_column_name: batch_column_name)
 
-        change_column_null(table, new, false) unless old_col.null
+        add_not_null_constraint(table, new) unless old_col.null
 
         copy_indexes(table, old, new)
         copy_foreign_keys(table, old, new)
diff --git a/lib/gitlab/database/partitioning_migration_helpers.rb b/lib/gitlab/database/partitioning_migration_helpers.rb
new file mode 100644
index 00000000000..55649ebbf8a
--- /dev/null
+++ b/lib/gitlab/database/partitioning_migration_helpers.rb
@@ -0,0 +1,122 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module PartitioningMigrationHelpers
+      include SchemaHelpers
+
+      def add_partitioned_foreign_key(from_table, to_table, column: nil, primary_key: :id, on_delete: :cascade)
+        cascade_delete = extract_cascade_option(on_delete)
+
+        update_foreign_keys(from_table, to_table, column, primary_key, cascade_delete) do |current_keys, existing_key, specified_key|
+          if existing_key.nil?
+            unless specified_key.save
+              raise "failed to create foreign key: #{specified_key.errors.full_messages.to_sentence}"
+            end
+
+            current_keys << specified_key
+          else
+            Rails.logger.warn "foreign key not added because it already exists: #{specified_key}" # rubocop:disable Gitlab/RailsLogger
+            current_keys
+          end
+        end
+      end
+
+      def remove_partitioned_foreign_key(from_table, to_table, column: nil, primary_key: :id)
+        update_foreign_keys(from_table, to_table, column, primary_key) do |current_keys, existing_key, specified_key|
+          if existing_key
+            existing_key.destroy!
+            current_keys.delete(existing_key)
+          else
+            Rails.logger.warn "foreign key not removed because it doesn't exist: #{specified_key}" # rubocop:disable Gitlab/RailsLogger
+          end
+
+          current_keys
+        end
+      end
+
+      def fk_function_name(table)
+        object_name(table, 'fk_cascade_function')
+      end
+
+      def fk_trigger_name(table)
+        object_name(table, 'fk_cascade_trigger')
+      end
+
+      private
+
+      def fk_from_spec(from_table, to_table, from_column, to_column, cascade_delete)
+        PartitionedForeignKey.new(from_table: from_table.to_s, to_table: to_table.to_s, from_column: from_column.to_s,
+                                  to_column: to_column.to_s, cascade_delete: cascade_delete)
+      end
+
+      def update_foreign_keys(from_table, to_table, from_column, to_column, cascade_delete = nil)
+        if transaction_open?
+          raise 'partitioned foreign key operations can not be run inside a transaction block, ' \
+            'you can disable transaction blocks by calling disable_ddl_transaction! ' \
+            'in the body of your migration class'
+        end
+
+        from_column ||= "#{to_table.to_s.singularize}_id"
+        specified_key = fk_from_spec(from_table, to_table, from_column, to_column, cascade_delete)
+
+        current_keys = PartitionedForeignKey.by_referenced_table(to_table).to_a
+        existing_key = find_existing_key(current_keys, specified_key)
+
+        final_keys = yield current_keys, existing_key, specified_key
+
+        fn_name = fk_function_name(to_table)
+        trigger_name = fk_trigger_name(to_table)
+
+        with_lock_retries do
+          drop_trigger(to_table, trigger_name, if_exists: true)
+
+          if final_keys.empty?
+            drop_function(fn_name, if_exists: true)
+          else
+            create_or_replace_fk_function(fn_name, final_keys)
+            create_function_trigger(trigger_name, fn_name, fires: "AFTER DELETE ON #{to_table}")
+          end
+        end
+      end
+
+      def extract_cascade_option(on_delete)
+        case on_delete
+        when :cascade then true
+        when :nullify then false
+        else raise ArgumentError, "invalid option #{on_delete} for :on_delete"
+        end
+      end
+
+      def with_lock_retries(&block)
+        Gitlab::Database::WithLockRetries.new({
+          klass: self.class,
+          logger: Gitlab::BackgroundMigration::Logger
+        }).run(&block)
+      end
+
+      def find_existing_key(keys, key)
+        keys.find { |k| k.from_table == key.from_table && k.from_column == key.from_column }
+      end
+
+      def create_or_replace_fk_function(fn_name, fk_specs)
+        create_trigger_function(fn_name, replace: true) do
+          cascade_statements = build_cascade_statements(fk_specs)
+          cascade_statements << 'RETURN OLD;'
+
+          cascade_statements.join("\n")
+        end
+      end
+
+      def build_cascade_statements(foreign_keys)
+        foreign_keys.map do |fks|
+          if fks.cascade_delete?
+            "DELETE FROM #{fks.from_table} WHERE #{fks.from_column} = OLD.#{fks.to_column};"
+          else
+            "UPDATE #{fks.from_table} SET #{fks.from_column} = NULL WHERE #{fks.from_column} = OLD.#{fks.to_column};"
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/partitioning_migration_helpers/partitioned_foreign_key.rb b/lib/gitlab/database/partitioning_migration_helpers/partitioned_foreign_key.rb
new file mode 100644
index 00000000000..f9a90511f9b
--- /dev/null
+++ b/lib/gitlab/database/partitioning_migration_helpers/partitioned_foreign_key.rb
@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module PartitioningMigrationHelpers
+      class PartitionedForeignKey < ApplicationRecord
+        validates_with PartitionedForeignKeyValidator
+
+        scope :by_referenced_table, ->(table) { where(to_table: table) }
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/partitioning_migration_helpers/partitioned_foreign_key_validator.rb b/lib/gitlab/database/partitioning_migration_helpers/partitioned_foreign_key_validator.rb
new file mode 100644
index 00000000000..089cf2b8931
--- /dev/null
+++ b/lib/gitlab/database/partitioning_migration_helpers/partitioned_foreign_key_validator.rb
@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module PartitioningMigrationHelpers
+      class PartitionedForeignKeyValidator < ActiveModel::Validator
+        def validate(record)
+          validate_key_part(record, :from_table, :from_column)
+          validate_key_part(record, :to_table, :to_column)
+        end
+
+        private
+
+        def validate_key_part(record, table_field, column_field)
+          if !connection.table_exists?(record[table_field])
+            record.errors.add(table_field, 'must be a valid table')
+          elsif !connection.column_exists?(record[table_field], record[column_field])
+            record.errors.add(column_field, 'must be a valid column')
+          end
+        end
+
+        def connection
+          ActiveRecord::Base.connection
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/rename_reserved_paths_migration/v1/rename_base.rb b/lib/gitlab/database/rename_reserved_paths_migration/v1/rename_base.rb
index 565f34b78b7..2c9d0d6c0d1 100644
--- a/lib/gitlab/database/rename_reserved_paths_migration/v1/rename_base.rb
+++ b/lib/gitlab/database/rename_reserved_paths_migration/v1/rename_base.rb
@@ -157,7 +157,7 @@ module Gitlab
               failed_reverts = []
 
               while rename_info = redis.lpop(key)
-                path_before_rename, path_after_rename = JSON.parse(rename_info)
+                path_before_rename, path_after_rename = Gitlab::Json.parse(rename_info)
                 say "renaming #{type} from #{path_after_rename} back to #{path_before_rename}"
                 begin
                   yield(path_before_rename, path_after_rename)
diff --git a/lib/gitlab/database/schema_helpers.rb b/lib/gitlab/database/schema_helpers.rb
new file mode 100644
index 00000000000..f8d01c78ae8
--- /dev/null
+++ b/lib/gitlab/database/schema_helpers.rb
@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module SchemaHelpers
+      def create_trigger_function(name, replace: true)
+        replace_clause = optional_clause(replace, "OR REPLACE")
+        execute(<<~SQL)
+          CREATE #{replace_clause} FUNCTION #{name}()
+          RETURNS TRIGGER AS
+          $$
+          BEGIN
+          #{yield}
+          END
+          $$ LANGUAGE PLPGSQL
+        SQL
+      end
+
+      def create_function_trigger(name, fn_name, fires: nil)
+        execute(<<~SQL)
+          CREATE TRIGGER #{name}
+          #{fires}
+          FOR EACH ROW
+          EXECUTE PROCEDURE #{fn_name}()
+        SQL
+      end
+
+      def drop_function(name, if_exists: true)
+        exists_clause = optional_clause(if_exists, "IF EXISTS")
+        execute("DROP FUNCTION #{exists_clause} #{name}()")
+      end
+
+      def drop_trigger(table_name, name, if_exists: true)
+        exists_clause = optional_clause(if_exists, "IF EXISTS")
+        execute("DROP TRIGGER #{exists_clause} #{name} ON #{table_name}")
+      end
+
+      def object_name(table, type)
+        identifier = "#{table}_#{type}"
+        hashed_identifier = Digest::SHA256.hexdigest(identifier).first(10)
+
+        "#{type}_#{hashed_identifier}"
+      end
+
+      private
+
+      def optional_clause(flag, clause)
+        flag ? clause : ""
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/with_lock_retries.rb b/lib/gitlab/database/with_lock_retries.rb
index 2f36bfa1480..bebcba6f42e 100644
--- a/lib/gitlab/database/with_lock_retries.rb
+++ b/lib/gitlab/database/with_lock_retries.rb
@@ -78,12 +78,18 @@ module Gitlab
           run_block_with_transaction
         rescue ActiveRecord::LockWaitTimeout
           if retry_with_lock_timeout?
+            disable_idle_in_transaction_timeout
             wait_until_next_retry
+            reset_db_settings
 
             retry
           else
+            reset_db_settings
             run_block_without_lock_timeout
           end
+
+        ensure
+          reset_db_settings
         end
       end
 
@@ -153,6 +159,14 @@ module Gitlab
       def current_sleep_time_in_seconds
         timing_configuration[current_iteration - 1][1].to_f
       end
+
+      def disable_idle_in_transaction_timeout
+        execute("SET LOCAL idle_in_transaction_session_timeout TO '0'")
+      end
+
+      def reset_db_settings
+        execute('RESET idle_in_transaction_session_timeout; RESET lock_timeout')
+      end
     end
   end
 end