Add latest changes from gitlab-org/gitlab@14-5-stable-eev14.5.0-rc42

3 files changed, 234 insertions, 0 deletions

diff --git a/app/services/loose_foreign_keys/batch_cleaner_service.rb b/app/services/loose_foreign_keys/batch_cleaner_service.rb
new file mode 100644
index 00000000000..06c05e8ff54
--- /dev/null
+++ b/app/services/loose_foreign_keys/batch_cleaner_service.rb
@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+module LooseForeignKeys
+  class BatchCleanerService
+    def initialize(parent_klass:, deleted_parent_records:, modification_tracker: LooseForeignKeys::ModificationTracker.new, models_by_table_name:)
+      @parent_klass = parent_klass
+      @deleted_parent_records = deleted_parent_records
+      @modification_tracker = modification_tracker
+      @models_by_table_name = models_by_table_name
+      @deleted_records_counter = Gitlab::Metrics.counter(
+        :loose_foreign_key_processed_deleted_records,
+        'The number of processed loose foreign key deleted records'
+      )
+    end
+
+    def execute
+      parent_klass.loose_foreign_key_definitions.each do |foreign_key_definition|
+        run_cleaner_service(foreign_key_definition, with_skip_locked: true)
+        break if modification_tracker.over_limit?
+
+        run_cleaner_service(foreign_key_definition, with_skip_locked: false)
+        break if modification_tracker.over_limit?
+      end
+
+      return if modification_tracker.over_limit?
+
+      # At this point, all associations are cleaned up, we can update the status of the parent records
+      update_count = LooseForeignKeys::DeletedRecord.mark_records_processed(deleted_parent_records)
+
+      deleted_records_counter.increment({ table: parent_klass.table_name, db_config_name: LooseForeignKeys::DeletedRecord.connection.pool.db_config.name }, update_count)
+    end
+
+    private
+
+    attr_reader :parent_klass, :deleted_parent_records, :modification_tracker, :models_by_table_name, :deleted_records_counter
+
+    def record_result(cleaner, result)
+      if cleaner.async_delete?
+        modification_tracker.add_deletions(result[:table], result[:affected_rows])
+      elsif cleaner.async_nullify?
+        modification_tracker.add_updates(result[:table], result[:affected_rows])
+      end
+    end
+
+    def run_cleaner_service(foreign_key_definition, with_skip_locked:)
+      cleaner = CleanerService.new(
+        model: models_by_table_name.fetch(foreign_key_definition.to_table),
+        foreign_key_definition: foreign_key_definition,
+        deleted_parent_records: deleted_parent_records,
+        with_skip_locked: with_skip_locked
+      )
+
+      loop do
+        result = cleaner.execute
+        record_result(cleaner, result)
+
+        break if modification_tracker.over_limit? || result[:affected_rows] == 0
+      end
+    end
+  end
+end
diff --git a/app/services/loose_foreign_keys/cleaner_service.rb b/app/services/loose_foreign_keys/cleaner_service.rb
new file mode 100644
index 00000000000..8fe053e2edf
--- /dev/null
+++ b/app/services/loose_foreign_keys/cleaner_service.rb
@@ -0,0 +1,99 @@
+# frozen_string_literal: true
+
+module LooseForeignKeys
+  # rubocop: disable CodeReuse/ActiveRecord
+  class CleanerService
+    DELETE_LIMIT = 1000
+    UPDATE_LIMIT = 500
+
+    delegate :connection, to: :model
+
+    def initialize(model:, foreign_key_definition:, deleted_parent_records:, with_skip_locked: false)
+      @model = model
+      @foreign_key_definition = foreign_key_definition
+      @deleted_parent_records = deleted_parent_records
+      @with_skip_locked = with_skip_locked
+    end
+
+    def execute
+      result = connection.execute(build_query)
+
+      { affected_rows: result.cmd_tuples, table: foreign_key_definition.to_table }
+    end
+
+    def async_delete?
+      foreign_key_definition.on_delete == :async_delete
+    end
+
+    def async_nullify?
+      foreign_key_definition.on_delete == :async_nullify
+    end
+
+    private
+
+    attr_reader :model, :foreign_key_definition, :deleted_parent_records, :with_skip_locked
+
+    def build_query
+      query = if async_delete?
+                delete_query
+              elsif async_nullify?
+                update_query
+              else
+                raise "Invalid on_delete argument: #{foreign_key_definition.on_delete}"
+              end
+
+      unless query.include?(%{"#{foreign_key_definition.column}" IN (})
+        raise("FATAL: foreign key condition is missing from the generated query: #{query}")
+      end
+
+      query
+    end
+
+    def arel_table
+      @arel_table ||= model.arel_table
+    end
+
+    def primary_keys
+      @primary_keys ||= connection.primary_keys(model.table_name).map { |key| arel_table[key] }
+    end
+
+    def quoted_table_name
+      @quoted_table_name ||= Arel.sql(connection.quote_table_name(model.table_name))
+    end
+
+    def delete_query
+      query = Arel::DeleteManager.new
+      query.from(quoted_table_name)
+
+      add_in_query_with_limit(query, DELETE_LIMIT)
+    end
+
+    def update_query
+      query = Arel::UpdateManager.new
+      query.table(quoted_table_name)
+      query.set([[arel_table[foreign_key_definition.column], nil]])
+
+      add_in_query_with_limit(query, UPDATE_LIMIT)
+    end
+
+    # IN query with one or composite primary key
+    # WHERE (primary_key1, primary_key2) IN (subselect)
+    def add_in_query_with_limit(query, limit)
+      columns = Arel::Nodes::Grouping.new(primary_keys)
+      query.where(columns.in(in_query_with_limit(limit))).to_sql
+    end
+
+    # Builds the following sub-query
+    # SELECT primary_keys FROM table WHERE foreign_key IN (1, 2, 3) LIMIT N
+    def in_query_with_limit(limit)
+      in_query = Arel::SelectManager.new
+      in_query.from(quoted_table_name)
+      in_query.where(arel_table[foreign_key_definition.column].in(deleted_parent_records.map(&:primary_key_value)))
+      in_query.projections = primary_keys
+      in_query.take(limit)
+      in_query.lock(Arel.sql('FOR UPDATE SKIP LOCKED')) if with_skip_locked
+      in_query
+    end
+  end
+  # rubocop: enable CodeReuse/ActiveRecord
+end
diff --git a/app/services/loose_foreign_keys/process_deleted_records_service.rb b/app/services/loose_foreign_keys/process_deleted_records_service.rb
new file mode 100644
index 00000000000..735fc8a2415
--- /dev/null
+++ b/app/services/loose_foreign_keys/process_deleted_records_service.rb
@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+
+module LooseForeignKeys
+  class ProcessDeletedRecordsService
+    BATCH_SIZE = 1000
+
+    def initialize(connection:)
+      @connection = connection
+    end
+
+    def execute
+      modification_tracker = ModificationTracker.new
+
+      tracked_tables.cycle do |table|
+        records = load_batch_for_table(table)
+
+        if records.empty?
+          tracked_tables.delete(table)
+          next
+        end
+
+        break if modification_tracker.over_limit?
+
+        model = find_parent_model!(table)
+
+        LooseForeignKeys::BatchCleanerService
+          .new(parent_klass: model,
+               deleted_parent_records: records,
+               modification_tracker: modification_tracker,
+               models_by_table_name: models_by_table_name)
+          .execute
+
+        break if modification_tracker.over_limit?
+      end
+
+      modification_tracker.stats
+    end
+
+    private
+
+    attr_reader :connection
+
+    def load_batch_for_table(table)
+      fully_qualified_table_name = "#{current_schema}.#{table}"
+      LooseForeignKeys::DeletedRecord.load_batch_for_table(fully_qualified_table_name, BATCH_SIZE)
+    end
+
+    def find_parent_model!(table)
+      models_by_table_name.fetch(table)
+    end
+
+    def current_schema
+      @current_schema = connection.current_schema
+    end
+
+    def tracked_tables
+      @tracked_tables ||= models_by_table_name
+        .select { |table_name, model| model.respond_to?(:loose_foreign_key_definitions) }
+        .keys
+    end
+
+    def models_by_table_name
+      @models_by_table_name ||= begin
+        all_models
+          .select(&:base_class?)
+          .index_by(&:table_name)
+      end
+    end
+
+    def all_models
+      ApplicationRecord.descendants
+    end
+  end
+end