1 files changed, 13 insertions, 8 deletions
diff --git a/lib/gitlab/background_migration/set_merge_request_diff_files_count.rb b/lib/gitlab/background_migration/set_merge_request_diff_files_count.rb
index 9f765d03d62..527dd2a0a83 100644
--- a/lib/gitlab/background_migration/set_merge_request_diff_files_count.rb
+++ b/lib/gitlab/background_migration/set_merge_request_diff_files_count.rb
@@ -4,13 +4,18 @@ module Gitlab
   module BackgroundMigration
     # Sets the MergeRequestDiff#files_count value for old rows
     class SetMergeRequestDiffFilesCount
-      COUNT_SUBQUERY = <<~SQL
-        files_count = (
-          SELECT count(*)
-          FROM merge_request_diff_files
-          WHERE merge_request_diff_files.merge_request_diff_id = merge_request_diffs.id
-        )
-      SQL
+      # Some historic data has a *lot* of files. Apply a sentinel to these cases
+      FILES_COUNT_SENTINEL = 2**15 - 1
+
+      def self.count_subquery
+        <<~SQL
+          files_count = (
+            SELECT LEAST(#{FILES_COUNT_SENTINEL}, count(*))
+            FROM merge_request_diff_files
+            WHERE merge_request_diff_files.merge_request_diff_id = merge_request_diffs.id
+          )
+        SQL
+      end
 
       class MergeRequestDiff < ActiveRecord::Base # rubocop:disable Style/Documentation
         include EachBatch
@@ -20,7 +25,7 @@ module Gitlab
 
       def perform(start_id, end_id)
         MergeRequestDiff.where(id: start_id..end_id).each_batch do |relation|
-          relation.update_all(COUNT_SUBQUERY)
+          relation.update_all(self.class.count_subquery)
         end
       end
     end