1 files changed, 13 insertions, 13 deletions
diff --git a/app/models/loose_foreign_keys/deleted_record.rb b/app/models/loose_foreign_keys/deleted_record.rb
index 01d6d068ca4..c3b3e76f67b 100644
--- a/app/models/loose_foreign_keys/deleted_record.rb
+++ b/app/models/loose_foreign_keys/deleted_record.rb
@@ -4,29 +4,29 @@ class LooseForeignKeys::DeletedRecord < ApplicationRecord
   self.primary_key = :id
 
   scope :for_table, -> (table) { where(fully_qualified_table_name: table) }
-  scope :ordered_by_id, -> { order(:id, :primary_key_value) }
-  # This needs to be parameterized once we start adding partitions
-  scope :for_partition, -> { where(partition: 1) }
+  scope :consume_order, -> { order(:partition, :consume_after, :id) }
 
   enum status: { pending: 1, processed: 2 }, _prefix: :status
 
   def self.load_batch_for_table(table, batch_size)
     for_table(table)
-      .for_partition
       .status_pending
-      .ordered_by_id
+      .consume_order
       .limit(batch_size)
       .to_a
   end
 
-  def self.mark_records_processed_for_table_between(table, from_record, to_record)
-    from = from_record.id
-    to = to_record.id
+  def self.mark_records_processed(all_records)
+    # Run a query for each partition to optimize the row lookup by primary key (partition, id)
+    update_count = 0
 
-    for_table(table)
-      .for_partition
-      .status_pending
-      .where(id: from..to)
-      .update_all(status: :processed)
+    all_records.group_by(&:partition).each do |partition, records_within_partition|
+      update_count += status_pending
+        .where(partition: partition)
+        .where(id: records_within_partition.pluck(:id))
+        .update_all(status: :processed)
+    end
+
+    update_count
   end
 end